Fix Agents to support code and rag simultaneously (#908)

# What does this PR do?

Fixes a bug where agents were not working when both rag and
code-interpreter were added as tools.


## Test Plan

Added a new client_sdk test which tests for this scenario 
```
LLAMA_STACK_CONFIG=together pytest -s -v  tests/client-sdk -k 'test_rag_and_code_agent'
```

---------

Co-authored-by: Hardik Shah <hjshah@fb.com>
This commit is contained in:
Hardik Shah 2025-01-30 17:09:34 -08:00 committed by GitHub
parent 94051cfe9e
commit 97eb3eecea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 83 additions and 8 deletions

View file

@ -66,6 +66,7 @@ from llama_stack.apis.vector_io import VectorIO
from llama_stack.providers.utils.kvstore import KVStore
from llama_stack.providers.utils.memory.vector_store import concat_interleaved_content
from llama_stack.providers.utils.telemetry import tracing
from .persistence import AgentPersistence
from .safety import SafetyException, ShieldRunnerMixin
@ -476,9 +477,12 @@ class ChatAgent(ShieldRunnerMixin):
)
span.set_attribute("output", retrieved_context)
span.set_attribute("tool_name", MEMORY_QUERY_TOOL)
if retrieved_context:
last_message = input_messages[-1]
last_message.context = retrieved_context
# append retrieved_context to the last user message
for message in input_messages[::-1]:
if isinstance(message, UserMessage):
message.context = retrieved_context
break
output_attachments = []