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

add back RAG logic

Browse source
This commit is contained in:
Hardik Shah 2025-01-30 14:03:41 -08:00
parent 66956277a0
commit e2bfaf9cd3
2 changed files with 16 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llama_stack/providers/inline/agents/meta_reference/agent_instance.py
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
@ -477,6 +478,12 @@ class ChatAgent(ShieldRunnerMixin):
span.set_attribute("output", retrieved_context)
span.set_attribute("tool_name", MEMORY_QUERY_TOOL)
# 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 = []
n_iter = 0

15
tests/client-sdk/agents/test_agents.py
View file

@ -211,7 +211,7 @@ def test_code_interpreter_for_attachments(llama_stack_client, agent_config):
}
codex_agent = Agent(llama_stack_client, agent_config)
session_id = codex_agent.create_session("test-session")
session_id = codex_agent.create_session(f"test-session-{uuid4()}")
inflation_doc = AgentDocument(
content="https://raw.githubusercontent.com/meta-llama/llama-stack-apps/main/examples/resources/inflation.csv",
mime_type="text/csv",
@ -299,7 +299,7 @@ def test_rag_agent(llama_stack_client, agent_config):
],
}
rag_agent = Agent(llama_stack_client, agent_config)
session_id = rag_agent.create_session("test-session")
session_id = rag_agent.create_session(f"test-session-{uuid4()}")
user_prompts = [
"What are the top 5 topics that were explained? Only list succinct bullet points.",
]
@ -334,7 +334,7 @@ def test_rag_and_code_agent(llama_stack_client, agent_config):
llama_stack_client.tool_runtime.rag_tool.insert(
documents=documents,
vector_db_id=vector_db_id,
chunk_size_in_tokens=512,
chunk_size_in_tokens=128,
)
agent_config = {
**agent_config,
@ -347,7 +347,6 @@ def test_rag_and_code_agent(llama_stack_client, agent_config):
],
}
agent = Agent(llama_stack_client, agent_config)
session_id = agent.create_session("test-session")
inflation_doc = Document(
document_id="test_csv",
content="https://raw.githubusercontent.com/meta-llama/llama-stack-apps/main/examples/resources/inflation.csv",
@ -355,22 +354,26 @@ def test_rag_and_code_agent(llama_stack_client, agent_config):
metadata={},
)
user_prompts = [
(
"Here is a csv file, can you describe it?",
[inflation_doc],
"code_interpreter",
),
(
"What are the top 5 topics that were explained? Only list succinct bullet points.",
[],
"query_from_memory",
),
("What is the average yearly inflation?", [inflation_doc], "code_interpreter"),
]
for prompt, docs, tool_name in user_prompts:
print(f"User> {prompt}")
session_id = agent.create_session(f"test-session-{uuid4()}")
response = agent.create_turn(
messages=[{"role": "user", "content": prompt}],
session_id=session_id,
documents=docs,
)
logs = [str(log) for log in EventLogger().log(response) if log is not None]
logs_str = "".join(logs)
assert f"Tool:{tool_name}" in logs_str