From f0225016d6af3a31a62f31274d37b2fd224f8c59 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 11:35:26 -0400 Subject: [PATCH 01/16] implements referenced documents on /query --- src/app/endpoints/query.py | 96 +++++++++++- src/models/responses.py | 29 +++- tests/unit/app/endpoints/test_query.py | 204 +++++++++++++++++++++---- 3 files changed, 289 insertions(+), 40 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 3199a503..24334a88 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -1,11 +1,13 @@ """Handler for REST API call to provide answer to query.""" +import ast from datetime import datetime, UTC import json import logging import os from pathlib import Path -from typing import Annotated, Any +import re +from typing import Annotated, Any, cast from llama_stack_client import APIConnectionError from llama_stack_client import AsyncLlamaStackClient # type: ignore @@ -41,10 +43,79 @@ router = APIRouter(tags=["query"]) auth_dependency = get_auth_dependency() +METADATA_PATTERN = re.compile(r"\nMetadata: (\{.+})\n") + + +def _process_knowledge_search_content( + tool_response: Any, metadata_map: dict[str, dict[str, Any]] +) -> None: + """Process knowledge search tool response content for metadata.""" + for text_content_item in tool_response.content: + if not hasattr(text_content_item, "text"): + continue + + for match in METADATA_PATTERN.findall(text_content_item.text): + try: + meta = ast.literal_eval(match) + if "document_id" in meta: + metadata_map[meta["document_id"]] = meta + except Exception: # pylint: disable=broad-except + logger.debug( + "An exception was thrown in processing %s", + match, + ) + + +def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]]: + """Extract referenced documents from tool execution steps. + + Args: + steps: List of response steps from the agent + + Returns: + List of referenced documents with doc_url and doc_title + """ + metadata_map: dict[str, dict[str, Any]] = {} + + for step in steps: + if step.step_type != "tool_execution" or not hasattr(step, "tool_responses"): + continue + + for tool_response in step.tool_responses: + if ( + tool_response.tool_name != "knowledge_search" + or not tool_response.content + ): + continue + + _process_knowledge_search_content(tool_response, metadata_map) + + # Extract referenced documents from metadata + return [ + { + "doc_url": v["docs_url"], + "doc_title": v["title"], + } + for v in filter( + lambda v: ("docs_url" in v) and ("title" in v), + metadata_map.values(), + ) + ] + + query_response: dict[int | str, dict[str, Any]] = { 200: { "conversation_id": "123e4567-e89b-12d3-a456-426614174000", "response": "LLM answer", + "referenced_documents": [ + { + "doc_url": ( + "https://docs.openshift.com/container-platform/" + "4.15/operators/olm/index.html" + ), + "doc_title": "Operator Lifecycle Manager (OLM)", + } + ], }, 400: { "description": "Missing or invalid credentials provided by client", @@ -189,7 +260,7 @@ async def query_endpoint_handler( user_conversation=user_conversation, query_request=query_request ), ) - response, conversation_id = await retrieve_response( + response, conversation_id, referenced_documents = await retrieve_response( client, llama_stack_model_id, query_request, @@ -223,7 +294,11 @@ async def query_endpoint_handler( provider_id=provider_id, ) - return QueryResponse(conversation_id=conversation_id, response=response) + return QueryResponse( + conversation_id=conversation_id, + response=response, + referenced_documents=referenced_documents, + ) # connection to Llama Stack server except APIConnectionError as e: @@ -322,7 +397,7 @@ async def retrieve_response( # pylint: disable=too-many-locals query_request: QueryRequest, token: str, mcp_headers: dict[str, dict[str, str]] | None = None, -) -> tuple[str, str]: +) -> tuple[str, str, list[dict[str, str]]]: """Retrieve response from LLMs and agents.""" available_input_shields = [ shield.identifier @@ -402,7 +477,7 @@ async def retrieve_response( # pylint: disable=too-many-locals toolgroups=toolgroups, ) - # Check for validation errors in the response + # Check for validation errors and extract referenced documents steps = getattr(response, "steps", []) for step in steps: if step.step_type == "shield_call" and step.violation: @@ -410,7 +485,16 @@ async def retrieve_response( # pylint: disable=too-many-locals metrics.llm_calls_validation_errors_total.inc() break - return str(response.output_message.content), conversation_id # type: ignore[union-attr] + # Extract referenced documents from tool execution steps + referenced_documents = extract_referenced_documents_from_steps(steps) + + # When stream=False, response should have output_message attribute + response_obj = cast(Any, response) + return ( + str(response_obj.output_message.content), + conversation_id, + referenced_documents, + ) def validate_attachments_metadata(attachments: list[Attachment]) -> None: diff --git a/src/models/responses.py b/src/models/responses.py index cb8ee09c..e2a1d891 100644 --- a/src/models/responses.py +++ b/src/models/responses.py @@ -36,8 +36,6 @@ class ModelsResponse(BaseModel): # TODO(lucasagomes): a lot of fields to add to QueryResponse. For now # we are keeping it simple. The missing fields are: -# - referenced_documents: The optional URLs and titles for the documents used -# to generate the response. # - truncated: Set to True if conversation history was truncated to be within context window. # - input_tokens: Number of tokens sent to LLM # - output_tokens: Number of tokens received from LLM @@ -51,6 +49,8 @@ class QueryResponse(BaseModel): Attributes: conversation_id: The optional conversation ID (UUID). response: The response. + referenced_documents: The optional URLs and titles for the documents used + to generate the response. """ conversation_id: Optional[str] = Field( @@ -66,6 +66,22 @@ class QueryResponse(BaseModel): ], ) + referenced_documents: list[dict[str, str]] = Field( + default_factory=list, + description="List of documents referenced in generating the response", + examples=[ + [ + { + "doc_url": ( + "https://docs.openshift.com/container-platform/" + "4.15/operators/olm/index.html" + ), + "doc_title": "Operator Lifecycle Manager (OLM)", + } + ] + ], + ) + # provides examples for /docs endpoint model_config = { "json_schema_extra": { @@ -73,6 +89,15 @@ class QueryResponse(BaseModel): { "conversation_id": "123e4567-e89b-12d3-a456-426614174000", "response": "Operator Lifecycle Manager (OLM) helps users install...", + "referenced_documents": [ + { + "doc_url": ( + "https://docs.openshift.com/container-platform/" + "4.15/operators/olm/index.html" + ), + "doc_title": "Operator Lifecycle Manager (OLM)", + } + ], } ] } diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 1f52260a..dfe4a5a1 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -8,6 +8,8 @@ from llama_stack_client import APIConnectionError from llama_stack_client.types import UserMessage # type: ignore +from llama_stack_client.types.shared.interleaved_content_item import TextContentItem +from llama_stack_client.types.tool_response import ToolResponse from configuration import AppConfig from app.endpoints.query import ( @@ -28,6 +30,34 @@ MOCK_AUTH = ("mock_user_id", "mock_username", "mock_token") +SAMPLE_KNOWLEDGE_SEARCH_RESULTS = [ + """knowledge_search tool found 2 chunks: +BEGIN of knowledge_search tool results. +""", + """Result 1 +Content: ABC +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1', \ +'source': None} +""", + """Result 2 +Content: ABC +Metadata: {'docs_url': 'https://example.com/doc2', 'title': 'Doc2', 'document_id': 'doc-2', \ +'source': None} +""", + """END of knowledge_search tool results. +""", + # Following metadata contains an intentionally incorrect keyword "Title" (instead of "title") + # and it is not picked as a referenced document. + """Result 3 +Content: ABC +Metadata: {'docs_url': 'https://example.com/doc3', 'Title': 'Doc3', 'document_id': 'doc-3', \ +'source': None} +""", + """The above results were retrieved to help answer the user\'s query: "Sample Query". +Use them as supporting information only in answering this query. +""", +] + def mock_database_operations(mocker): """Helper function to mock database operations for query endpoints.""" @@ -128,7 +158,7 @@ async def _test_query_endpoint_handler(mocker, store_transcript_to_file=False): mocker.patch( "app.endpoints.query.retrieve_response", - return_value=(llm_response, conversation_id), + return_value=(llm_response, conversation_id, []), ) mocker.patch( "app.endpoints.query.select_model_and_provider_id", @@ -185,6 +215,59 @@ async def test_query_endpoint_handler_store_transcript(mocker): await _test_query_endpoint_handler(mocker, store_transcript_to_file=True) +@pytest.mark.asyncio +async def test_query_endpoint_handler_with_referenced_documents(mocker): + """Test the query endpoint handler returns referenced documents.""" + mock_metric = mocker.patch("metrics.llm_calls_total") + mock_client = mocker.AsyncMock() + mock_lsc = mocker.patch("client.AsyncLlamaStackClientHolder.get_client") + mock_lsc.return_value = mock_client + mock_client.models.list.return_value = [ + mocker.Mock(identifier="model1", model_type="llm", provider_id="provider1"), + mocker.Mock(identifier="model2", model_type="llm", provider_id="provider2"), + ] + + mock_config = mocker.Mock() + mock_config.user_data_collection_configuration.transcripts_enabled = False + mocker.patch("app.endpoints.query.configuration", mock_config) + + llm_response = "LLM answer with referenced documents" + conversation_id = "fake_conversation_id" + referenced_documents = [ + {"doc_url": "https://example.com/doc1", "doc_title": "Doc1"}, + {"doc_url": "https://example.com/doc2", "doc_title": "Doc2"}, + ] + query = "What is OpenStack?" + + mocker.patch( + "app.endpoints.query.retrieve_response", + return_value=(llm_response, conversation_id, referenced_documents), + ) + mocker.patch( + "app.endpoints.query.select_model_and_provider_id", + return_value=("fake_model_id", "fake_model_id", "fake_provider_id"), + ) + mocker.patch("app.endpoints.query.is_transcripts_enabled", return_value=False) + + # Mock database operations + mock_database_operations(mocker) + + query_request = QueryRequest(query=query) + + response = await query_endpoint_handler(query_request, auth=MOCK_AUTH) + + # Assert the response contains referenced documents + assert response.response == llm_response + assert response.conversation_id == conversation_id + assert response.referenced_documents == referenced_documents + assert len(response.referenced_documents) == 2 + assert response.referenced_documents[0]["doc_title"] == "Doc1" + assert response.referenced_documents[1]["doc_title"] == "Doc2" + + # Assert the metric for successful LLM calls is incremented + mock_metric.labels("fake_provider_id", "fake_model_id").inc.assert_called_once() + + def test_select_model_and_provider_id_from_request(mocker): """Test the select_model_and_provider_id function.""" mocker.patch( @@ -408,7 +491,7 @@ async def test_retrieve_response_vector_db_available(prepare_agent_mocks, mocker model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, referenced_documents = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -416,6 +499,7 @@ async def test_retrieve_response_vector_db_available(prepare_agent_mocks, mocker mock_metric.inc.assert_not_called() assert response == "LLM answer" assert conversation_id == "fake_conversation_id" + assert referenced_documents == [] # No knowledge_search in this test, so empty list mock_agent.create_turn.assert_called_once_with( messages=[UserMessage(content="What is OpenStack?", role="user")], session_id="fake_session_id", @@ -425,6 +509,64 @@ async def test_retrieve_response_vector_db_available(prepare_agent_mocks, mocker ) +@pytest.mark.asyncio +async def test_retrieve_response_with_knowledge_search_extracts_referenced_documents( + prepare_agent_mocks, mocker +): + """Test the retrieve_response function extracts referenced documents from knowledge_search.""" + mock_client, mock_agent = prepare_agent_mocks + mock_agent.create_turn.return_value.output_message.content = "LLM answer" + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock the response with tool execution steps containing knowledge_search results + mock_tool_response = ToolResponse( + call_id="c1", + tool_name="knowledge_search", + content=[ + TextContentItem(text=s, type="text") + for s in SAMPLE_KNOWLEDGE_SEARCH_RESULTS + ], + ) + + mock_tool_execution_step = mocker.Mock() + mock_tool_execution_step.step_type = "tool_execution" + mock_tool_execution_step.tool_responses = [mock_tool_response] + + mock_agent.create_turn.return_value.steps = [mock_tool_execution_step] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, referenced_documents = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + assert response == "LLM answer" + assert conversation_id == "fake_conversation_id" + + # Assert referenced documents were extracted correctly + assert len(referenced_documents) == 2 + assert referenced_documents[0]["doc_url"] == "https://example.com/doc1" + assert referenced_documents[0]["doc_title"] == "Doc1" + assert referenced_documents[1]["doc_url"] == "https://example.com/doc2" + assert referenced_documents[1]["doc_title"] == "Doc2" + + # Doc3 should not be included because it has "Title" instead of "title" + doc_titles = [doc["doc_title"] for doc in referenced_documents] + assert "Doc3" not in doc_titles + + @pytest.mark.asyncio async def test_retrieve_response_no_available_shields(prepare_agent_mocks, mocker): """Test the retrieve_response function.""" @@ -446,7 +588,7 @@ async def test_retrieve_response_no_available_shields(prepare_agent_mocks, mocke model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -495,7 +637,7 @@ def __repr__(self): model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -547,7 +689,7 @@ def __repr__(self): model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -601,7 +743,7 @@ def __repr__(self): model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -657,7 +799,7 @@ async def test_retrieve_response_with_one_attachment(prepare_agent_mocks, mocker model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -711,7 +853,7 @@ async def test_retrieve_response_with_two_attachments(prepare_agent_mocks, mocke model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -766,7 +908,7 @@ async def test_retrieve_response_with_mcp_servers(prepare_agent_mocks, mocker): model_id = "fake_model_id" access_token = "test_token_123" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -835,7 +977,7 @@ async def test_retrieve_response_with_mcp_servers_empty_token( model_id = "fake_model_id" access_token = "" # Empty token - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -893,8 +1035,6 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( ) query_request = QueryRequest(query="What is OpenStack?") - model_id = "fake_model_id" - access_token = "" mcp_headers = { "filesystem-server": {"Authorization": "Bearer test_token_123"}, "git-server": {"Authorization": "Bearer test_token_456"}, @@ -906,11 +1046,11 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( }, } - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, - model_id, + "fake_model_id", query_request, - access_token, + "", mcp_headers=mcp_headers, ) @@ -920,7 +1060,7 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( # Verify get_agent was called with the correct parameters mock_get_agent.assert_called_once_with( mock_client, - model_id, + "fake_model_id", mocker.ANY, # system_prompt [], # available_input_shields [], # available_output_shields @@ -928,20 +1068,20 @@ async def test_retrieve_response_with_mcp_servers_and_mcp_headers( False, # no_tools ) - expected_mcp_headers = { - "http://localhost:3000": {"Authorization": "Bearer test_token_123"}, - "https://git.example.com/mcp": {"Authorization": "Bearer test_token_456"}, - "http://another-server-mcp-server:3000": { - "Authorization": "Bearer test_token_789" - }, - # we do not put "unknown-mcp-server" url as it's unknown to lightspeed-stack - } - # Check that the agent's extra_headers property was set correctly expected_extra_headers = { "X-LlamaStack-Provider-Data": json.dumps( { - "mcp_headers": expected_mcp_headers, + "mcp_headers": { + "http://localhost:3000": {"Authorization": "Bearer test_token_123"}, + "https://git.example.com/mcp": { + "Authorization": "Bearer test_token_456" + }, + "http://another-server-mcp-server:3000": { + "Authorization": "Bearer test_token_789" + }, + # we do not put "unknown-mcp-server" url as it's unknown to lightspeed-stack + }, } ) } @@ -987,7 +1127,7 @@ async def test_retrieve_response_shield_violation(prepare_agent_mocks, mocker): query_request = QueryRequest(query="What is OpenStack?") - _, conversation_id = await retrieve_response( + _, conversation_id, _ = await retrieve_response( mock_client, "fake_model_id", query_request, "test_token" ) @@ -1139,7 +1279,7 @@ async def test_auth_tuple_unpacking_in_query_endpoint_handler(mocker): mock_retrieve_response = mocker.patch( "app.endpoints.query.retrieve_response", - return_value=("test response", "test_conversation_id"), + return_value=("test response", "test_conversation_id", []), ) mocker.patch( @@ -1179,7 +1319,7 @@ async def test_query_endpoint_handler_no_tools_true(mocker): mocker.patch( "app.endpoints.query.retrieve_response", - return_value=(llm_response, conversation_id), + return_value=(llm_response, conversation_id, []), ) mocker.patch( "app.endpoints.query.select_model_and_provider_id", @@ -1218,7 +1358,7 @@ async def test_query_endpoint_handler_no_tools_false(mocker): mocker.patch( "app.endpoints.query.retrieve_response", - return_value=(llm_response, conversation_id), + return_value=(llm_response, conversation_id, []), ) mocker.patch( "app.endpoints.query.select_model_and_provider_id", @@ -1267,7 +1407,7 @@ async def test_retrieve_response_no_tools_bypasses_mcp_and_rag( model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) @@ -1317,7 +1457,7 @@ async def test_retrieve_response_no_tools_false_preserves_functionality( model_id = "fake_model_id" access_token = "test_token" - response, conversation_id = await retrieve_response( + response, conversation_id, _ = await retrieve_response( mock_client, model_id, query_request, access_token ) From 8c0a1d0af7fffbd941b014fdda93400c0cfc621e Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 13:29:21 -0400 Subject: [PATCH 02/16] code rabbit fixes and additional tests --- src/app/endpoints/query.py | 25 ++--- tests/unit/app/endpoints/test_query.py | 139 +++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 17 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 24334a88..e77f85da 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -43,7 +43,7 @@ router = APIRouter(tags=["query"]) auth_dependency = get_auth_dependency() -METADATA_PATTERN = re.compile(r"\nMetadata: (\{.+})\n") +METADATA_PATTERN = re.compile(r"^Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) def _process_knowledge_search_content( @@ -78,28 +78,20 @@ def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]] metadata_map: dict[str, dict[str, Any]] = {} for step in steps: - if step.step_type != "tool_execution" or not hasattr(step, "tool_responses"): + if getattr(step, "step_type", "") != "tool_execution" or not hasattr(step, "tool_responses"): continue - for tool_response in step.tool_responses: - if ( - tool_response.tool_name != "knowledge_search" - or not tool_response.content - ): + for tool_response in getattr(step, "tool_responses", []) or []: + if getattr(tool_response, "tool_name", "") != "knowledge_search" or not getattr(tool_response, "content", []): continue _process_knowledge_search_content(tool_response, metadata_map) # Extract referenced documents from metadata return [ - { - "doc_url": v["docs_url"], - "doc_title": v["title"], - } - for v in filter( - lambda v: ("docs_url" in v) and ("title" in v), - metadata_map.values(), - ) + {"doc_url": v["docs_url"], "doc_title": v["title"]} + for v in metadata_map.values() + if "docs_url" in v and "title" in v ] @@ -480,10 +472,9 @@ async def retrieve_response( # pylint: disable=too-many-locals # Check for validation errors and extract referenced documents steps = getattr(response, "steps", []) for step in steps: - if step.step_type == "shield_call" and step.violation: + if getattr(step, "step_type", "") == "shield_call" and getattr(step, "violation", False): # Metric for LLM validation errors metrics.llm_calls_validation_errors_total.inc() - break # Extract referenced documents from tool execution steps referenced_documents = extract_referenced_documents_from_steps(steps) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index dfe4a5a1..561e2854 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -22,6 +22,7 @@ store_transcript, get_rag_toolgroups, evaluate_model_hints, + _process_knowledge_search_content, ) from models.requests import QueryRequest, Attachment @@ -1583,3 +1584,141 @@ def test_evaluate_model_hints( assert provider_id == expected_provider assert model_id == expected_model + + +def test_process_knowledge_search_content_with_valid_metadata(mocker): + """Test _process_knowledge_search_content with valid metadata.""" + # Mock tool response with valid metadata + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Test Doc', 'document_id': 'doc-1'} +""" + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata was correctly parsed and added + assert "doc-1" in metadata_map + assert metadata_map["doc-1"]["docs_url"] == "https://example.com/doc1" + assert metadata_map["doc-1"]["title"] == "Test Doc" + assert metadata_map["doc-1"]["document_id"] == "doc-1" + + +def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): + """Test _process_knowledge_search_content handles SyntaxError from invalid metadata.""" + mock_logger = mocker.patch("app.endpoints.query.logger") + + # Mock tool response with invalid metadata (invalid Python syntax) + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +Metadata: {'docs_url': 'https://example.com/doc1' 'title': 'Test Doc', 'document_id': 'doc-1'} +""" # Missing comma between 'doc1' and 'title' - will cause SyntaxError + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty due to exception + assert len(metadata_map) == 0 + + # Verify debug logging was called + mock_logger.debug.assert_called_once() + args = mock_logger.debug.call_args[0] + assert "An exception was thrown in processing" in args[0] + + +def test_process_knowledge_search_content_with_invalid_metadata_value_error(mocker): + """Test _process_knowledge_search_content handles ValueError from invalid metadata.""" + mock_logger = mocker.patch("app.endpoints.query.logger") + + # Mock tool response with invalid metadata containing complex expressions + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +Metadata: {func_call(): 'value', 'title': 'Test Doc', 'document_id': 'doc-1'} +""" # Function call in dict - will cause ValueError since it's not a literal + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty due to exception + assert len(metadata_map) == 0 + + # Verify debug logging was called + mock_logger.debug.assert_called_once() + args = mock_logger.debug.call_args[0] + assert "An exception was thrown in processing" in args[0] + + +def test_process_knowledge_search_content_with_non_dict_metadata(mocker): + """Test _process_knowledge_search_content handles non-dict metadata gracefully.""" + mock_logger = mocker.patch("app.endpoints.query.logger") + + # Mock tool response with metadata that's not a dict + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +Metadata: "just a string" +""" + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty (no document_id in string) + assert len(metadata_map) == 0 + + # No exception should be logged since string is valid literal + mock_logger.debug.assert_not_called() + + +def test_process_knowledge_search_content_with_metadata_missing_document_id(mocker): + """Test _process_knowledge_search_content skips metadata without document_id.""" + # Mock tool response with valid metadata but missing document_id + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Test Doc'} +""" # No document_id field + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty since document_id is missing + assert len(metadata_map) == 0 + + +def test_process_knowledge_search_content_with_no_text_attribute(mocker): + """Test _process_knowledge_search_content skips content items without text attribute.""" + # Mock tool response with content item that has no text attribute + text_content_item = mocker.Mock(spec=[]) # spec=[] means no attributes + + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty since text attribute is missing + assert len(metadata_map) == 0 From d6e147537d905b8bbc10e24aa5f5b23c8a8e81c2 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 13:31:34 -0400 Subject: [PATCH 03/16] formatting --- src/app/endpoints/query.py | 12 ++++-- tests/unit/app/endpoints/test_query.py | 60 +++++++++++++------------- 2 files changed, 39 insertions(+), 33 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index e77f85da..1d77764a 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -78,11 +78,15 @@ def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]] metadata_map: dict[str, dict[str, Any]] = {} for step in steps: - if getattr(step, "step_type", "") != "tool_execution" or not hasattr(step, "tool_responses"): + if getattr(step, "step_type", "") != "tool_execution" or not hasattr( + step, "tool_responses" + ): continue for tool_response in getattr(step, "tool_responses", []) or []: - if getattr(tool_response, "tool_name", "") != "knowledge_search" or not getattr(tool_response, "content", []): + if getattr( + tool_response, "tool_name", "" + ) != "knowledge_search" or not getattr(tool_response, "content", []): continue _process_knowledge_search_content(tool_response, metadata_map) @@ -472,7 +476,9 @@ async def retrieve_response( # pylint: disable=too-many-locals # Check for validation errors and extract referenced documents steps = getattr(response, "steps", []) for step in steps: - if getattr(step, "step_type", "") == "shield_call" and getattr(step, "violation", False): + if getattr(step, "step_type", "") == "shield_call" and getattr( + step, "violation", False + ): # Metric for LLM validation errors metrics.llm_calls_validation_errors_total.inc() diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 561e2854..3badae51 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -1594,14 +1594,14 @@ def test_process_knowledge_search_content_with_valid_metadata(mocker): Content: Test content Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Test Doc', 'document_id': 'doc-1'} """ - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata was correctly parsed and added assert "doc-1" in metadata_map assert metadata_map["doc-1"]["docs_url"] == "https://example.com/doc1" @@ -1612,24 +1612,24 @@ def test_process_knowledge_search_content_with_valid_metadata(mocker): def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): """Test _process_knowledge_search_content handles SyntaxError from invalid metadata.""" mock_logger = mocker.patch("app.endpoints.query.logger") - + # Mock tool response with invalid metadata (invalid Python syntax) text_content_item = mocker.Mock() text_content_item.text = """Result 1 Content: Test content Metadata: {'docs_url': 'https://example.com/doc1' 'title': 'Test Doc', 'document_id': 'doc-1'} """ # Missing comma between 'doc1' and 'title' - will cause SyntaxError - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 - + # Verify debug logging was called mock_logger.debug.assert_called_once() args = mock_logger.debug.call_args[0] @@ -1639,24 +1639,24 @@ def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(moc def test_process_knowledge_search_content_with_invalid_metadata_value_error(mocker): """Test _process_knowledge_search_content handles ValueError from invalid metadata.""" mock_logger = mocker.patch("app.endpoints.query.logger") - + # Mock tool response with invalid metadata containing complex expressions text_content_item = mocker.Mock() text_content_item.text = """Result 1 Content: Test content Metadata: {func_call(): 'value', 'title': 'Test Doc', 'document_id': 'doc-1'} """ # Function call in dict - will cause ValueError since it's not a literal - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 - + # Verify debug logging was called mock_logger.debug.assert_called_once() args = mock_logger.debug.call_args[0] @@ -1666,24 +1666,24 @@ def test_process_knowledge_search_content_with_invalid_metadata_value_error(mock def test_process_knowledge_search_content_with_non_dict_metadata(mocker): """Test _process_knowledge_search_content handles non-dict metadata gracefully.""" mock_logger = mocker.patch("app.endpoints.query.logger") - + # Mock tool response with metadata that's not a dict text_content_item = mocker.Mock() text_content_item.text = """Result 1 Content: Test content Metadata: "just a string" """ - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata_map remains empty (no document_id in string) assert len(metadata_map) == 0 - + # No exception should be logged since string is valid literal mock_logger.debug.assert_not_called() @@ -1696,14 +1696,14 @@ def test_process_knowledge_search_content_with_metadata_missing_document_id(mock Content: Test content Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Test Doc'} """ # No document_id field - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata_map remains empty since document_id is missing assert len(metadata_map) == 0 @@ -1712,13 +1712,13 @@ def test_process_knowledge_search_content_with_no_text_attribute(mocker): """Test _process_knowledge_search_content skips content items without text attribute.""" # Mock tool response with content item that has no text attribute text_content_item = mocker.Mock(spec=[]) # spec=[] means no attributes - + tool_response = mocker.Mock() tool_response.content = [text_content_item] - + metadata_map = {} - + _process_knowledge_search_content(tool_response, metadata_map) - + # Verify metadata_map remains empty since text attribute is missing assert len(metadata_map) == 0 From 866bb5c64692a67c525b57bc8a11e219da5c4f6c Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 14:25:22 -0400 Subject: [PATCH 04/16] more coderabbit changes --- src/app/endpoints/query.py | 11 +- tests/unit/app/endpoints/test_query.py | 173 +++++++++++++++++++++++-- 2 files changed, 169 insertions(+), 15 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 1d77764a..16c64247 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -43,7 +43,7 @@ router = APIRouter(tags=["query"]) auth_dependency = get_auth_dependency() -METADATA_PATTERN = re.compile(r"^Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) +METADATA_PATTERN = re.compile(r"^\s*Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) def _process_knowledge_search_content( @@ -59,10 +59,11 @@ def _process_knowledge_search_content( meta = ast.literal_eval(match) if "document_id" in meta: metadata_map[meta["document_id"]] = meta - except Exception: # pylint: disable=broad-except + except (SyntaxError, ValueError): # only expected from literal_eval logger.debug( "An exception was thrown in processing %s", match, + exc_info=True, ) @@ -121,7 +122,7 @@ def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]] "description": "User is not authorized", "model": ForbiddenResponse, }, - 503: { + 500: { "detail": { "response": "Unable to connect to Llama Stack", "cause": "Connection error.", @@ -487,8 +488,10 @@ async def retrieve_response( # pylint: disable=too-many-locals # When stream=False, response should have output_message attribute response_obj = cast(Any, response) + content = getattr(getattr(response_obj, "output_message", None), "content", "") + content_str = "" if content is None else str(content) return ( - str(response_obj.output_message.content), + content_str, conversation_id, referenced_documents, ) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 3badae51..4522ce7a 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -44,6 +44,11 @@ Content: ABC Metadata: {'docs_url': 'https://example.com/doc2', 'title': 'Doc2', 'document_id': 'doc-2', \ 'source': None} +""", + """Result 2b +Content: ABC + Metadata: {'docs_url': 'https://example.com/doc2b', 'title': 'Doc2b', 'document_id': 'doc-2b', \ +'source': None} """, """END of knowledge_search tool results. """, @@ -101,10 +106,6 @@ def setup_configuration_fixture(): async def test_query_endpoint_handler_configuration_not_loaded(mocker): """Test the query endpoint handler if configuration is not loaded.""" # simulate state when no configuration is loaded - mocker.patch( - "app.endpoints.query.configuration", - return_value=mocker.Mock(), - ) mocker.patch("app.endpoints.query.configuration", None) query = "What is OpenStack?" @@ -272,12 +273,11 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): def test_select_model_and_provider_id_from_request(mocker): """Test the select_model_and_provider_id function.""" mocker.patch( - "metrics.utils.configuration.inference.default_provider", + "app.endpoints.query.configuration.inference.default_provider", "default_provider", ) mocker.patch( - "metrics.utils.configuration.inference.default_model", - "default_model", + "app.endpoints.query.configuration.inference.default_model", "default_model" ) model_list = [ @@ -312,12 +312,11 @@ def test_select_model_and_provider_id_from_request(mocker): def test_select_model_and_provider_id_from_configuration(mocker): """Test the select_model_and_provider_id function.""" mocker.patch( - "metrics.utils.configuration.inference.default_provider", + "app.endpoints.query.configuration.inference.default_provider", "default_provider", ) mocker.patch( - "metrics.utils.configuration.inference.default_model", - "default_model", + "app.endpoints.query.configuration.inference.default_model", "default_model" ) model_list = [ @@ -557,11 +556,13 @@ async def test_retrieve_response_with_knowledge_search_extracts_referenced_docum assert conversation_id == "fake_conversation_id" # Assert referenced documents were extracted correctly - assert len(referenced_documents) == 2 + assert len(referenced_documents) == 3 assert referenced_documents[0]["doc_url"] == "https://example.com/doc1" assert referenced_documents[0]["doc_title"] == "Doc1" assert referenced_documents[1]["doc_url"] == "https://example.com/doc2" assert referenced_documents[1]["doc_title"] == "Doc2" + assert referenced_documents[2]["doc_url"] == "https://example.com/doc2b" + assert referenced_documents[2]["doc_title"] == "Doc2b" # Doc3 should not be included because it has "Title" instead of "title" doc_titles = [doc["doc_title"] for doc in referenced_documents] @@ -1722,3 +1723,153 @@ def test_process_knowledge_search_content_with_no_text_attribute(mocker): # Verify metadata_map remains empty since text attribute is missing assert len(metadata_map) == 0 + + +@pytest.mark.asyncio +async def test_retrieve_response_with_none_content(prepare_agent_mocks, mocker): + """Test retrieve_response handles None content gracefully.""" + mock_client, mock_agent = prepare_agent_mocks + + # Mock response with None content + mock_response = mocker.Mock() + mock_response.output_message.content = None + mock_response.steps = [] + mock_agent.create_turn.return_value = mock_response + + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, _ = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + # Should return empty string instead of "None" + assert response == "" + assert conversation_id == "fake_conversation_id" + + +@pytest.mark.asyncio +async def test_retrieve_response_with_missing_output_message( + prepare_agent_mocks, mocker +): + """Test retrieve_response handles missing output_message gracefully.""" + mock_client, mock_agent = prepare_agent_mocks + + # Mock response without output_message attribute + mock_response = mocker.Mock(spec=["steps"]) # Only has steps attribute + mock_response.steps = [] + mock_agent.create_turn.return_value = mock_response + + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, _ = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + # Should return empty string when output_message is missing + assert response == "" + assert conversation_id == "fake_conversation_id" + + +@pytest.mark.asyncio +async def test_retrieve_response_with_missing_content_attribute( + prepare_agent_mocks, mocker +): + """Test retrieve_response handles missing content attribute gracefully.""" + mock_client, mock_agent = prepare_agent_mocks + + # Mock response with output_message but no content attribute + mock_response = mocker.Mock() + mock_response.output_message = mocker.Mock(spec=[]) # No content attribute + mock_response.steps = [] + mock_agent.create_turn.return_value = mock_response + + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, _ = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + # Should return empty string when content attribute is missing + assert response == "" + assert conversation_id == "fake_conversation_id" + + +@pytest.mark.asyncio +async def test_retrieve_response_with_structured_content_object( + prepare_agent_mocks, mocker +): + """Test retrieve_response handles structured content objects properly.""" + mock_client, mock_agent = prepare_agent_mocks + + # Mock response with a structured content object + structured_content = {"type": "text", "value": "This is structured content"} + mock_response = mocker.Mock() + mock_response.output_message.content = structured_content + mock_response.steps = [] + mock_agent.create_turn.return_value = mock_response + + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, _ = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + # Should convert the structured object to string representation + assert response == str(structured_content) + assert conversation_id == "fake_conversation_id" From 30263972610afabaf013f900f3b4286b87761031 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 15:59:11 -0400 Subject: [PATCH 05/16] coderabbit changes also implements the pydantic object for streaming_query since it wouldn't make sense to behave differently across endpoints --- src/app/endpoints/query.py | 53 ++++++--- src/app/endpoints/streaming_query.py | 38 +++--- src/models/responses.py | 13 ++- tests/unit/app/endpoints/test_query.py | 49 +++++--- .../app/endpoints/test_streaming_query.py | 108 ++++++++++++++++++ 5 files changed, 216 insertions(+), 45 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 16c64247..4e63b418 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -27,7 +27,12 @@ from app.database import get_session import metrics from models.database.conversations import UserConversation -from models.responses import QueryResponse, UnauthorizedResponse, ForbiddenResponse +from models.responses import ( + QueryResponse, + UnauthorizedResponse, + ForbiddenResponse, + ReferencedDocument, +) from models.requests import QueryRequest, Attachment import constants from utils.endpoints import ( @@ -50,24 +55,40 @@ def _process_knowledge_search_content( tool_response: Any, metadata_map: dict[str, dict[str, Any]] ) -> None: """Process knowledge search tool response content for metadata.""" - for text_content_item in tool_response.content: - if not hasattr(text_content_item, "text"): + # Guard against missing tool_response or content + if not tool_response: + return + + content = getattr(tool_response, "content", None) + if not content: + return + + # Ensure content is iterable + try: + iter(content) + except TypeError: + return + + for text_content_item in content: + # Skip items that lack a non-empty "text" attribute + text = getattr(text_content_item, "text", None) + if not text: continue - for match in METADATA_PATTERN.findall(text_content_item.text): + for match in METADATA_PATTERN.findall(text): try: meta = ast.literal_eval(match) - if "document_id" in meta: + # Verify the result is a dict before accessing keys + if isinstance(meta, dict) and "document_id" in meta: metadata_map[meta["document_id"]] = meta except (SyntaxError, ValueError): # only expected from literal_eval - logger.debug( + logger.exception( "An exception was thrown in processing %s", match, - exc_info=True, ) -def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]]: +def extract_referenced_documents_from_steps(steps: list) -> list[ReferencedDocument]: """Extract referenced documents from tool execution steps. Args: @@ -94,7 +115,7 @@ def extract_referenced_documents_from_steps(steps: list) -> list[dict[str, str]] # Extract referenced documents from metadata return [ - {"doc_url": v["docs_url"], "doc_title": v["title"]} + ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) for v in metadata_map.values() if "docs_url" in v and "title" in v ] @@ -388,13 +409,13 @@ def is_input_shield(shield: Shield) -> bool: return _is_inout_shield(shield) or not is_output_shield(shield) -async def retrieve_response( # pylint: disable=too-many-locals +async def retrieve_response( # pylint: disable=too-many-locals,too-many-branches client: AsyncLlamaStackClient, model_id: str, query_request: QueryRequest, token: str, mcp_headers: dict[str, dict[str, str]] | None = None, -) -> tuple[str, str, list[dict[str, str]]]: +) -> tuple[str, str, list[ReferencedDocument]]: """Retrieve response from LLMs and agents.""" available_input_shields = [ shield.identifier @@ -488,8 +509,14 @@ async def retrieve_response( # pylint: disable=too-many-locals # When stream=False, response should have output_message attribute response_obj = cast(Any, response) - content = getattr(getattr(response_obj, "output_message", None), "content", "") - content_str = "" if content is None else str(content) + + # Safely guard access to output_message and content + output_message = getattr(response_obj, "output_message", None) + if output_message and getattr(output_message, "content", None) is not None: + content_str = str(output_message.content) + else: + content_str = "" + return ( content_str, conversation_id, diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 8b00ef61..2f875dca 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -24,6 +24,7 @@ import metrics from models.requests import QueryRequest from models.database.conversations import UserConversation +from models.responses import ReferencedDocument from utils.endpoints import check_configuration_loaded, get_agent, get_system_prompt from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups @@ -72,20 +73,27 @@ def stream_start_event(conversation_id: str) -> str: def stream_end_event(metadata_map: dict) -> str: """Yield the end of the data stream.""" + # Create ReferencedDocument objects and convert them to serializable dict format + referenced_documents = [] + for v in filter( + lambda v: ("docs_url" in v) and ("title" in v), + metadata_map.values(), + ): + doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) + referenced_documents.append( + { + "doc_url": str( + doc.doc_url + ), # Convert AnyUrl to string for JSON serialization + "doc_title": doc.doc_title, + } + ) + return format_stream_data( { "event": "end", "data": { - "referenced_documents": [ - { - "doc_url": v["docs_url"], - "doc_title": v["title"], - } - for v in filter( - lambda v: ("docs_url" in v) and ("title" in v), - metadata_map.values(), - ) - ], + "referenced_documents": referenced_documents, "truncated": None, # TODO(jboos): implement truncated "input_tokens": 0, # TODO(jboos): implement input tokens "output_tokens": 0, # TODO(jboos): implement output tokens @@ -330,10 +338,14 @@ def _handle_tool_execution_event( for match in METADATA_PATTERN.findall(text_content_item.text): try: meta = ast.literal_eval(match) - if "document_id" in meta: + # Verify the result is a dict before accessing keys + if isinstance(meta, dict) and "document_id" in meta: metadata_map[meta["document_id"]] = meta - except Exception: # pylint: disable=broad-except - logger.debug( + except ( + SyntaxError, + ValueError, + ): # only expected from literal_eval + logger.exception( "An exception was thrown in processing %s", match, ) diff --git a/src/models/responses.py b/src/models/responses.py index e2a1d891..9cf0c0d3 100644 --- a/src/models/responses.py +++ b/src/models/responses.py @@ -2,7 +2,7 @@ from typing import Any, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, AnyUrl class ModelsResponse(BaseModel): @@ -43,6 +43,15 @@ class ModelsResponse(BaseModel): # - tool_calls: List of tool requests. # - tool_results: List of tool results. # See LLMResponse in ols-service for more details. + + +class ReferencedDocument(BaseModel): + """Model representing a document referenced in generating a response.""" + + doc_url: AnyUrl = Field(description="URL of the referenced document") + doc_title: str = Field(description="Title of the referenced document") + + class QueryResponse(BaseModel): """Model representing LLM response to a query. @@ -66,7 +75,7 @@ class QueryResponse(BaseModel): ], ) - referenced_documents: list[dict[str, str]] = Field( + referenced_documents: list[ReferencedDocument] = Field( default_factory=list, description="List of documents referenced in generating the response", examples=[ diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 4522ce7a..6ce01eba 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -27,6 +27,7 @@ from models.requests import QueryRequest, Attachment from models.config import ModelContextProtocolServer +from models.responses import ReferencedDocument from models.database.conversations import UserConversation MOCK_AUTH = ("mock_user_id", "mock_username", "mock_token") @@ -236,8 +237,8 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): llm_response = "LLM answer with referenced documents" conversation_id = "fake_conversation_id" referenced_documents = [ - {"doc_url": "https://example.com/doc1", "doc_title": "Doc1"}, - {"doc_url": "https://example.com/doc2", "doc_title": "Doc2"}, + ReferencedDocument(doc_url="https://example.com/doc1", doc_title="Doc1"), + ReferencedDocument(doc_url="https://example.com/doc2", doc_title="Doc2"), ] query = "What is OpenStack?" @@ -263,8 +264,8 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): assert response.conversation_id == conversation_id assert response.referenced_documents == referenced_documents assert len(response.referenced_documents) == 2 - assert response.referenced_documents[0]["doc_title"] == "Doc1" - assert response.referenced_documents[1]["doc_title"] == "Doc2" + assert response.referenced_documents[0].doc_title == "Doc1" + assert response.referenced_documents[1].doc_title == "Doc2" # Assert the metric for successful LLM calls is incremented mock_metric.labels("fake_provider_id", "fake_model_id").inc.assert_called_once() @@ -557,15 +558,15 @@ async def test_retrieve_response_with_knowledge_search_extracts_referenced_docum # Assert referenced documents were extracted correctly assert len(referenced_documents) == 3 - assert referenced_documents[0]["doc_url"] == "https://example.com/doc1" - assert referenced_documents[0]["doc_title"] == "Doc1" - assert referenced_documents[1]["doc_url"] == "https://example.com/doc2" - assert referenced_documents[1]["doc_title"] == "Doc2" - assert referenced_documents[2]["doc_url"] == "https://example.com/doc2b" - assert referenced_documents[2]["doc_title"] == "Doc2b" + assert str(referenced_documents[0].doc_url) == "https://example.com/doc1" + assert referenced_documents[0].doc_title == "Doc1" + assert str(referenced_documents[1].doc_url) == "https://example.com/doc2" + assert referenced_documents[1].doc_title == "Doc2" + assert str(referenced_documents[2].doc_url) == "https://example.com/doc2b" + assert referenced_documents[2].doc_title == "Doc2b" # Doc3 should not be included because it has "Title" instead of "title" - doc_titles = [doc["doc_title"] for doc in referenced_documents] + doc_titles = [doc.doc_title for doc in referenced_documents] assert "Doc3" not in doc_titles @@ -1631,9 +1632,9 @@ def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(moc # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 - # Verify debug logging was called - mock_logger.debug.assert_called_once() - args = mock_logger.debug.call_args[0] + # Verify exception logging was called + mock_logger.exception.assert_called_once() + args = mock_logger.exception.call_args[0] assert "An exception was thrown in processing" in args[0] @@ -1658,9 +1659,9 @@ def test_process_knowledge_search_content_with_invalid_metadata_value_error(mock # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 - # Verify debug logging was called - mock_logger.debug.assert_called_once() - args = mock_logger.debug.call_args[0] + # Verify exception logging was called + mock_logger.exception.assert_called_once() + args = mock_logger.exception.call_args[0] assert "An exception was thrown in processing" in args[0] @@ -1725,6 +1726,20 @@ def test_process_knowledge_search_content_with_no_text_attribute(mocker): assert len(metadata_map) == 0 +def test_process_knowledge_search_content_with_none_content(mocker): + """Test _process_knowledge_search_content handles tool_response with content=None.""" + # Mock tool response with content = None + tool_response = mocker.Mock() + tool_response.content = None + + metadata_map = {} + + _process_knowledge_search_content(tool_response, metadata_map) + + # Verify metadata_map remains empty when content is None + assert len(metadata_map) == 0 + + @pytest.mark.asyncio async def test_retrieve_response_with_none_content(prepare_agent_mocks, mocker): """Test retrieve_response handles None content gracefully.""" diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index 8e03aa9c..78d32d10 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -40,6 +40,7 @@ streaming_query_endpoint_handler, retrieve_response, stream_build_event, + stream_end_event, ) from models.requests import QueryRequest, Attachment @@ -996,6 +997,113 @@ def test_stream_build_event_step_complete(): assert '"id": 0' in result +def test_stream_build_event_knowledge_search_with_invalid_metadata(mocker): + """Test stream_build_event handles invalid metadata in knowledge_search tool response.""" + mock_logger = mocker.patch("app.endpoints.streaming_query.logger") + + # Create a mock chunk with knowledge_search tool response containing invalid metadata + chunk = AgentTurnResponseStreamChunk( + event=TurnResponseEvent( + payload=AgentTurnResponseStepCompletePayload( + event_type="step_complete", + step_id="s1", + step_type="tool_execution", + step_details=ToolExecutionStep( + turn_id="t1", + step_id="s2", + step_type="tool_execution", + tool_responses=[ + ToolResponse( + call_id="c1", + tool_name="knowledge_search", + content=[ + TextContentItem( + text="""Result 1 +Content: Test content +Metadata: {'docs_url': 'https://example.com/doc1' 'title': 'Test Doc', 'document_id': 'doc-1'} +""", # Missing comma - invalid syntax + type="text", + ) + ], + ) + ], + tool_calls=[ + ToolCall( + call_id="t1", tool_name="knowledge_search", arguments={} + ) + ], + ), + ) + ) + ) + + metadata_map = {} + result_list = list(stream_build_event(chunk, 0, metadata_map)) + + # Verify metadata_map remains empty due to invalid metadata + assert len(metadata_map) == 0 + + # Verify the function still returns tool execution events + assert len(result_list) == 2 # One for tool_calls, one for tool_responses + + # Verify exception logging was called + mock_logger.exception.assert_called_once() + args = mock_logger.exception.call_args[0] + assert "An exception was thrown in processing" in args[0] + + +def test_stream_end_event_with_referenced_documents(): + """Test stream_end_event creates proper JSON with ReferencedDocument validation.""" + metadata_map = { + "doc-1": { + "docs_url": "https://example.com/doc1", + "title": "Test Document 1", + "document_id": "doc-1", + }, + "doc-2": { + "docs_url": "https://example.com/doc2", + "title": "Test Document 2", + "document_id": "doc-2", + }, + "doc-3": { + # Missing title - should be filtered out + "docs_url": "https://example.com/doc3", + "document_id": "doc-3", + }, + "doc-4": { + # Missing docs_url - should be filtered out + "title": "Test Document 4", + "document_id": "doc-4", + }, + } + + result = stream_end_event(metadata_map) + + # Parse the JSON response + parsed = json.loads(result.replace("data: ", "")) + + # Verify structure + assert parsed["event"] == "end" + assert "referenced_documents" in parsed["data"] + + # Verify only valid documents are included + referenced_docs = parsed["data"]["referenced_documents"] + assert len(referenced_docs) == 2 + + # Verify document structure and URL validation + doc_urls = [doc["doc_url"] for doc in referenced_docs] + doc_titles = [doc["doc_title"] for doc in referenced_docs] + + assert "https://example.com/doc1" in doc_urls + assert "https://example.com/doc2" in doc_urls + assert "Test Document 1" in doc_titles + assert "Test Document 2" in doc_titles + + # Verify filtered documents are not included + assert "https://example.com/doc3" not in doc_urls + assert "Test Document 4" not in doc_titles + + def test_stream_build_event_error(): """Test stream_build_event function returns a 'error' when chunk contains error information.""" # Create a mock chunk without an expected payload structure From d5e26225071950e605cb90e3e64b67cf2cd5593b Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Thu, 14 Aug 2025 16:17:24 -0400 Subject: [PATCH 06/16] more coderabbit --- src/app/endpoints/query.py | 49 ++++--- src/app/endpoints/streaming_query.py | 63 +++++---- src/utils/metadata.py | 34 +++++ tests/unit/app/endpoints/test_query.py | 104 ++++++++++++++ .../app/endpoints/test_streaming_query.py | 38 +++++ tests/unit/utils/test_metadata.py | 130 ++++++++++++++++++ 6 files changed, 369 insertions(+), 49 deletions(-) create mode 100644 src/utils/metadata.py create mode 100644 tests/unit/utils/test_metadata.py diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 4e63b418..db559365 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -1,14 +1,14 @@ """Handler for REST API call to provide answer to query.""" -import ast from datetime import datetime, UTC import json import logging import os from pathlib import Path -import re from typing import Annotated, Any, cast +import pydantic + from llama_stack_client import APIConnectionError from llama_stack_client import AsyncLlamaStackClient # type: ignore from llama_stack_client.types import UserMessage, Shield # type: ignore @@ -43,13 +43,12 @@ ) from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups from utils.suid import get_suid +from utils.metadata import parse_knowledge_search_metadata logger = logging.getLogger("app.endpoints.handlers") router = APIRouter(tags=["query"]) auth_dependency = get_auth_dependency() -METADATA_PATTERN = re.compile(r"^\s*Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) - def _process_knowledge_search_content( tool_response: Any, metadata_map: dict[str, dict[str, Any]] @@ -75,17 +74,14 @@ def _process_knowledge_search_content( if not text: continue - for match in METADATA_PATTERN.findall(text): - try: - meta = ast.literal_eval(match) - # Verify the result is a dict before accessing keys - if isinstance(meta, dict) and "document_id" in meta: - metadata_map[meta["document_id"]] = meta - except (SyntaxError, ValueError): # only expected from literal_eval - logger.exception( - "An exception was thrown in processing %s", - match, - ) + try: + parsed_metadata = parse_knowledge_search_metadata(text) + metadata_map.update(parsed_metadata) + except ValueError: + logger.exception( + "An exception was thrown in processing metadata from text: %s", + text[:200] + "..." if len(text) > 200 else text, + ) def extract_referenced_documents_from_steps(steps: list) -> list[ReferencedDocument]: @@ -113,12 +109,23 @@ def extract_referenced_documents_from_steps(steps: list) -> list[ReferencedDocum _process_knowledge_search_content(tool_response, metadata_map) - # Extract referenced documents from metadata - return [ - ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) - for v in metadata_map.values() - if "docs_url" in v and "title" in v - ] + # Extract referenced documents from metadata with error handling + referenced_documents = [] + for v in metadata_map.values(): + if "docs_url" in v and "title" in v: + try: + doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) + referenced_documents.append(doc) + except (pydantic.ValidationError, ValueError, Exception) as e: + logger.warning( + "Skipping invalid referenced document with docs_url='%s', title='%s': %s", + v.get("docs_url", ""), + v.get("title", ""), + str(e), + ) + continue + + return referenced_documents query_response: dict[int | str, dict[str, Any]] = { diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 2f875dca..5c5fe0de 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -1,11 +1,11 @@ """Handler for REST API call to provide answer to streaming query.""" -import ast import json -import re import logging from typing import Annotated, Any, AsyncIterator, Iterator +import pydantic + from llama_stack_client import APIConnectionError from llama_stack_client import AsyncLlamaStackClient # type: ignore from llama_stack_client.types import UserMessage # type: ignore @@ -27,6 +27,7 @@ from models.responses import ReferencedDocument from utils.endpoints import check_configuration_loaded, get_agent, get_system_prompt from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups +from utils.metadata import parse_knowledge_search_metadata from app.endpoints.query import ( get_rag_toolgroups, @@ -46,9 +47,6 @@ auth_dependency = get_auth_dependency() -METADATA_PATTERN = re.compile(r"\nMetadata: (\{.+})\n") - - def format_stream_data(d: dict) -> str: """Format outbound data in the Event Stream Format.""" data = json.dumps(d) @@ -79,15 +77,24 @@ def stream_end_event(metadata_map: dict) -> str: lambda v: ("docs_url" in v) and ("title" in v), metadata_map.values(), ): - doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) - referenced_documents.append( - { - "doc_url": str( - doc.doc_url - ), # Convert AnyUrl to string for JSON serialization - "doc_title": doc.doc_title, - } - ) + try: + doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) + referenced_documents.append( + { + "doc_url": str( + doc.doc_url + ), # Convert AnyUrl to string for JSON serialization + "doc_title": doc.doc_title, + } + ) + except (pydantic.ValidationError, ValueError, Exception) as e: + logger.warning( + "Skipping invalid referenced document with docs_url='%s', title='%s': %s", + v.get("docs_url", ""), + v.get("title", ""), + str(e), + ) + continue return format_stream_data( { @@ -335,20 +342,20 @@ def _handle_tool_execution_event( newline_pos = summary.find("\n") if newline_pos > 0: summary = summary[:newline_pos] - for match in METADATA_PATTERN.findall(text_content_item.text): - try: - meta = ast.literal_eval(match) - # Verify the result is a dict before accessing keys - if isinstance(meta, dict) and "document_id" in meta: - metadata_map[meta["document_id"]] = meta - except ( - SyntaxError, - ValueError, - ): # only expected from literal_eval - logger.exception( - "An exception was thrown in processing %s", - match, - ) + try: + parsed_metadata = parse_knowledge_search_metadata( + text_content_item.text + ) + metadata_map.update(parsed_metadata) + except ValueError: + logger.exception( + "An exception was thrown in processing metadata from text: %s", + ( + text_content_item.text[:200] + "..." + if len(text_content_item.text) > 200 + else text_content_item.text + ), + ) yield format_stream_data( { diff --git a/src/utils/metadata.py b/src/utils/metadata.py new file mode 100644 index 00000000..8db13d76 --- /dev/null +++ b/src/utils/metadata.py @@ -0,0 +1,34 @@ +"""Shared utilities for parsing metadata from knowledge search responses.""" + +import ast +import re +from typing import Any + + +METADATA_PATTERN = re.compile(r"^\s*Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) + + +def parse_knowledge_search_metadata(text: str) -> dict[str, Any]: + """Parse metadata from knowledge search text content. + + Args: + text: Text content that may contain metadata patterns + + Returns: + Dictionary of document_id -> metadata mappings + + Raises: + ValueError: If metadata parsing fails due to invalid JSON or syntax + """ + metadata_map: dict[str, Any] = {} + + for match in METADATA_PATTERN.findall(text): + try: + meta = ast.literal_eval(match) + # Verify the result is a dict before accessing keys + if isinstance(meta, dict) and "document_id" in meta: + metadata_map[meta["document_id"]] = meta + except (SyntaxError, ValueError) as e: + raise ValueError(f"Failed to parse metadata '{match}': {e}") from e + + return metadata_map diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 6ce01eba..a7bc4f59 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -1888,3 +1888,107 @@ async def test_retrieve_response_with_structured_content_object( # Should convert the structured object to string representation assert response == str(structured_content) assert conversation_id == "fake_conversation_id" + + +@pytest.mark.asyncio +async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, mocker): + """Test that retrieve_response skips entries with invalid docs_url.""" + mock_client, mock_agent = prepare_agent_mocks + mock_agent.create_turn.return_value.output_message.content = "LLM answer" + mock_client.shields.list.return_value = [] + mock_client.vector_dbs.list.return_value = [] + + # Mock tool response with valid and invalid docs_url entries + invalid_docs_url_results = [ + """knowledge_search tool found 2 chunks: +BEGIN of knowledge_search tool results. +""", + """Result 1 +Content: Valid content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Valid Doc', 'document_id': 'doc-1'} +""", + """Result 2 +Content: Invalid content +Metadata: {'docs_url': 'not-a-valid-url', 'title': 'Invalid Doc', 'document_id': 'doc-2'} +""", + """END of knowledge_search tool results. +""", + ] + + mock_tool_response = mocker.Mock() + mock_tool_response.call_id = "c1" + mock_tool_response.tool_name = "knowledge_search" + mock_tool_response.content = [ + mocker.Mock(text=s, type="text") for s in invalid_docs_url_results + ] + + mock_tool_execution_step = mocker.Mock() + mock_tool_execution_step.step_type = "tool_execution" + mock_tool_execution_step.tool_responses = [mock_tool_response] + + mock_agent.create_turn.return_value.steps = [mock_tool_execution_step] + + # Mock configuration with empty MCP servers + mock_config = mocker.Mock() + mock_config.mcp_servers = [] + mocker.patch("app.endpoints.query.configuration", mock_config) + mocker.patch( + "app.endpoints.query.get_agent", + return_value=(mock_agent, "fake_conversation_id", "fake_session_id"), + ) + + query_request = QueryRequest(query="What is OpenStack?") + model_id = "fake_model_id" + access_token = "test_token" + + response, conversation_id, referenced_documents = await retrieve_response( + mock_client, model_id, query_request, access_token + ) + + assert response == "LLM answer" + assert conversation_id == "fake_conversation_id" + + # Assert only the valid document is included, invalid one is skipped + assert len(referenced_documents) == 1 + assert str(referenced_documents[0].doc_url) == "https://example.com/doc1" + assert referenced_documents[0].doc_title == "Valid Doc" + + +@pytest.mark.asyncio +async def test_extract_referenced_documents_from_steps_handles_validation_errors( + mocker, +): + """Test that extract_referenced_documents_from_steps handles validation errors gracefully.""" + # Mock tool response with invalid docs_url that will cause pydantic validation error + mock_tool_response = mocker.Mock() + mock_tool_response.tool_name = "knowledge_search" + mock_tool_response.content = [ + mocker.Mock( + text="""Result 1 +Content: Valid content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Valid Doc', 'document_id': 'doc-1'} +""" + ), + mocker.Mock( + text="""Result 2 +Content: Invalid content +Metadata: {'docs_url': 'invalid-url', 'title': 'Invalid Doc', 'document_id': 'doc-2'} +""" + ), + ] + + mock_tool_execution_step = mocker.Mock() + mock_tool_execution_step.step_type = "tool_execution" + mock_tool_execution_step.tool_responses = [mock_tool_response] + + steps = [mock_tool_execution_step] + + # Import the function directly to test it + from app.endpoints.query import extract_referenced_documents_from_steps + + referenced_documents = extract_referenced_documents_from_steps(steps) + + # Should only return the valid document, skipping the invalid one + assert len(referenced_documents) == 1 + assert str(referenced_documents[0].doc_url) == "https://example.com/doc1" + assert referenced_documents[0].doc_title == "Valid Doc" diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index 78d32d10..9166d505 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -1104,6 +1104,44 @@ def test_stream_end_event_with_referenced_documents(): assert "Test Document 4" not in doc_titles +def test_stream_end_event_skips_invalid_docs_url(): + """Test stream_end_event skips entries with invalid docs_url.""" + metadata_map = { + "doc-1": { + "docs_url": "https://example.com/doc1", + "title": "Valid Document", + "document_id": "doc-1", + }, + "doc-2": { + "docs_url": "not-a-valid-url", # Invalid URL that will cause ValidationError + "title": "Invalid Document", + "document_id": "doc-2", + }, + "doc-3": { + "docs_url": "", # Empty URL that will cause ValidationError + "title": "Empty URL Document", + "document_id": "doc-3", + }, + } + + result = stream_end_event(metadata_map) + + # Parse the JSON response + parsed = json.loads(result.replace("data: ", "")) + + # Verify structure + assert parsed["event"] == "end" + assert "referenced_documents" in parsed["data"] + + # Verify only valid documents are included, invalid ones are skipped + referenced_docs = parsed["data"]["referenced_documents"] + assert len(referenced_docs) == 1 + + # Verify the valid document is included + assert referenced_docs[0]["doc_url"] == "https://example.com/doc1" + assert referenced_docs[0]["doc_title"] == "Valid Document" + + def test_stream_build_event_error(): """Test stream_build_event function returns a 'error' when chunk contains error information.""" # Create a mock chunk without an expected payload structure diff --git a/tests/unit/utils/test_metadata.py b/tests/unit/utils/test_metadata.py new file mode 100644 index 00000000..136a6e52 --- /dev/null +++ b/tests/unit/utils/test_metadata.py @@ -0,0 +1,130 @@ +"""Unit tests for utils.metadata module.""" + +import pytest + +from utils.metadata import parse_knowledge_search_metadata, METADATA_PATTERN + + +def test_metadata_pattern_exists(): + """Test that METADATA_PATTERN is properly defined.""" + assert METADATA_PATTERN is not None + assert hasattr(METADATA_PATTERN, "findall") + + +def test_parse_knowledge_search_metadata_valid_single(): + """Test parsing valid metadata with single entry.""" + text = """Result 1 +Content: Some content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 1 + assert "doc-1" in result + assert result["doc-1"]["docs_url"] == "https://example.com/doc1" + assert result["doc-1"]["title"] == "Doc1" + assert result["doc-1"]["document_id"] == "doc-1" + + +def test_parse_knowledge_search_metadata_valid_multiple(): + """Test parsing valid metadata with multiple entries.""" + text = """Result 1 +Content: Some content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} + +Result 2 +Content: More content +Metadata: {'docs_url': 'https://example.com/doc2', 'title': 'Doc2', 'document_id': 'doc-2'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 2 + assert "doc-1" in result + assert "doc-2" in result + assert result["doc-1"]["title"] == "Doc1" + assert result["doc-2"]["title"] == "Doc2" + + +def test_parse_knowledge_search_metadata_no_metadata(): + """Test parsing text with no metadata.""" + text = """Result 1 +Content: Some content without metadata +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 0 + + +def test_parse_knowledge_search_metadata_missing_document_id(): + """Test parsing metadata without document_id is ignored.""" + text = """Result 1 +Content: Some content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 0 + + +def test_parse_knowledge_search_metadata_malformed_json(): + """Test parsing malformed JSON raises ValueError.""" + text = """Result 1 +Content: Some content +Metadata: {'docs_url': 'https://example.com/doc1' 'title': 'Doc1', 'document_id': 'doc-1'} +""" + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text) + + assert "Failed to parse metadata" in str(exc_info.value) + + +def test_parse_knowledge_search_metadata_invalid_syntax(): + """Test parsing invalid Python syntax raises ValueError.""" + text = """Result 1 +Content: Some content +Metadata: {func_call(): 'value', 'title': 'Doc1', 'document_id': 'doc-1'} +""" + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text) + + assert "Failed to parse metadata" in str(exc_info.value) + + +def test_parse_knowledge_search_metadata_non_dict(): + """Test parsing non-dict metadata is ignored.""" + text = """Result 1 +Content: Some content +Metadata: "just a string" +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 0 + + +def test_parse_knowledge_search_metadata_mixed_valid_invalid(): + """Test parsing text with both valid and invalid metadata.""" + text = """Result 1 +Content: Some content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} + +Result 2 +Content: Bad content +Metadata: {'docs_url': 'https://example.com/doc2' 'title': 'Doc2', 'document_id': 'doc-2'} +""" + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text) + + assert "Failed to parse metadata" in str(exc_info.value) + + +def test_parse_knowledge_search_metadata_whitespace_handling(): + """Test parsing metadata with various whitespace patterns.""" + text = """Result 1 +Content: Some content + Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 1 + assert "doc-1" in result + assert result["doc-1"]["title"] == "Doc1" From a7f815e2867edfe83589a0a1b2f8edc7e6e08d8f Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Fri, 15 Aug 2025 08:38:30 -0400 Subject: [PATCH 07/16] more coderabbit --- src/app/endpoints/query.py | 6 ++- src/app/endpoints/streaming_query.py | 2 +- src/utils/metadata.py | 16 +++++-- tests/unit/app/endpoints/test_query.py | 4 +- tests/unit/utils/test_metadata.py | 62 ++++++++++++++++++++++++++ 5 files changed, 80 insertions(+), 10 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index db559365..50bda154 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -84,7 +84,9 @@ def _process_knowledge_search_content( ) -def extract_referenced_documents_from_steps(steps: list) -> list[ReferencedDocument]: +def extract_referenced_documents_from_steps( + steps: list[Any], +) -> list[ReferencedDocument]: """Extract referenced documents from tool execution steps. Args: @@ -116,7 +118,7 @@ def extract_referenced_documents_from_steps(steps: list) -> list[ReferencedDocum try: doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) referenced_documents.append(doc) - except (pydantic.ValidationError, ValueError, Exception) as e: + except (pydantic.ValidationError, ValueError) as e: logger.warning( "Skipping invalid referenced document with docs_url='%s', title='%s': %s", v.get("docs_url", ""), diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 5c5fe0de..bfd722bb 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -87,7 +87,7 @@ def stream_end_event(metadata_map: dict) -> str: "doc_title": doc.doc_title, } ) - except (pydantic.ValidationError, ValueError, Exception) as e: + except (pydantic.ValidationError, ValueError) as e: logger.warning( "Skipping invalid referenced document with docs_url='%s', title='%s': %s", v.get("docs_url", ""), diff --git a/src/utils/metadata.py b/src/utils/metadata.py index 8db13d76..09150247 100644 --- a/src/utils/metadata.py +++ b/src/utils/metadata.py @@ -8,19 +8,24 @@ METADATA_PATTERN = re.compile(r"^\s*Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) -def parse_knowledge_search_metadata(text: str) -> dict[str, Any]: +def parse_knowledge_search_metadata( + text: str, *, strict: bool = True +) -> dict[str, dict[str, Any]]: """Parse metadata from knowledge search text content. Args: text: Text content that may contain metadata patterns + strict: If True (default), raise ValueError on first parsing error. + If False, skip invalid blocks and continue parsing. Returns: Dictionary of document_id -> metadata mappings Raises: - ValueError: If metadata parsing fails due to invalid JSON or syntax + ValueError: If metadata parsing fails due to invalid Python-literal or JSON-like syntax + (only in strict mode) """ - metadata_map: dict[str, Any] = {} + metadata_map: dict[str, dict[str, Any]] = {} for match in METADATA_PATTERN.findall(text): try: @@ -29,6 +34,9 @@ def parse_knowledge_search_metadata(text: str) -> dict[str, Any]: if isinstance(meta, dict) and "document_id" in meta: metadata_map[meta["document_id"]] = meta except (SyntaxError, ValueError) as e: - raise ValueError(f"Failed to parse metadata '{match}': {e}") from e + if strict: + raise ValueError(f"Failed to parse metadata '{match}': {e}") from e + # non-strict mode: skip bad blocks, keep the rest + continue return metadata_map diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index a7bc4f59..914f48ce 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -23,6 +23,7 @@ get_rag_toolgroups, evaluate_model_hints, _process_knowledge_search_content, + extract_referenced_documents_from_steps, ) from models.requests import QueryRequest, Attachment @@ -1983,9 +1984,6 @@ async def test_extract_referenced_documents_from_steps_handles_validation_errors steps = [mock_tool_execution_step] - # Import the function directly to test it - from app.endpoints.query import extract_referenced_documents_from_steps - referenced_documents = extract_referenced_documents_from_steps(steps) # Should only return the valid document, skipping the invalid one diff --git a/tests/unit/utils/test_metadata.py b/tests/unit/utils/test_metadata.py index 136a6e52..2836a992 100644 --- a/tests/unit/utils/test_metadata.py +++ b/tests/unit/utils/test_metadata.py @@ -128,3 +128,65 @@ def test_parse_knowledge_search_metadata_whitespace_handling(): assert len(result) == 1 assert "doc-1" in result assert result["doc-1"]["title"] == "Doc1" + + +def test_parse_metadata_duplicate_document_id_last_wins(): + """Test that duplicate document_id entries overwrite (last wins).""" + text = ( + "Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1a', " + "'document_id': 'dup'}\n" + "Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1b', " + "'document_id': 'dup'}" + ) + result = parse_knowledge_search_metadata(text) + + assert list(result.keys()) == ["dup"] + assert result["dup"]["title"] == "Doc1b" + + +def test_parse_knowledge_search_metadata_non_strict_mode(): + """Test non-strict mode skips invalid blocks and continues parsing.""" + text = """Result 1 +Content: Valid content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} + +Result 2 +Content: Bad content +Metadata: {'docs_url': 'https://example.com/doc2' 'title': 'Doc2', 'document_id': 'doc-2'} + +Result 3 +Content: More valid content +Metadata: {'docs_url': 'https://example.com/doc3', 'title': 'Doc3', 'document_id': 'doc-3'} +""" + result = parse_knowledge_search_metadata(text, strict=False) + + # Should have 2 valid documents, skipping the malformed one + assert len(result) == 2 + assert "doc-1" in result + assert "doc-3" in result + assert "doc-2" not in result # malformed entry should be skipped + assert result["doc-1"]["title"] == "Doc1" + assert result["doc-3"]["title"] == "Doc3" + + +def test_parse_knowledge_search_metadata_strict_mode_default(): + """Test that strict mode is the default behavior.""" + text = """Result 1 +Content: Valid content +Metadata: {'docs_url': 'https://example.com/doc1', 'title': 'Doc1', 'document_id': 'doc-1'} + +Result 2 +Content: Bad content +Metadata: {'docs_url': 'https://example.com/doc2' 'title': 'Doc2', 'document_id': 'doc-2'} +""" + # Should raise ValueError in strict mode (default) + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text) + + assert "Failed to parse metadata" in str(exc_info.value) + + # Explicitly setting strict=True should behave the same + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text, strict=True) + + assert "Failed to parse metadata" in str(exc_info.value) From af646c65b6f25e2b34dfb1d60afa39503340b15c Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Fri, 15 Aug 2025 09:49:29 -0400 Subject: [PATCH 08/16] =?UTF-8?q?=F0=9F=90=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/utils/metadata.py | 62 +++++++++- tests/unit/utils/test_metadata.py | 196 ++++++++++++++++++++++++++++-- 2 files changed, 247 insertions(+), 11 deletions(-) diff --git a/src/utils/metadata.py b/src/utils/metadata.py index 09150247..6289d2c9 100644 --- a/src/utils/metadata.py +++ b/src/utils/metadata.py @@ -5,7 +5,40 @@ from typing import Any -METADATA_PATTERN = re.compile(r"^\s*Metadata:\s*(\{.*?\})\s*$", re.MULTILINE) +# Case-insensitive pattern to find "Metadata:" labels +METADATA_LABEL_PATTERN = re.compile(r"^\s*metadata:\s*", re.MULTILINE | re.IGNORECASE) + + +def _extract_balanced_braces(text: str, start_pos: int) -> str: + """Extract a balanced brace substring starting from start_pos. + + Args: + text: The text to search + start_pos: Position where the opening brace should be + + Returns: + The balanced brace substring including the braces + + Raises: + ValueError: If no balanced braces are found + """ + if start_pos >= len(text) or text[start_pos] != "{": + raise ValueError("No opening brace found at start position") + + brace_count = 0 + pos = start_pos + + while pos < len(text): + char = text[pos] + if char == "{": + brace_count += 1 + elif char == "}": + brace_count -= 1 + if brace_count == 0: + return text[start_pos : pos + 1] + pos += 1 + + raise ValueError("Unmatched opening brace - no closing brace found") def parse_knowledge_search_metadata( @@ -27,15 +60,36 @@ def parse_knowledge_search_metadata( """ metadata_map: dict[str, dict[str, Any]] = {} - for match in METADATA_PATTERN.findall(text): + # Find all "Metadata:" labels (case-insensitive) + for match in METADATA_LABEL_PATTERN.finditer(text): try: - meta = ast.literal_eval(match) + # Find the position right after the "Metadata:" label + label_end = match.end() + + # Skip any whitespace after the label + pos = label_end + while pos < len(text) and text[pos].isspace(): + pos += 1 + + # Look for opening brace + if pos >= len(text) or text[pos] != "{": + continue # No brace found, skip this match + + # Extract balanced brace content + brace_content = _extract_balanced_braces(text, pos) + + # Parse the extracted content + meta = ast.literal_eval(brace_content) + # Verify the result is a dict before accessing keys if isinstance(meta, dict) and "document_id" in meta: metadata_map[meta["document_id"]] = meta + except (SyntaxError, ValueError) as e: if strict: - raise ValueError(f"Failed to parse metadata '{match}': {e}") from e + raise ValueError( + f"Failed to parse metadata at position {match.start()}: {e}" + ) from e # non-strict mode: skip bad blocks, keep the rest continue diff --git a/tests/unit/utils/test_metadata.py b/tests/unit/utils/test_metadata.py index 2836a992..40f2d43a 100644 --- a/tests/unit/utils/test_metadata.py +++ b/tests/unit/utils/test_metadata.py @@ -2,13 +2,22 @@ import pytest -from utils.metadata import parse_knowledge_search_metadata, METADATA_PATTERN +from utils.metadata import parse_knowledge_search_metadata, METADATA_LABEL_PATTERN def test_metadata_pattern_exists(): - """Test that METADATA_PATTERN is properly defined.""" - assert METADATA_PATTERN is not None - assert hasattr(METADATA_PATTERN, "findall") + """Test that METADATA_LABEL_PATTERN is properly defined and captures labels correctly.""" + assert METADATA_LABEL_PATTERN is not None + assert hasattr(METADATA_LABEL_PATTERN, "finditer") + + # Test that the pattern captures metadata labels case-insensitively + sample = "Foo\nMetadata: {'a': 1}\nMETADATA: {'b': 2}\nBar" + matches = list(METADATA_LABEL_PATTERN.finditer(sample)) + assert len(matches) == 2 + + # Check that the matches are at the expected positions + assert sample[matches[0].start() : matches[0].end()] == "Metadata: " + assert sample[matches[1].start() : matches[1].end()] == "METADATA: " def test_parse_knowledge_search_metadata_valid_single(): @@ -66,8 +75,8 @@ def test_parse_knowledge_search_metadata_missing_document_id(): assert len(result) == 0 -def test_parse_knowledge_search_metadata_malformed_json(): - """Test parsing malformed JSON raises ValueError.""" +def test_parse_knowledge_search_metadata_malformed_literal(): + """Test parsing malformed Python literal raises ValueError.""" text = """Result 1 Content: Some content Metadata: {'docs_url': 'https://example.com/doc1' 'title': 'Doc1', 'document_id': 'doc-1'} @@ -140,7 +149,8 @@ def test_parse_metadata_duplicate_document_id_last_wins(): ) result = parse_knowledge_search_metadata(text) - assert list(result.keys()) == ["dup"] + assert len(result) == 1 + assert set(result.keys()) == {"dup"} assert result["dup"]["title"] == "Doc1b" @@ -190,3 +200,175 @@ def test_parse_knowledge_search_metadata_strict_mode_default(): parse_knowledge_search_metadata(text, strict=True) assert "Failed to parse metadata" in str(exc_info.value) + + +def test_metadata_pattern_case_insensitive_and_nested(): + """Test case-insensitive matching and nested payloads.""" + text = """Result +Content +METADATA: {'document_id': 'doc-1', 'nested': {'k': [1, 2, 3]}, 'title': 'Nested Doc'} +Another result +metadata: {'document_id': 'doc-2', 'complex': {'a': {'b': {'c': 42}}, 'list': [{'x': 1}, {'y': 2}]}, 'title': 'Complex Doc'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 2 + assert "doc-1" in result + assert "doc-2" in result + + # Verify the nested structure was parsed correctly + assert result["doc-1"]["nested"]["k"] == [1, 2, 3] + assert result["doc-1"]["title"] == "Nested Doc" + + assert result["doc-2"]["complex"]["a"]["b"]["c"] == 42 + assert result["doc-2"]["complex"]["list"][0]["x"] == 1 + assert result["doc-2"]["complex"]["list"][1]["y"] == 2 + assert result["doc-2"]["title"] == "Complex Doc" + + +def test_metadata_pattern_various_case_variations(): + """Test different case variations of metadata label.""" + text = """ +Metadata: {'document_id': 'doc-1', 'title': 'Standard'} +METADATA: {'document_id': 'doc-2', 'title': 'Uppercase'} +metadata: {'document_id': 'doc-3', 'title': 'Lowercase'} +MetaData: {'document_id': 'doc-4', 'title': 'Mixed Case'} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 4 + assert result["doc-1"]["title"] == "Standard" + assert result["doc-2"]["title"] == "Uppercase" + assert result["doc-3"]["title"] == "Lowercase" + assert result["doc-4"]["title"] == "Mixed Case" + + +def test_balanced_braces_with_nested_dicts_and_strings(): + """Test balanced brace parsing with complex nested structures.""" + text = """ +Metadata: {'document_id': 'doc-1', 'data': {'nested': 'value with {braces} in string'}, 'array': [{'inner': 'val'}]} +""" + result = parse_knowledge_search_metadata(text) + + assert len(result) == 1 + assert result["doc-1"]["data"]["nested"] == "value with {braces} in string" + assert result["doc-1"]["array"][0]["inner"] == "val" + + +def test_unmatched_braces_handling(): + """Test handling of unmatched braces in strict and non-strict modes.""" + text = """ +Metadata: {'document_id': 'doc-1', 'incomplete': 'missing brace' +Valid after: some text +Metadata: {'document_id': 'doc-2', 'title': 'Valid Doc'} +""" + # Strict mode should raise error + with pytest.raises(ValueError) as exc_info: + parse_knowledge_search_metadata(text, strict=True) + + assert "Failed to parse metadata" in str(exc_info.value) + + # Non-strict mode should skip the invalid entry and parse the valid one + result = parse_knowledge_search_metadata(text, strict=False) + assert len(result) == 1 + assert "doc-2" in result + assert result["doc-2"]["title"] == "Valid Doc" + + +def test_no_opening_brace_after_metadata_label(): + """Test handling when no opening brace follows metadata label.""" + text = """ +Metadata: not a dict +Some other content +Metadata: {'document_id': 'doc-1', 'title': 'Valid'} +""" + result = parse_knowledge_search_metadata(text) + + # Should only find the valid metadata entry + assert len(result) == 1 + assert "doc-1" in result + assert result["doc-1"]["title"] == "Valid" + + +@pytest.mark.parametrize( + "text, strict, expected_ids, description", + [ + # Valid cases + ( + "Metadata: {'document_id': 'a', 'title': 'Doc A'}", + True, + {"a"}, + "single valid metadata", + ), + ( + "Metadata: {'document_id': 'a', 'title': 'Doc A'}\n" + "Metadata: {'document_id': 'b', 'title': 'Doc B'}", + True, + {"a", "b"}, + "multiple valid metadata", + ), + ( + "METADATA: {'document_id': 'upper', 'title': 'Upper'}\n" + "metadata: {'document_id': 'lower', 'title': 'Lower'}", + True, + {"upper", "lower"}, + "case-insensitive labels", + ), + # Error handling - strict mode + ( + "Metadata: {'document_id': 'a', 'title': 'Doc A'}\n" + "Metadata: {'document_id': 'b' 'oops': 1}", + False, + {"a"}, + "malformed metadata skipped in non-strict mode", + ), + ( + "Metadata: not_a_dict\nMetadata: {'document_id': 'valid', 'title': 'Valid'}", + True, + {"valid"}, + "non-dict content ignored", + ), + # No metadata cases + ( + "Some text without metadata", + True, + set(), + "no metadata found", + ), + ( + "Metadata: {'title': 'No ID'}", + True, + set(), + "metadata without document_id ignored", + ), + ], +) +def test_parse_metadata_parametrized(text, strict, expected_ids, description): + """Parametrized test for various metadata parsing scenarios.""" + if strict and "malformed" in description: + # Should raise in strict mode for malformed content + with pytest.raises(ValueError): + parse_knowledge_search_metadata(text, strict=strict) + else: + result = parse_knowledge_search_metadata(text, strict=strict) + assert set(result.keys()) == expected_ids, f"Failed for: {description}" + + +@pytest.mark.parametrize( + "metadata_label, expected_matches", + [ + ("Metadata:", 1), + ("METADATA:", 1), + ("metadata:", 1), + ("MetaData:", 1), + ("MetaDaTa:", 1), + (" Metadata: ", 1), # with whitespace + ("NotMetadata:", 0), # should not match + ("metadata", 0), # missing colon + ], +) +def test_label_pattern_matching(metadata_label, expected_matches): + """Test that the label pattern matches various case variations correctly.""" + sample = f"Some text\n{metadata_label} {{'document_id': 'test'}}\nMore text" + matches = list(METADATA_LABEL_PATTERN.finditer(sample)) + assert len(matches) == expected_matches From 06cba91b046018982669e59dba324fd7ea098304 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Fri, 15 Aug 2025 10:35:01 -0400 Subject: [PATCH 09/16] fixes in-place modification issue --- src/app/endpoints/query.py | 26 ++++++++++++++++-------- tests/unit/app/endpoints/test_query.py | 28 +++++++------------------- 2 files changed, 25 insertions(+), 29 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 50bda154..dfb89467 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -50,23 +50,30 @@ auth_dependency = get_auth_dependency() -def _process_knowledge_search_content( - tool_response: Any, metadata_map: dict[str, dict[str, Any]] -) -> None: - """Process knowledge search tool response content for metadata.""" +def _process_knowledge_search_content(tool_response: Any) -> dict[str, dict[str, Any]]: + """Process knowledge search tool response content for metadata. + + Args: + tool_response: Tool response object containing content to parse + + Returns: + Dictionary mapping document_id to metadata dict + """ + metadata_map: dict[str, dict[str, Any]] = {} + # Guard against missing tool_response or content if not tool_response: - return + return metadata_map content = getattr(tool_response, "content", None) if not content: - return + return metadata_map # Ensure content is iterable try: iter(content) except TypeError: - return + return metadata_map for text_content_item in content: # Skip items that lack a non-empty "text" attribute @@ -83,6 +90,8 @@ def _process_knowledge_search_content( text[:200] + "..." if len(text) > 200 else text, ) + return metadata_map + def extract_referenced_documents_from_steps( steps: list[Any], @@ -109,7 +118,8 @@ def extract_referenced_documents_from_steps( ) != "knowledge_search" or not getattr(tool_response, "content", []): continue - _process_knowledge_search_content(tool_response, metadata_map) + response_metadata = _process_knowledge_search_content(tool_response) + metadata_map.update(response_metadata) # Extract referenced documents from metadata with error handling referenced_documents = [] diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 914f48ce..63b7f73e 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -1601,9 +1601,7 @@ def test_process_knowledge_search_content_with_valid_metadata(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata was correctly parsed and added assert "doc-1" in metadata_map @@ -1626,9 +1624,7 @@ def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(moc tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 @@ -1653,9 +1649,7 @@ def test_process_knowledge_search_content_with_invalid_metadata_value_error(mock tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty due to exception assert len(metadata_map) == 0 @@ -1680,9 +1674,7 @@ def test_process_knowledge_search_content_with_non_dict_metadata(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty (no document_id in string) assert len(metadata_map) == 0 @@ -1703,9 +1695,7 @@ def test_process_knowledge_search_content_with_metadata_missing_document_id(mock tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty since document_id is missing assert len(metadata_map) == 0 @@ -1719,9 +1709,7 @@ def test_process_knowledge_search_content_with_no_text_attribute(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty since text attribute is missing assert len(metadata_map) == 0 @@ -1733,9 +1721,7 @@ def test_process_knowledge_search_content_with_none_content(mocker): tool_response = mocker.Mock() tool_response.content = None - metadata_map = {} - - _process_knowledge_search_content(tool_response, metadata_map) + metadata_map = _process_knowledge_search_content(tool_response) # Verify metadata_map remains empty when content is None assert len(metadata_map) == 0 From 1f8fbb9f2028aa172afd83468f0f076a6b886507 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Fri, 15 Aug 2025 11:10:44 -0400 Subject: [PATCH 10/16] codrabbit test change requests --- tests/unit/app/endpoints/test_query.py | 76 ++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 63b7f73e..49350a08 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -264,6 +264,7 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): assert response.response == llm_response assert response.conversation_id == conversation_id assert response.referenced_documents == referenced_documents + assert all(isinstance(d, ReferencedDocument) for d in response.referenced_documents) assert len(response.referenced_documents) == 2 assert response.referenced_documents[0].doc_title == "Doc1" assert response.referenced_documents[1].doc_title == "Doc2" @@ -1727,6 +1728,71 @@ def test_process_knowledge_search_content_with_none_content(mocker): assert len(metadata_map) == 0 +def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker): + """The last metadata block for a given document_id should win.""" + text_items = [ + mocker.Mock( + text="Metadata: {'docs_url': 'https://example.com/first', " + "'title': 'First', 'document_id': 'doc-x'}" + ), + mocker.Mock( + text="Metadata: {'docs_url': 'https://example.com/second', " + "'title': 'Second', 'document_id': 'doc-x'}" + ), + ] + tool_response = mocker.Mock() + tool_response.tool_name = "knowledge_search" + tool_response.content = text_items + + # Process content + metadata_map = _process_knowledge_search_content(tool_response) + assert metadata_map["doc-x"]["docs_url"] == "https://example.com/second" + assert metadata_map["doc-x"]["title"] == "Second" + + # Ensure extraction reflects last-wins as well + step = mocker.Mock() + step.step_type = "tool_execution" + step.tool_responses = [tool_response] + docs = extract_referenced_documents_from_steps([step]) + assert len(docs) == 1 + assert str(docs[0].doc_url) == "https://example.com/second" + assert docs[0].doc_title == "Second" + + +def test_process_knowledge_search_content_with_braces_inside_strings(mocker): + """Test that braces inside strings are handled correctly.""" + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +metadata: {'document_id': 'doc-100', 'title': 'A {weird} title', " +"'docs_url': 'https://example.com/100', 'extra': {'note': 'contains {braces}'}}""" + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = _process_knowledge_search_content(tool_response) + assert "doc-100" in metadata_map + assert metadata_map["doc-100"]["title"] == "A {weird} title" + assert metadata_map["doc-100"]["docs_url"] == "https://example.com/100" + assert metadata_map["doc-100"]["extra"]["note"] == "contains {braces}" + + +def test_process_knowledge_search_content_with_nested_objects(mocker): + """Test that nested objects are parsed correctly.""" + text_content_item = mocker.Mock() + text_content_item.text = """Result 1 +Content: Test content +MetaData: {"document_id": "doc-200", "title": "Nested JSON", " +""docs_url": "https://example.com/200", "meta": {"k": {"inner": 1}}}""" + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = _process_knowledge_search_content(tool_response) + assert "doc-200" in metadata_map + assert metadata_map["doc-200"]["title"] == "Nested JSON" + assert metadata_map["doc-200"]["docs_url"] == "https://example.com/200" + assert metadata_map["doc-200"]["meta"]["k"]["inner"] == 1 + + @pytest.mark.asyncio async def test_retrieve_response_with_none_content(prepare_agent_mocks, mocker): """Test retrieve_response handles None content gracefully.""" @@ -1880,6 +1946,9 @@ async def test_retrieve_response_with_structured_content_object( @pytest.mark.asyncio async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, mocker): """Test that retrieve_response skips entries with invalid docs_url.""" + # Mock logger to capture warning logs + mock_logger = mocker.patch("app.endpoints.query.logger") + mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" mock_client.shields.list.return_value = [] @@ -1940,6 +2009,13 @@ async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, moc assert str(referenced_documents[0].doc_url) == "https://example.com/doc1" assert referenced_documents[0].doc_title == "Valid Doc" + # Ensure we logged a warning for the invalid docs_url + assert any( + call[0][0].startswith("Skipping invalid referenced document") + or "Skipping invalid referenced document" in str(call) + for call in mock_logger.warning.call_args_list + ) + @pytest.mark.asyncio async def test_extract_referenced_documents_from_steps_handles_validation_errors( From 125619323d95b83e36ddd2542920f05c6d2e6779 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Fri, 15 Aug 2025 11:28:11 -0400 Subject: [PATCH 11/16] fixes some issues with tests --- tests/unit/app/endpoints/test_query.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 49350a08..52419b14 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -1762,10 +1762,12 @@ def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker def test_process_knowledge_search_content_with_braces_inside_strings(mocker): """Test that braces inside strings are handled correctly.""" text_content_item = mocker.Mock() - text_content_item.text = """Result 1 -Content: Test content -metadata: {'document_id': 'doc-100', 'title': 'A {weird} title', " -"'docs_url': 'https://example.com/100', 'extra': {'note': 'contains {braces}'}}""" + text_content_item.text = ( + "Result 1\n" + "Content: Test content\n" + "Metadata: {'document_id': 'doc-100', 'title': 'A {weird} title', " + "'docs_url': 'https://example.com/100', 'extra': {'note': 'contains {braces}'}}" + ) tool_response = mocker.Mock() tool_response.content = [text_content_item] @@ -1779,10 +1781,12 @@ def test_process_knowledge_search_content_with_braces_inside_strings(mocker): def test_process_knowledge_search_content_with_nested_objects(mocker): """Test that nested objects are parsed correctly.""" text_content_item = mocker.Mock() - text_content_item.text = """Result 1 -Content: Test content -MetaData: {"document_id": "doc-200", "title": "Nested JSON", " -""docs_url": "https://example.com/200", "meta": {"k": {"inner": 1}}}""" + text_content_item.text = ( + "Result 1\n" + "Content: Test content\n" + 'Metadata: {"document_id": "doc-200", "title": "Nested JSON", ' + '"docs_url": "https://example.com/200", "meta": {"k": {"inner": 1}}}' + ) tool_response = mocker.Mock() tool_response.content = [text_content_item] From 1fa7b87662810aa136d0539072db6832539c8947 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Mon, 18 Aug 2025 10:00:33 -0400 Subject: [PATCH 12/16] avoid brittle equality --- tests/unit/app/endpoints/test_query.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 52419b14..7a195beb 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -263,11 +263,13 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): # Assert the response contains referenced documents assert response.response == llm_response assert response.conversation_id == conversation_id - assert response.referenced_documents == referenced_documents + # Avoid brittle equality on Pydantic models; compare fields instead + assert [(str(d.doc_url), d.doc_title) for d in response.referenced_documents] == [ + ("https://example.com/doc1", "Doc1"), + ("https://example.com/doc2", "Doc2"), + ] assert all(isinstance(d, ReferencedDocument) for d in response.referenced_documents) assert len(response.referenced_documents) == 2 - assert response.referenced_documents[0].doc_title == "Doc1" - assert response.referenced_documents[1].doc_title == "Doc2" # Assert the metric for successful LLM calls is incremented mock_metric.labels("fake_provider_id", "fake_model_id").inc.assert_called_once() @@ -1680,8 +1682,8 @@ def test_process_knowledge_search_content_with_non_dict_metadata(mocker): # Verify metadata_map remains empty (no document_id in string) assert len(metadata_map) == 0 - # No exception should be logged since string is valid literal - mock_logger.debug.assert_not_called() + # No exception should be logged since string is a valid literal and simply ignored + mock_logger.exception.assert_not_called() def test_process_knowledge_search_content_with_metadata_missing_document_id(mocker): @@ -2019,6 +2021,10 @@ async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, moc or "Skipping invalid referenced document" in str(call) for call in mock_logger.warning.call_args_list ) + # Verify the bad URL is included in the log message for extra confidence + assert any( + "not-a-valid-url" in str(call) for call in mock_logger.warning.call_args_list + ) @pytest.mark.asyncio From 12becc5ccdea307e8b2c49d9aaf90811bea0b902 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Mon, 18 Aug 2025 10:46:50 -0400 Subject: [PATCH 13/16] coderabbit --- src/app/endpoints/query.py | 106 ++---------------------- src/app/endpoints/streaming_query.py | 12 +-- src/utils/metadata.py | 109 +++++++++++++++++++++++++ tests/unit/app/endpoints/test_query.py | 76 +++++++---------- 4 files changed, 153 insertions(+), 150 deletions(-) diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 2e85be77..ce67b1c6 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -7,7 +7,6 @@ from pathlib import Path from typing import Annotated, Any, cast -import pydantic from llama_stack_client import APIConnectionError from llama_stack_client import AsyncLlamaStackClient # type: ignore @@ -43,103 +42,15 @@ ) from utils.mcp_headers import mcp_headers_dependency, handle_mcp_headers_with_toolgroups from utils.suid import get_suid -from utils.metadata import parse_knowledge_search_metadata +from utils.metadata import ( + extract_referenced_documents_from_steps, +) logger = logging.getLogger("app.endpoints.handlers") router = APIRouter(tags=["query"]) auth_dependency = get_auth_dependency() -def _process_knowledge_search_content(tool_response: Any) -> dict[str, dict[str, Any]]: - """Process knowledge search tool response content for metadata. - - Args: - tool_response: Tool response object containing content to parse - - Returns: - Dictionary mapping document_id to metadata dict - """ - metadata_map: dict[str, dict[str, Any]] = {} - - # Guard against missing tool_response or content - if not tool_response: - return metadata_map - - content = getattr(tool_response, "content", None) - if not content: - return metadata_map - - # Ensure content is iterable - try: - iter(content) - except TypeError: - return metadata_map - - for text_content_item in content: - # Skip items that lack a non-empty "text" attribute - text = getattr(text_content_item, "text", None) - if not text: - continue - - try: - parsed_metadata = parse_knowledge_search_metadata(text) - metadata_map.update(parsed_metadata) - except ValueError: - logger.exception( - "An exception was thrown in processing metadata from text: %s", - text[:200] + "..." if len(text) > 200 else text, - ) - - return metadata_map - - -def extract_referenced_documents_from_steps( - steps: list[Any], -) -> list[ReferencedDocument]: - """Extract referenced documents from tool execution steps. - - Args: - steps: List of response steps from the agent - - Returns: - List of referenced documents with doc_url and doc_title - """ - metadata_map: dict[str, dict[str, Any]] = {} - - for step in steps: - if getattr(step, "step_type", "") != "tool_execution" or not hasattr( - step, "tool_responses" - ): - continue - - for tool_response in getattr(step, "tool_responses", []) or []: - if getattr( - tool_response, "tool_name", "" - ) != "knowledge_search" or not getattr(tool_response, "content", []): - continue - - response_metadata = _process_knowledge_search_content(tool_response) - metadata_map.update(response_metadata) - - # Extract referenced documents from metadata with error handling - referenced_documents = [] - for v in metadata_map.values(): - if "docs_url" in v and "title" in v: - try: - doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) - referenced_documents.append(doc) - except (pydantic.ValidationError, ValueError) as e: - logger.warning( - "Skipping invalid referenced document with docs_url='%s', title='%s': %s", - v.get("docs_url", ""), - v.get("title", ""), - str(e), - ) - continue - - return referenced_documents - - query_response: dict[int | str, dict[str, Any]] = { 200: { "conversation_id": "123e4567-e89b-12d3-a456-426614174000", @@ -516,8 +427,9 @@ async def retrieve_response( # pylint: disable=too-many-locals,too-many-branche mcp_headers (dict[str, dict[str, str]], optional): Headers for multi-component processing. Returns: - tuple[str, str]: A tuple containing the LLM or agent's response content - and the conversation ID. + tuple[str, str, list[ReferencedDocument]]: A tuple containing the response + content, the conversation ID, and the list of referenced documents parsed + from tool execution steps. """ available_input_shields = [ shield.identifier @@ -615,12 +527,12 @@ async def retrieve_response( # pylint: disable=too-many-locals,too-many-branche # Safely guard access to output_message and content output_message = getattr(response_obj, "output_message", None) if output_message and getattr(output_message, "content", None) is not None: - content_str = str(output_message.content) + response_text = str(output_message.content) else: - content_str = "" + response_text = "" return ( - content_str, + response_text, conversation_id, referenced_documents, ) diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index 73169b37..bf6c0137 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -452,17 +452,13 @@ def _handle_tool_execution_event( summary = summary[:newline_pos] try: parsed_metadata = parse_knowledge_search_metadata( - text_content_item.text + text_content_item.text, strict=False ) metadata_map.update(parsed_metadata) - except ValueError: + except ValueError as e: logger.exception( - "An exception was thrown in processing metadata from text: %s", - ( - text_content_item.text[:200] + "..." - if len(text_content_item.text) > 200 - else text_content_item.text - ), + "Error processing metadata from text; position=%s", + getattr(e, "position", "unknown"), ) yield format_stream_data( diff --git a/src/utils/metadata.py b/src/utils/metadata.py index 6289d2c9..6bc49736 100644 --- a/src/utils/metadata.py +++ b/src/utils/metadata.py @@ -1,9 +1,17 @@ """Shared utilities for parsing metadata from knowledge search responses.""" import ast +import json +import logging import re from typing import Any +import pydantic + +from models.responses import ReferencedDocument + +logger = logging.getLogger(__name__) + # Case-insensitive pattern to find "Metadata:" labels METADATA_LABEL_PATTERN = re.compile(r"^\s*metadata:\s*", re.MULTILINE | re.IGNORECASE) @@ -94,3 +102,104 @@ def parse_knowledge_search_metadata( continue return metadata_map + + +def process_knowledge_search_content(tool_response: Any) -> dict[str, dict[str, Any]]: + """Process knowledge search tool response content for metadata. + + Args: + tool_response: Tool response object containing content to parse + + Returns: + Dictionary mapping document_id to metadata dict + """ + metadata_map: dict[str, dict[str, Any]] = {} + + # Guard against missing tool_response or content + if not tool_response: + return metadata_map + + content = getattr(tool_response, "content", None) + if not content: + return metadata_map + + # Handle string content by attempting JSON parsing + if isinstance(content, str): + try: + content = json.loads(content, strict=False) + except (json.JSONDecodeError, TypeError): + # If JSON parsing fails or content is still a string, return empty + if isinstance(content, str): + return metadata_map + + # Ensure content is iterable (but not a string) + if isinstance(content, str): + return metadata_map + try: + iter(content) + except TypeError: + return metadata_map + + for text_content_item in content: + # Skip items that lack a non-empty "text" attribute + text = getattr(text_content_item, "text", None) + if not text: + continue + + try: + parsed_metadata = parse_knowledge_search_metadata(text, strict=False) + metadata_map.update(parsed_metadata) + except ValueError as e: + logger.exception( + "Error processing metadata from text; position=%s", + getattr(e, "position", "unknown"), + ) + + return metadata_map + + +def extract_referenced_documents_from_steps( + steps: list[Any], +) -> list[ReferencedDocument]: + """Extract referenced documents from tool execution steps. + + Args: + steps: List of response steps from the agent + + Returns: + List of referenced documents with doc_url and doc_title, sorted deterministically + """ + metadata_map: dict[str, dict[str, Any]] = {} + + for step in steps: + if getattr(step, "step_type", "") != "tool_execution" or not hasattr( + step, "tool_responses" + ): + continue + + for tool_response in getattr(step, "tool_responses", []) or []: + if getattr( + tool_response, "tool_name", "" + ) != "knowledge_search" or not getattr(tool_response, "content", []): + continue + + response_metadata = process_knowledge_search_content(tool_response) + metadata_map.update(response_metadata) + + # Extract referenced documents from metadata with error handling + referenced_documents = [] + for v in metadata_map.values(): + if "docs_url" in v and "title" in v: + try: + doc = ReferencedDocument(doc_url=v["docs_url"], doc_title=v["title"]) + referenced_documents.append(doc) + except (pydantic.ValidationError, ValueError) as e: + logger.warning( + "Skipping invalid referenced document with docs_url='%s', title='%s': %s", + v.get("docs_url", ""), + v.get("title", ""), + str(e), + ) + continue + + return sorted(referenced_documents, key=lambda d: (d.doc_title, str(d.doc_url))) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 7a195beb..d8513d02 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -22,9 +22,9 @@ store_transcript, get_rag_toolgroups, evaluate_model_hints, - _process_knowledge_search_content, extract_referenced_documents_from_steps, ) +from utils.metadata import process_knowledge_search_content from models.requests import QueryRequest, Attachment from models.config import ModelContextProtocolServer @@ -1592,8 +1592,8 @@ def test_evaluate_model_hints( assert model_id == expected_model -def test_process_knowledge_search_content_with_valid_metadata(mocker): - """Test _process_knowledge_search_content with valid metadata.""" +def testprocess_knowledge_search_content_with_valid_metadata(mocker): + """Test process_knowledge_search_content with valid metadata.""" # Mock tool response with valid metadata text_content_item = mocker.Mock() text_content_item.text = """Result 1 @@ -1604,7 +1604,7 @@ def test_process_knowledge_search_content_with_valid_metadata(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) # Verify metadata was correctly parsed and added assert "doc-1" in metadata_map @@ -1613,10 +1613,8 @@ def test_process_knowledge_search_content_with_valid_metadata(mocker): assert metadata_map["doc-1"]["document_id"] == "doc-1" -def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): - """Test _process_knowledge_search_content handles SyntaxError from invalid metadata.""" - mock_logger = mocker.patch("app.endpoints.query.logger") - +def testprocess_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): + """Test process_knowledge_search_content gracefully handles SyntaxError.""" # Mock tool response with invalid metadata (invalid Python syntax) text_content_item = mocker.Mock() text_content_item.text = """Result 1 @@ -1627,21 +1625,14 @@ def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(moc tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) - # Verify metadata_map remains empty due to exception + # Verify metadata_map remains empty due to invalid syntax (gracefully handled) assert len(metadata_map) == 0 - # Verify exception logging was called - mock_logger.exception.assert_called_once() - args = mock_logger.exception.call_args[0] - assert "An exception was thrown in processing" in args[0] - - -def test_process_knowledge_search_content_with_invalid_metadata_value_error(mocker): - """Test _process_knowledge_search_content handles ValueError from invalid metadata.""" - mock_logger = mocker.patch("app.endpoints.query.logger") +def testprocess_knowledge_search_content_with_invalid_metadata_value_error(mocker): + """Test process_knowledge_search_content gracefully handles ValueError from invalid metadata.""" # Mock tool response with invalid metadata containing complex expressions text_content_item = mocker.Mock() text_content_item.text = """Result 1 @@ -1652,19 +1643,14 @@ def test_process_knowledge_search_content_with_invalid_metadata_value_error(mock tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) - # Verify metadata_map remains empty due to exception + # Verify metadata_map remains empty due to invalid expression (gracefully handled) assert len(metadata_map) == 0 - # Verify exception logging was called - mock_logger.exception.assert_called_once() - args = mock_logger.exception.call_args[0] - assert "An exception was thrown in processing" in args[0] - -def test_process_knowledge_search_content_with_non_dict_metadata(mocker): - """Test _process_knowledge_search_content handles non-dict metadata gracefully.""" +def testprocess_knowledge_search_content_with_non_dict_metadata(mocker): + """Test process_knowledge_search_content handles non-dict metadata gracefully.""" mock_logger = mocker.patch("app.endpoints.query.logger") # Mock tool response with metadata that's not a dict @@ -1677,7 +1663,7 @@ def test_process_knowledge_search_content_with_non_dict_metadata(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) # Verify metadata_map remains empty (no document_id in string) assert len(metadata_map) == 0 @@ -1686,8 +1672,8 @@ def test_process_knowledge_search_content_with_non_dict_metadata(mocker): mock_logger.exception.assert_not_called() -def test_process_knowledge_search_content_with_metadata_missing_document_id(mocker): - """Test _process_knowledge_search_content skips metadata without document_id.""" +def testprocess_knowledge_search_content_with_metadata_missing_document_id(mocker): + """Test process_knowledge_search_content skips metadata without document_id.""" # Mock tool response with valid metadata but missing document_id text_content_item = mocker.Mock() text_content_item.text = """Result 1 @@ -1698,39 +1684,39 @@ def test_process_knowledge_search_content_with_metadata_missing_document_id(mock tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) # Verify metadata_map remains empty since document_id is missing assert len(metadata_map) == 0 -def test_process_knowledge_search_content_with_no_text_attribute(mocker): - """Test _process_knowledge_search_content skips content items without text attribute.""" +def testprocess_knowledge_search_content_with_no_text_attribute(mocker): + """Test process_knowledge_search_content skips content items without text attribute.""" # Mock tool response with content item that has no text attribute text_content_item = mocker.Mock(spec=[]) # spec=[] means no attributes tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) # Verify metadata_map remains empty since text attribute is missing assert len(metadata_map) == 0 -def test_process_knowledge_search_content_with_none_content(mocker): - """Test _process_knowledge_search_content handles tool_response with content=None.""" +def testprocess_knowledge_search_content_with_none_content(mocker): + """Test process_knowledge_search_content handles tool_response with content=None.""" # Mock tool response with content = None tool_response = mocker.Mock() tool_response.content = None - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) # Verify metadata_map remains empty when content is None assert len(metadata_map) == 0 -def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker): +def testprocess_knowledge_search_content_duplicate_document_id_last_wins(mocker): """The last metadata block for a given document_id should win.""" text_items = [ mocker.Mock( @@ -1747,7 +1733,7 @@ def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker tool_response.content = text_items # Process content - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) assert metadata_map["doc-x"]["docs_url"] == "https://example.com/second" assert metadata_map["doc-x"]["title"] == "Second" @@ -1761,7 +1747,7 @@ def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker assert docs[0].doc_title == "Second" -def test_process_knowledge_search_content_with_braces_inside_strings(mocker): +def testprocess_knowledge_search_content_with_braces_inside_strings(mocker): """Test that braces inside strings are handled correctly.""" text_content_item = mocker.Mock() text_content_item.text = ( @@ -1773,14 +1759,14 @@ def test_process_knowledge_search_content_with_braces_inside_strings(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) assert "doc-100" in metadata_map assert metadata_map["doc-100"]["title"] == "A {weird} title" assert metadata_map["doc-100"]["docs_url"] == "https://example.com/100" assert metadata_map["doc-100"]["extra"]["note"] == "contains {braces}" -def test_process_knowledge_search_content_with_nested_objects(mocker): +def testprocess_knowledge_search_content_with_nested_objects(mocker): """Test that nested objects are parsed correctly.""" text_content_item = mocker.Mock() text_content_item.text = ( @@ -1792,7 +1778,7 @@ def test_process_knowledge_search_content_with_nested_objects(mocker): tool_response = mocker.Mock() tool_response.content = [text_content_item] - metadata_map = _process_knowledge_search_content(tool_response) + metadata_map = process_knowledge_search_content(tool_response) assert "doc-200" in metadata_map assert metadata_map["doc-200"]["title"] == "Nested JSON" assert metadata_map["doc-200"]["docs_url"] == "https://example.com/200" @@ -1953,7 +1939,7 @@ async def test_retrieve_response_with_structured_content_object( async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, mocker): """Test that retrieve_response skips entries with invalid docs_url.""" # Mock logger to capture warning logs - mock_logger = mocker.patch("app.endpoints.query.logger") + mock_logger = mocker.patch("utils.metadata.logger") mock_client, mock_agent = prepare_agent_mocks mock_agent.create_turn.return_value.output_message.content = "LLM answer" From 0291e11712006af7c1e69c373995b6b387f0c24d Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Mon, 18 Aug 2025 15:52:23 -0400 Subject: [PATCH 14/16] coderabbit --- src/utils/metadata.py | 13 +++-- tests/unit/app/endpoints/test_query.py | 47 +++++++++++++++---- .../app/endpoints/test_streaming_query.py | 6 +-- 3 files changed, 49 insertions(+), 17 deletions(-) diff --git a/src/utils/metadata.py b/src/utils/metadata.py index 6bc49736..991946ff 100644 --- a/src/utils/metadata.py +++ b/src/utils/metadata.py @@ -128,9 +128,16 @@ def process_knowledge_search_content(tool_response: Any) -> dict[str, dict[str, try: content = json.loads(content, strict=False) except (json.JSONDecodeError, TypeError): - # If JSON parsing fails or content is still a string, return empty - if isinstance(content, str): - return metadata_map + # If JSON parsing fails, try parsing as metadata text + try: + parsed_metadata = parse_knowledge_search_metadata(content, strict=False) + metadata_map.update(parsed_metadata) + except ValueError as e: + logger.exception( + "Error processing string content as metadata; position=%s", + getattr(e, "position", "unknown"), + ) + return metadata_map # Ensure content is iterable (but not a string) if isinstance(content, str): diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index d8513d02..c71dedf8 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -1592,7 +1592,7 @@ def test_evaluate_model_hints( assert model_id == expected_model -def testprocess_knowledge_search_content_with_valid_metadata(mocker): +def test_process_knowledge_search_content_with_valid_metadata(mocker): """Test process_knowledge_search_content with valid metadata.""" # Mock tool response with valid metadata text_content_item = mocker.Mock() @@ -1613,7 +1613,7 @@ def testprocess_knowledge_search_content_with_valid_metadata(mocker): assert metadata_map["doc-1"]["document_id"] == "doc-1" -def testprocess_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): +def test_process_knowledge_search_content_with_invalid_metadata_syntax_error(mocker): """Test process_knowledge_search_content gracefully handles SyntaxError.""" # Mock tool response with invalid metadata (invalid Python syntax) text_content_item = mocker.Mock() @@ -1631,7 +1631,7 @@ def testprocess_knowledge_search_content_with_invalid_metadata_syntax_error(mock assert len(metadata_map) == 0 -def testprocess_knowledge_search_content_with_invalid_metadata_value_error(mocker): +def test_process_knowledge_search_content_with_invalid_metadata_value_error(mocker): """Test process_knowledge_search_content gracefully handles ValueError from invalid metadata.""" # Mock tool response with invalid metadata containing complex expressions text_content_item = mocker.Mock() @@ -1649,7 +1649,7 @@ def testprocess_knowledge_search_content_with_invalid_metadata_value_error(mocke assert len(metadata_map) == 0 -def testprocess_knowledge_search_content_with_non_dict_metadata(mocker): +def test_process_knowledge_search_content_with_non_dict_metadata(mocker): """Test process_knowledge_search_content handles non-dict metadata gracefully.""" mock_logger = mocker.patch("app.endpoints.query.logger") @@ -1672,7 +1672,7 @@ def testprocess_knowledge_search_content_with_non_dict_metadata(mocker): mock_logger.exception.assert_not_called() -def testprocess_knowledge_search_content_with_metadata_missing_document_id(mocker): +def test_process_knowledge_search_content_with_metadata_missing_document_id(mocker): """Test process_knowledge_search_content skips metadata without document_id.""" # Mock tool response with valid metadata but missing document_id text_content_item = mocker.Mock() @@ -1690,7 +1690,7 @@ def testprocess_knowledge_search_content_with_metadata_missing_document_id(mocke assert len(metadata_map) == 0 -def testprocess_knowledge_search_content_with_no_text_attribute(mocker): +def test_process_knowledge_search_content_with_no_text_attribute(mocker): """Test process_knowledge_search_content skips content items without text attribute.""" # Mock tool response with content item that has no text attribute text_content_item = mocker.Mock(spec=[]) # spec=[] means no attributes @@ -1704,7 +1704,7 @@ def testprocess_knowledge_search_content_with_no_text_attribute(mocker): assert len(metadata_map) == 0 -def testprocess_knowledge_search_content_with_none_content(mocker): +def test_process_knowledge_search_content_with_none_content(mocker): """Test process_knowledge_search_content handles tool_response with content=None.""" # Mock tool response with content = None tool_response = mocker.Mock() @@ -1716,7 +1716,7 @@ def testprocess_knowledge_search_content_with_none_content(mocker): assert len(metadata_map) == 0 -def testprocess_knowledge_search_content_duplicate_document_id_last_wins(mocker): +def test_process_knowledge_search_content_duplicate_document_id_last_wins(mocker): """The last metadata block for a given document_id should win.""" text_items = [ mocker.Mock( @@ -1747,7 +1747,7 @@ def testprocess_knowledge_search_content_duplicate_document_id_last_wins(mocker) assert docs[0].doc_title == "Second" -def testprocess_knowledge_search_content_with_braces_inside_strings(mocker): +def test_process_knowledge_search_content_with_braces_inside_strings(mocker): """Test that braces inside strings are handled correctly.""" text_content_item = mocker.Mock() text_content_item.text = ( @@ -1766,7 +1766,7 @@ def testprocess_knowledge_search_content_with_braces_inside_strings(mocker): assert metadata_map["doc-100"]["extra"]["note"] == "contains {braces}" -def testprocess_knowledge_search_content_with_nested_objects(mocker): +def test_process_knowledge_search_content_with_nested_objects(mocker): """Test that nested objects are parsed correctly.""" text_content_item = mocker.Mock() text_content_item.text = ( @@ -1785,6 +1785,33 @@ def testprocess_knowledge_search_content_with_nested_objects(mocker): assert metadata_map["doc-200"]["meta"]["k"]["inner"] == 1 +def test_process_knowledge_search_content_with_string_fallback_parsing(mocker): + """Test that string content uses parse_knowledge_search_metadata as fallback.""" + # Create a tool response with string content containing metadata + string_content = """Result 1 +Content: Test content +Metadata: {'docs_url': 'https://example.com/fallback', 'title': 'Fallback Doc', 'document_id': 'fallback-1'} + +Result 2 +Content: More content +Metadata: {'docs_url': 'https://example.com/fallback2', 'title': 'Fallback Doc 2', 'document_id': 'fallback-2'} +""" + + tool_response = mocker.Mock() + tool_response.content = string_content # String instead of list + + metadata_map = process_knowledge_search_content(tool_response) + + # Verify fallback parsing worked correctly + assert len(metadata_map) == 2 + assert "fallback-1" in metadata_map + assert "fallback-2" in metadata_map + assert metadata_map["fallback-1"]["title"] == "Fallback Doc" + assert metadata_map["fallback-1"]["docs_url"] == "https://example.com/fallback" + assert metadata_map["fallback-2"]["title"] == "Fallback Doc 2" + assert metadata_map["fallback-2"]["docs_url"] == "https://example.com/fallback2" + + @pytest.mark.asyncio async def test_retrieve_response_with_none_content(prepare_agent_mocks, mocker): """Test retrieve_response handles None content gracefully.""" diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index 9166d505..94678af6 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -1046,10 +1046,8 @@ def test_stream_build_event_knowledge_search_with_invalid_metadata(mocker): # Verify the function still returns tool execution events assert len(result_list) == 2 # One for tool_calls, one for tool_responses - # Verify exception logging was called - mock_logger.exception.assert_called_once() - args = mock_logger.exception.call_args[0] - assert "An exception was thrown in processing" in args[0] + # Verify no exception logging was called in non-strict mode + mock_logger.exception.assert_not_called() def test_stream_end_event_with_referenced_documents(): From a463e6f63172e530aedca4cde8a2a45628915201 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Mon, 18 Aug 2025 16:33:41 -0400 Subject: [PATCH 15/16] coderabbit --- tests/unit/app/endpoints/test_query.py | 32 ++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index c71dedf8..f988a38f 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -22,9 +22,11 @@ store_transcript, get_rag_toolgroups, evaluate_model_hints, +) +from utils.metadata import ( + process_knowledge_search_content, extract_referenced_documents_from_steps, ) -from utils.metadata import process_knowledge_search_content from models.requests import QueryRequest, Attachment from models.config import ModelContextProtocolServer @@ -270,6 +272,10 @@ async def test_query_endpoint_handler_with_referenced_documents(mocker): ] assert all(isinstance(d, ReferencedDocument) for d in response.referenced_documents) assert len(response.referenced_documents) == 2 + # Titles should be sorted deterministically by doc_title + assert [d.doc_title for d in response.referenced_documents] == sorted( + [d.doc_title for d in response.referenced_documents] + ) # Assert the metric for successful LLM calls is incremented mock_metric.labels("fake_provider_id", "fake_model_id").inc.assert_called_once() @@ -1317,7 +1323,6 @@ async def test_query_endpoint_handler_no_tools_true(mocker): ] mock_config = mocker.Mock() - mock_config.user_data_collection_configuration.transcripts_disabled = True mocker.patch("app.endpoints.query.configuration", mock_config) llm_response = "LLM answer without tools" @@ -1356,7 +1361,6 @@ async def test_query_endpoint_handler_no_tools_false(mocker): ] mock_config = mocker.Mock() - mock_config.user_data_collection_configuration.transcripts_disabled = True mocker.patch("app.endpoints.query.configuration", mock_config) llm_response = "LLM answer with tools" @@ -1651,7 +1655,7 @@ def test_process_knowledge_search_content_with_invalid_metadata_value_error(mock def test_process_knowledge_search_content_with_non_dict_metadata(mocker): """Test process_knowledge_search_content handles non-dict metadata gracefully.""" - mock_logger = mocker.patch("app.endpoints.query.logger") + mock_logger = mocker.patch("utils.metadata.logger") # Mock tool response with metadata that's not a dict text_content_item = mocker.Mock() @@ -1812,6 +1816,24 @@ def test_process_knowledge_search_content_with_string_fallback_parsing(mocker): assert metadata_map["fallback-2"]["docs_url"] == "https://example.com/fallback2" +def test_process_knowledge_search_content_metadata_label_case_insensitive(mocker): + """Test that metadata labels are detected case-insensitively.""" + text_content_item = mocker.Mock() + text_content_item.text = ( + "Result 1\n" + "Content: Test content\n" + "metadata: {'document_id': 'doc-ci', 'title': 'Case Insensitive', 'docs_url': 'https://example.com/ci'}\n" + ) + tool_response = mocker.Mock() + tool_response.content = [text_content_item] + + metadata_map = process_knowledge_search_content(tool_response) + + assert "doc-ci" in metadata_map + assert metadata_map["doc-ci"]["title"] == "Case Insensitive" + assert metadata_map["doc-ci"]["docs_url"] == "https://example.com/ci" + + @pytest.mark.asyncio async def test_retrieve_response_with_none_content(prepare_agent_mocks, mocker): """Test retrieve_response handles None content gracefully.""" @@ -1994,7 +2016,7 @@ async def test_retrieve_response_skips_invalid_docs_url(prepare_agent_mocks, moc mock_tool_response.call_id = "c1" mock_tool_response.tool_name = "knowledge_search" mock_tool_response.content = [ - mocker.Mock(text=s, type="text") for s in invalid_docs_url_results + TextContentItem(text=s, type="text") for s in invalid_docs_url_results ] mock_tool_execution_step = mocker.Mock() From 622ae9d7916900255371a6376349688fabfa48f8 Mon Sep 17 00:00:00 2001 From: Erik M Jacobs Date: Tue, 19 Aug 2025 08:34:47 -0400 Subject: [PATCH 16/16] fix lint --- tests/unit/app/endpoints/test_query.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index fba130a1..4d948595 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -501,12 +501,13 @@ async def test_retrieve_response_no_returned_message(prepare_agent_mocks, mocker model_id = "fake_model_id" access_token = "test_token" - response, _ = await retrieve_response( + response, _, referenced_documents = await retrieve_response( mock_client, model_id, query_request, access_token ) # fallback mechanism: check that the response is empty assert response == "" + assert referenced_documents == [] @pytest.mark.asyncio @@ -532,12 +533,13 @@ async def test_retrieve_response_message_without_content(prepare_agent_mocks, mo model_id = "fake_model_id" access_token = "test_token" - response, _ = await retrieve_response( + response, _, referenced_documents = await retrieve_response( mock_client, model_id, query_request, access_token ) # fallback mechanism: check that the response is empty assert response == "" + assert referenced_documents == [] @pytest.mark.asyncio @@ -1884,7 +1886,8 @@ def test_process_knowledge_search_content_metadata_label_case_insensitive(mocker text_content_item.text = ( "Result 1\n" "Content: Test content\n" - "metadata: {'document_id': 'doc-ci', 'title': 'Case Insensitive', 'docs_url': 'https://example.com/ci'}\n" + "metadata: {'document_id': 'doc-ci', 'title': 'Case Insensitive', " + "'docs_url': 'https://example.com/ci'}\n" ) tool_response = mocker.Mock() tool_response.content = [text_content_item]