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 48

fix(a2a): merge fresh agent vectors into the live cache and drop entries of another dimension

Browse source
This commit is contained in:
mateo-berri 2026-08-27 20:58:14 -07:00
parent db02cf81e5
commit 4026aa6575
2 changed files with 31 additions and 3 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

@ -193,6 +193,14 @@ class AgentSearchIndex:
def __init__(self) -> None:
self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({})
def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]:
kept: Final = {
text: vector
for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items()
if len(vector) == len(embedded.query_vector)
}
return MappingProxyType({**kept, **embedded.vectors})
async def search(
self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder, embedding_model: str
) -> AgentSearchHits | AgentSearchEmbeddingFailed:
@ -207,9 +215,7 @@ class AgentSearchIndex:
return AgentSearchEmbeddingFailed(
reason=f"embedding model {embedding_model} returned vectors of mixed dimensions"
)
self._vectors = MappingProxyType(
{**self._vectors, embedding_model: MappingProxyType({**cached, **embedded.vectors})}
)
self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)})
ranked: Final = sorted(
(
AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.vectors[text]))

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

@ -1,3 +1,4 @@
import asyncio
from collections.abc import Sequence
from types import MappingProxyType
from typing import Final
@ -84,6 +85,7 @@ class FixedDimensionEmbedder:
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
self.calls.append(tuple(texts))
await asyncio.sleep(0)
return tuple((1.0,) * self.dimensions for _ in texts)
@ -156,6 +158,26 @@ class TestAgentSearchIndex:
("language translation", *(agent_search_text(agent) for agent in AGENTS)),
]
@pytest.mark.asyncio
async def test_re_embedding_a_subset_drops_the_other_agents_old_vectors(self) -> None:
index = AgentSearchIndex()
await index.search("language translation", AGENTS, top_k=5, embed=FakeEmbedder(), embedding_model="m")
wide = FixedDimensionEmbedder(2)
await index.search("language translation", AGENTS[:1], top_k=5, embed=wide, embedding_model="m")
await index.search("language translation", AGENTS, top_k=5, embed=wide, embedding_model="m")
assert wide.calls[-1] == ("language translation", *(agent_search_text(agent) for agent in AGENTS[1:]))
@pytest.mark.asyncio
async def test_concurrent_searches_keep_each_others_vectors(self) -> None:
index = AgentSearchIndex()
embedder = FixedDimensionEmbedder(3)
await asyncio.gather(
index.search("q", AGENTS[:1], top_k=5, embed=embedder, embedding_model="m"),
index.search("q", AGENTS[1:], top_k=5, embed=embedder, embedding_model="m"),
)
await index.search("q", AGENTS, top_k=5, embed=embedder, embedding_model="m")
assert embedder.calls[-1] == ("q",)
@pytest.mark.asyncio
async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None:
async def mixed(texts: Sequence[str]) -> Sequence[Vector]: