phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-08 13:00:52 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add vector_db_id to chunk metadata

Browse source 
Adding unit tests
This commit is contained in:
are-ces 2025-08-26 13:19:59 +02:00
parent eed25fc6e4
commit b50cb25b57
2 changed files with 81 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

14
llama_stack/providers/inline/tool_runtime/rag/memory.py
View file

@ -131,8 +131,18 @@ class MemoryToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, RAGToolRunti
for vector_db_id in vector_db_ids
]
results: list[QueryChunksResponse] = await asyncio.gather(*tasks)
chunks = [c for r in results for c in r.chunks]
scores = [s for r in results for s in r.scores]
chunks = []
scores = []
for vector_db_id, result in zip(vector_db_ids, results, strict=False):
for chunk, score in zip(result.chunks, result.scores, strict=False):
if not hasattr(chunk, "metadata") or chunk.metadata is None:
chunk.metadata = {}
chunk.metadata["vector_db_id"] = vector_db_id
chunks.append(chunk)
scores.append(score)
if not chunks:
return RAGQueryResult(content=None)