diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 896a22fbf9d..fe4b255235f 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -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, diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index 4b674eb142b..c734d2fd9f3 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -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