fix: RAG with documents (#1337)

Summary:
This was broken by
https://github.com/meta-llama/llama-stack/pull/1015/files#r1975394190

Test Plan:

added e2e test
This commit is contained in:
ehhuang 2025-02-28 16:51:00 -08:00 committed by GitHub
parent 7854af8b52
commit 21ec67356c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 66 additions and 0 deletions

View file

@ -533,6 +533,11 @@ class ChatAgent(ShieldRunnerMixin):
if documents:
await self.handle_documents(session_id, documents, input_messages, tool_defs)
session_info = await self.storage.get_session_info(session_id)
# if the session has a memory bank id, let the memory tool use it
if session_info and session_info.vector_db_id:
toolgroup_args[RAG_TOOL_GROUP]["vector_db_ids"].append(session_info.vector_db_id)
output_attachments = []
n_iter = 0