From 8dba8b08e34be5927b89636a551cb2436699f3bd Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Tue, 14 Oct 2025 22:10:35 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=91=EF=B8=8F(mistral)=20attempt=20to?= =?UTF-8?q?=20fix=20Unmarshaller=20error?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mistral model sometime returns `{'type': 'reference', 'reference_ids': ['cuisineaz']}` which does not match the Mistral SDK format for `ReferenceChunk`: `reference_ids: List[int]``... The Mistral documentation provides a way to return citation (even if we don't want to use this systemp), so we expect the model to behave properly with this change. --- .../chat/tests/tools/test_web_search_brave.py | 52 +++++++++---------- src/backend/chat/tools/web_search_brave.py | 50 ++++++++---------- 2 files changed, 47 insertions(+), 55 deletions(-) diff --git a/src/backend/chat/tests/tools/test_web_search_brave.py b/src/backend/chat/tests/tools/test_web_search_brave.py index 3c53f59..d4acfab 100644 --- a/src/backend/chat/tests/tools/test_web_search_brave.py +++ b/src/backend/chat/tests/tools/test_web_search_brave.py @@ -69,18 +69,14 @@ def test_agent_web_search_brave_success_with_extra_snippets(): tool_return = web_search_brave("test query") assert hasattr(tool_return, "return_value") - assert tool_return.return_value == [ - { - "link": "https://example.com/a", + assert tool_return.return_value == { + "0": { + "snippets": ["Snippet A1", "Snippet A2"], "title": "Result A", - "extra_snippets": ["Snippet A1", "Snippet A2"], + "url": "https://example.com/a", }, - { - "link": "https://example.com/b", - "title": "Result B", - "extra_snippets": ["Snippet B1"], - }, - ] + "1": {"snippets": ["Snippet B1"], "title": "Result B", "url": "https://example.com/b"}, + } assert tool_return.metadata["sources"] == {"https://example.com/a", "https://example.com/b"} # Check request parameters @@ -120,13 +116,13 @@ def test_agent_web_search_brave_success_without_extra_snippets(): ) tool_return = web_search_brave("test query") - assert tool_return.return_value == [ - { - "link": "https://example.com/c", + assert tool_return.return_value == { + "0": { + "snippets": ["Extracted Content C\nlink"], "title": "Result C", - "extra_snippets": ["Extracted Content C\nlink"], + "url": "https://example.com/c", } - ] + } assert tool_return.metadata["sources"] == {"https://example.com/c"} @@ -159,13 +155,13 @@ def test_agent_web_search_brave_success_without_extra_snippets_summarization(set mock_llm_summarize.assert_called_with("test query", "Extracted Content C\nlink") - assert tool_return.return_value == [ - { - "link": "https://example.com/c", + assert tool_return.return_value == { + "0": { + "snippets": ["Summarized extracted Content C\nlink"], "title": "Result C", - "extra_snippets": ["Summarized extracted Content C\nlink"], + "url": "https://example.com/c", } - ] + } assert tool_return.metadata["sources"] == {"https://example.com/c"} @@ -179,7 +175,7 @@ def test_agent_web_search_brave_empty_results(): status=200, ) tool_return = web_search_brave("empty query") - assert tool_return.return_value == [] + assert tool_return.return_value == {} assert tool_return.metadata["sources"] == set() @@ -282,7 +278,7 @@ def test_agent_web_search_brave_single_worker(settings): tool_return = web_search_brave("single worker query") assert len(tool_return.return_value) == 1 - assert tool_return.return_value[0]["extra_snippets"] == ["Single Content"] + assert tool_return.return_value["0"]["snippets"] == ["Single Content"] @responses.activate @@ -377,10 +373,10 @@ def test_web_search_brave_with_document_backend_success(): tool_return = web_search_brave_with_document_backend(mock_ctx, "rag query") assert len(tool_return.return_value) == 2 - assert tool_return.return_value[0]["link"] == "https://example.com/doc1" - assert tool_return.return_value[0]["extra_snippets"] == ["RAG Content 1"] - assert tool_return.return_value[1]["link"] == "https://example.com/doc2" - assert tool_return.return_value[1]["extra_snippets"] == ["RAG Content 2"] + assert tool_return.return_value["0"]["url"] == "https://example.com/doc1" + assert tool_return.return_value["0"]["snippets"] == ["RAG Content 1"] + assert tool_return.return_value["1"]["url"] == "https://example.com/doc2" + assert tool_return.return_value["1"]["snippets"] == ["RAG Content 2"] assert tool_return.metadata["sources"] == { "https://example.com/doc1", "https://example.com/doc2", @@ -471,7 +467,7 @@ def test_web_search_brave_with_document_backend_fetch_error(settings): tool_return = web_search_brave_with_document_backend(mock_ctx, "error query") # Should complete despite error (error is caught and logged in multi-worker path) - assert tool_return.return_value == [] + assert tool_return.return_value == {} @responses.activate @@ -517,7 +513,7 @@ def test_web_search_brave_with_document_backend_no_matching_rag_results(): tool_return = web_search_brave_with_document_backend(mock_ctx, "query") # No results should be returned since RAG URL doesn't match search results - assert tool_return.return_value == [] + assert tool_return.return_value == {} assert tool_return.metadata["sources"] == set() diff --git a/src/backend/chat/tools/web_search_brave.py b/src/backend/chat/tools/web_search_brave.py index 7d15f02..bd54764 100644 --- a/src/backend/chat/tools/web_search_brave.py +++ b/src/backend/chat/tools/web_search_brave.py @@ -119,6 +119,27 @@ def _query_brave_api(query: str) -> List[dict]: return json_response.get("web", {}).get("results", []) +def format_tool_return(raw_search_results: List[dict]) -> ToolReturn: + """Format the raw search results into a ToolReturn object.""" + return ToolReturn( + # Format return value "mistral-like": https://docs.mistral.ai/capabilities/citations/ + return_value={ + str(idx): { + "url": result["url"], + "title": result["title"], + "snippets": result.get("extra_snippets", []), + } + for idx, result in enumerate(raw_search_results) + if result.get("extra_snippets", []) + }, + metadata={ + "sources": { + result["url"] for result in raw_search_results if result.get("extra_snippets", []) + } + }, + ) + + def web_search_brave(query: str) -> ToolReturn: """ Search the web for up-to-date information @@ -154,17 +175,7 @@ def web_search_brave(query: str) -> ToolReturn: idx = future_map[future] raw_search_results[idx]["extra_snippets"] = future.result() - return ToolReturn( - return_value=[ - { - "link": result["url"], - "title": result["title"], - "extra_snippets": result.get("extra_snippets", []), - } - for result in raw_search_results - ], - metadata={"sources": {result["url"] for result in raw_search_results}}, - ) + return format_tool_return(raw_search_results) def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn: @@ -214,19 +225,4 @@ def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolR result.setdefault("extra_snippets", []).append(rag_result.content) break - return ToolReturn( - return_value=[ - { - "link": result["url"], - "title": result["title"], - "extra_snippets": result.get("extra_snippets", []), - } - for result in raw_search_results - if result.get("extra_snippets", []) - ], - metadata={ - "sources": { - result["url"] for result in raw_search_results if result.get("extra_snippets", []) - } - }, - ) + return format_tool_return(raw_search_results)