fix: rag as attachment bug (#1392)

Summary:

Test Plan:
added new test
LLAMA_STACK_CONFIG=fireworks pytest -s -v
tests/api/agents/test_agents.py --safety-shield
meta-llama/Llama-Guard-3-8B
This commit is contained in:
ehhuang 2025-03-04 13:08:16 -08:00 committed by GitHub
parent e9a37bad63
commit fd8c991393
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 3830 additions and 403 deletions

View file

@ -401,7 +401,19 @@ def test_rag_agent(llama_stack_client_with_mocked_inference, agent_config, rag_t
assert expected_kw in response.output_message.content.lower()
def test_rag_agent_with_attachments(llama_stack_client_with_mocked_inference, agent_config):
@pytest.mark.parametrize(
"toolgroup",
[
dict(
name="builtin::rag/knowledge_search",
args={
"vector_db_ids": [],
},
),
"builtin::rag/knowledge_search",
],
)
def test_rag_agent_with_attachments(llama_stack_client_with_mocked_inference, agent_config, toolgroup):
urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]
documents = [
Document(
@ -414,14 +426,7 @@ def test_rag_agent_with_attachments(llama_stack_client_with_mocked_inference, ag
]
agent_config = {
**agent_config,
"toolgroups": [
dict(
name="builtin::rag/knowledge_search",
args={
"vector_db_ids": [],
},
)
],
"toolgroups": [toolgroup],
}
rag_agent = Agent(llama_stack_client_with_mocked_inference, agent_config)
session_id = rag_agent.create_session(f"test-session-{uuid4()}")