scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(agent_search): evict stale-dimension vectors when re-embedding a subset

Browse source 
The cache write in AgentSearchIndex.search merged the pre-await snapshot with the re-embed result. When a dimension change forced re-embedding only the current request's texts, old-size vectors for the other agents were preserved next to new-size ones, so the next broader search saw mixed dimensions and re-embedded everyone, and concurrent writes to the same model could also be dropped.
This commit is contained in:
Cursor Agent 2026-08-28 02:55:29 +00:00
parent db02cf81e5
commit efbeae4393
No known key found for this signature in database
2 changed files with 21 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/agent_endpoints/agent_search.py
View file

@ -199,17 +199,21 @@ class AgentSearchIndex:
if not agents: if not agents:
return AgentSearchHits(hits=()) return AgentSearchHits(hits=())
texts: Final = tuple(agent_search_text(agent) for agent in agents) texts: Final = tuple(agent_search_text(agent) for agent in agents)
cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) embedded: Final = await _embed_query_and_agents(
embedded: Final = await _embed_query_and_agents(embed, query, texts, cached) embed, query, texts, self._vectors.get(embedding_model, _NO_VECTORS)
)
if isinstance(embedded, AgentSearchEmbeddingFailed): if isinstance(embedded, AgentSearchEmbeddingFailed):
return embedded return embedded
if not _same_dimension(embedded.query_vector, embedded.vectors, texts): if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
return AgentSearchEmbeddingFailed( return AgentSearchEmbeddingFailed(
reason=f"embedding model {embedding_model} returned vectors of mixed dimensions" reason=f"embedding model {embedding_model} returned vectors of mixed dimensions"
) )
self._vectors = MappingProxyType( current: Final = self._vectors.get(embedding_model, _NO_VECTORS)
{**self._vectors, embedding_model: MappingProxyType({**cached, **embedded.vectors})} dim: Final = len(embedded.query_vector)
merged: Final = MappingProxyType(
{text: vector for text, vector in chain(current.items(), embedded.vectors.items()) if len(vector) == dim}
) )
self._vectors = MappingProxyType({**self._vectors, embedding_model: merged})
ranked: Final = sorted( ranked: Final = sorted(
( (
AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.vectors[text])) AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.vectors[text]))

13
tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
View file

@ -156,6 +156,19 @@ class TestAgentSearchIndex:
("language translation", *(agent_search_text(agent) for agent in AGENTS)), ("language translation", *(agent_search_text(agent) for agent in AGENTS)),
] ]
@pytest.mark.asyncio
async def test_subset_reembed_after_dimension_change_evicts_stale_vectors(self) -> None:
index = AgentSearchIndex()
await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
narrow = FixedDimensionEmbedder(2)
await index.search("language translation", (TRANSLATOR,), top_k=5, embed=narrow, embedding_model="m")
broader = FixedDimensionEmbedder(2)
outcome = await index.search("language translation", AGENTS, top_k=5, embed=broader, embedding_model="m")
assert isinstance(outcome, AgentSearchHits)
assert broader.calls == [
("language translation", agent_search_text(SQL_ANALYST), agent_search_text(TRIP_PLANNER)),
]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None: async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None:
async def mixed(texts: Sequence[str]) -> Sequence[Vector]: async def mixed(texts: Sequence[str]) -> Sequence[Vector]: