mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(a2a): drop cached agent vectors whose dimension no longer matches the query
When the configured agent_search_embedding_model or a router fallback switches to a model with a different embedding dimension, the query vector shape stops matching cached agent vectors. cosine_similarity uses zip(..., strict=True), so ranking raised ValueError outside the try/except that treats embedding failures as AgentSearchEmbeddingFailed. Because the cache never dropped the old vectors, every subsequent GET /v1/agents?query= and agent_search MCP call returned 500 until the worker restarted. After embedding the query, filter the cache to only vectors of the same dimension and re-embed any agent texts that got evicted.
This commit is contained in:
parent
e9cc9c9bc3
commit
2ed2486bf4
2 changed files with 29 additions and 2 deletions
|
|
@ -164,10 +164,24 @@ class AgentSearchIndex:
|
|||
return AgentSearchEmbeddingFailed(
|
||||
reason=f"embedding model returned {len(vectors)} vectors for {len(missing) + 1} inputs"
|
||||
)
|
||||
self._vectors = MappingProxyType(dict(chain(self._vectors.items(), zip(missing, vectors[1:], strict=True))))
|
||||
query_vector: Final = vectors[0]
|
||||
fresh: Final = dict(zip(missing, vectors[1:], strict=True))
|
||||
kept: Final = {text: vec for text, vec in self._vectors.items() if len(vec) == len(query_vector)}
|
||||
stale: Final = tuple(dict.fromkeys(text for text in texts if text not in kept and text not in fresh))
|
||||
try:
|
||||
refreshed: Final = await embed(stale) if stale else ()
|
||||
except (OpenAIError, ValueError, BudgetExceededError) as exc:
|
||||
return AgentSearchEmbeddingFailed(reason=f"re-embedding stale agent texts failed: {exc}")
|
||||
if len(refreshed) != len(stale):
|
||||
return AgentSearchEmbeddingFailed(
|
||||
reason=f"embedding model returned {len(refreshed)} vectors for {len(stale)} inputs"
|
||||
)
|
||||
self._vectors = MappingProxyType(
|
||||
dict(chain(kept.items(), fresh.items(), zip(stale, refreshed, strict=True)))
|
||||
)
|
||||
ranked: Final = sorted(
|
||||
(
|
||||
AgentSearchHit(agent=agent, score=cosine_similarity(vectors[0], self._vectors[text]))
|
||||
AgentSearchHit(agent=agent, score=cosine_similarity(query_vector, self._vectors[text]))
|
||||
for agent, text in zip(agents, texts, strict=True)
|
||||
),
|
||||
key=lambda hit: hit.score,
|
||||
|
|
|
|||
|
|
@ -147,6 +147,19 @@ class TestAgentSearchIndex:
|
|||
outcome = await AgentSearchIndex().search("q", AGENTS, top_k=5, embed=short)
|
||||
assert isinstance(outcome, AgentSearchEmbeddingFailed)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dimension_change_invalidates_cached_agent_vectors(self) -> None:
|
||||
index = AgentSearchIndex()
|
||||
warm = FakeEmbedder()
|
||||
await index.search("language translation", AGENTS, top_k=5, embed=warm)
|
||||
|
||||
async def wider(texts: Sequence[str]) -> Sequence[Vector]:
|
||||
return tuple((1.0, 0.0, 0.0, 0.0) for _ in texts)
|
||||
|
||||
outcome = await index.search("language translation", AGENTS, top_k=5, embed=wider)
|
||||
assert isinstance(outcome, AgentSearchHits)
|
||||
assert len(outcome.hits) == len(AGENTS)
|
||||
|
||||
|
||||
class TestSearchAgents:
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue