Skip to content

Commit 20a1ea8

Browse files
Add file_search tool invocation test with citation annotation assertions (#1136) (#1156)
* Add file_search tool invocation test with citation annotation assertions * Update test input for IBM earnings query in vector store tests Co-authored-by: Christian Zaccaria <73656840+ChristianZaccaria@users.noreply.github.com>
1 parent e60248d commit 20a1ea8

1 file changed

Lines changed: 79 additions & 0 deletions

File tree

tests/llama_stack/vector_io/test_vector_stores.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,3 +174,82 @@ def test_vector_stores_search(
174174
LOGGER.info(f"Search mode {search_mode!r}: {len(queries)} queries returned results")
175175

176176
LOGGER.info(f"Successfully tested vector store search across modes: {search_modes}")
177+
178+
@pytest.mark.smoke
179+
def test_response_file_search_tool_invocation(
180+
self,
181+
unprivileged_llama_stack_client: LlamaStackClient,
182+
llama_stack_models: ModelInfo,
183+
vector_store_with_example_docs: VectorStore,
184+
) -> None:
185+
"""
186+
Test that the file_search tool is invoked and returns results via the responses API.
187+
188+
Given: A vector store with pre-uploaded documentation files
189+
When: A question requiring document knowledge is asked with the file_search tool
190+
Then: The response contains a file_search_call output with status 'completed',
191+
results with file_id and filename, and message annotations with file citations
192+
"""
193+
response = unprivileged_llama_stack_client.responses.create(
194+
input=IBM_EARNINGS_SEARCH_QUERIES_BY_MODE["vector"][0],
195+
model=llama_stack_models.model_id,
196+
instructions="Always use the file_search tool to look up information before answering.",
197+
stream=False,
198+
tools=[
199+
{
200+
"type": "file_search",
201+
"vector_store_ids": [vector_store_with_example_docs.id],
202+
}
203+
],
204+
)
205+
206+
# Verify file_search_call output exists (model invoked the file_search tool)
207+
file_search_calls = [item for item in response.output if item.type == "file_search_call"]
208+
assert file_search_calls, (
209+
"Expected a file_search_call output item in the response, indicating the model "
210+
f"invoked the file_search tool. Output types: {[item.type for item in response.output]}"
211+
)
212+
213+
file_search_call = file_search_calls[0]
214+
assert file_search_call.status == "completed", (
215+
f"Expected file_search_call status 'completed', got '{file_search_call.status}'"
216+
)
217+
218+
# Verify file_search returned results with file metadata
219+
assert file_search_call.results, "file_search_call should contain search results"
220+
221+
for result in file_search_call.results:
222+
assert result.file_id, "Search result must include a non-empty file_id"
223+
assert result.filename, "Search result must include a non-empty filename"
224+
assert result.text, "Search result must include non-empty text content"
225+
226+
LOGGER.info(
227+
f"file_search_call returned {len(file_search_call.results)} result(s). "
228+
f"Filenames: {[r.filename for r in file_search_call.results]}"
229+
)
230+
231+
# Verify the message contains file_citation annotations
232+
annotations = []
233+
for item in response.output:
234+
if item.type != "message" or not isinstance(item.content, list):
235+
continue
236+
for content_item in item.content:
237+
if content_item.annotations:
238+
annotations.extend(content_item.annotations)
239+
240+
assert annotations, "Response message should contain file_citation annotations when file_search returns results"
241+
242+
for annotation in annotations:
243+
assert annotation.type == "file_citation", (
244+
f"Expected annotation type 'file_citation', got '{annotation.type}'"
245+
)
246+
assert annotation.file_id, "Annotation must include a non-empty file_id"
247+
assert annotation.filename, "Annotation must include a non-empty filename"
248+
assert annotation.index is not None, "Annotation must include an index"
249+
250+
LOGGER.info(
251+
f"Found {len(annotations)} file_citation annotation(s). "
252+
f"File IDs: {[a.file_id for a in annotations]}. "
253+
f"Filenames: {[a.filename for a in annotations]}. "
254+
f"Indexes: {[a.index for a in annotations]}. "
255+
)

0 commit comments

Comments
 (0)