From 8358ca9490abd1f693a623f06ba0807aa3678cd2 Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Thu, 31 Jul 2025 23:45:02 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B(rag-websearch)=20fix=20the=20Alber?= =?UTF-8?q?t=20API=20search=20query?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This was not properly tested and the code changed just before last commit... boom. --- .../chat/agent_rag/albert_api_constants.py | 33 +++--- .../chat/agent_rag/web_search/albert_api.py | 2 +- src/backend/chat/tests/agent_rag/__init__.py | 0 .../agent_rag/test_albert_api_constants.py | 60 ++++++++++ .../tests/agent_rag/web_search/__init__.py | 0 .../agent_rag/web_search/test_albert_api.py | 105 ++++++++++++++++++ 6 files changed, 182 insertions(+), 18 deletions(-) create mode 100644 src/backend/chat/tests/agent_rag/__init__.py create mode 100644 src/backend/chat/tests/agent_rag/test_albert_api_constants.py create mode 100644 src/backend/chat/tests/agent_rag/web_search/__init__.py create mode 100644 src/backend/chat/tests/agent_rag/web_search/test_albert_api.py diff --git a/src/backend/chat/agent_rag/albert_api_constants.py b/src/backend/chat/agent_rag/albert_api_constants.py index d055507..1eb3f00 100644 --- a/src/backend/chat/agent_rag/albert_api_constants.py +++ b/src/backend/chat/agent_rag/albert_api_constants.py @@ -1,9 +1,9 @@ """Constants and schemas for the Albert RAG agent from Albert API codebase.""" from enum import Enum -from typing import Any, Dict, List, Literal, Optional +from typing import Annotated, Any, Dict, List, Literal, Optional, Self -from pydantic import BaseModel, Field, constr, model_validator +from pydantic import BaseModel, Field, StringConstraints, model_validator # - app/schemas/chunks.py @@ -98,15 +98,13 @@ class SearchArgs(BaseModel): rff_k: int = Field(default=20, description="k constant in RFF algorithm") k: int = Field(gt=0, default=4, description="Number of results to return") method: SearchMethod = Field(default=SearchMethod.SEMANTIC) - score_threshold: Optional[float] = ( - Field( - default=0.0, - ge=0.0, - le=1.0, - description=( - "Score of cosine similarity threshold for filtering results, " - "only available for semantic search method." - ), + score_threshold: Optional[float] = Field( + default=0.0, + ge=0.0, + le=1.0, + description=( + "Score of cosine similarity threshold for filtering results, " + "only available for semantic search method." ), ) web_search: bool = Field( @@ -115,24 +113,25 @@ class SearchArgs(BaseModel): web_search_k: int = Field(default=5, description="Number of results to return for web search.") @model_validator(mode="after") - def score_threshold_filter(self, values): + def score_threshold_filter(self) -> Self: """Validate the score threshold based on the search method.""" - if values.score_threshold and values.method not in ( + if self.score_threshold and self.method not in ( SearchMethod.SEMANTIC, SearchMethod.MULTIAGENT, ): raise ValueError( "Score threshold is only available for semantic and multiagent search methods." ) - return values + return self class SearchRequest(SearchArgs): """Model representing a search request in the Albert API.""" - prompt: constr(strip_whitespace=True, min_length=1) = Field( - description="Prompt related to the search" - ) + prompt: Annotated[ + str, + StringConstraints(strip_whitespace=True, min_length=1), + ] = Field(description="Prompt related to the search") class Search(BaseModel): diff --git a/src/backend/chat/agent_rag/web_search/albert_api.py b/src/backend/chat/agent_rag/web_search/albert_api.py index 2840bb5..9c56db6 100644 --- a/src/backend/chat/agent_rag/web_search/albert_api.py +++ b/src/backend/chat/agent_rag/web_search/albert_api.py @@ -90,7 +90,7 @@ class AlbertWebSearchManager(BaseWebSearchManager): response = requests.post( self._search_endpoint, headers=self._headers, - json=search_request.model_dump(), + json=search_request.model_dump(mode="json", exclude_unset=True), timeout=settings.ALBERT_API_TIMEOUT, ) response.raise_for_status() diff --git a/src/backend/chat/tests/agent_rag/__init__.py b/src/backend/chat/tests/agent_rag/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/tests/agent_rag/test_albert_api_constants.py b/src/backend/chat/tests/agent_rag/test_albert_api_constants.py new file mode 100644 index 0000000..c253de5 --- /dev/null +++ b/src/backend/chat/tests/agent_rag/test_albert_api_constants.py @@ -0,0 +1,60 @@ +"""Test suite for albert_api_constants.py.""" + +import pytest +from pydantic import ValidationError + +from chat.agent_rag.albert_api_constants import ( + SearchArgs, + SearchMethod, + SearchRequest, +) + + +def test_search_request_model_dump_json(): + """Test that SearchRequest.model_dump(mode='json') works correctly.""" + search_request = SearchRequest(prompt="test prompt") + assert search_request.model_dump(mode="json") == { + "collections": [], + "k": 4, + "method": "semantic", + "prompt": "test prompt", + "rff_k": 20, + "score_threshold": 0.0, + "web_search": False, + "web_search_k": 5, + } + + +def test_search_args_score_threshold_valid(): + """Test that score_threshold is valid for semantic and multiagent search methods.""" + try: + SearchArgs(method=SearchMethod.SEMANTIC, score_threshold=0.5) + SearchArgs(method=SearchMethod.MULTIAGENT, score_threshold=0.5) + except ValidationError: + pytest.fail("ValidationError was raised unexpectedly for valid search methods.") + + +def test_search_args_score_threshold_invalid(): + """Test that score_threshold raises ValueError for hybrid and lexical search methods.""" + with pytest.raises(ValidationError) as excinfo: + SearchArgs(method=SearchMethod.HYBRID, score_threshold=0.5) + assert "Score threshold is only available for semantic and multiagent search methods." in str( + excinfo.value + ) + + with pytest.raises(ValidationError) as excinfo: + SearchArgs(method=SearchMethod.LEXICAL, score_threshold=0.5) + assert "Score threshold is only available for semantic and multiagent search methods." in str( + excinfo.value + ) + + +def test_search_args_no_score_threshold(): + """Test that no error is raised when score_threshold is not set.""" + try: + SearchArgs(method=SearchMethod.HYBRID) + SearchArgs(method=SearchMethod.LEXICAL) + SearchArgs(method=SearchMethod.SEMANTIC) + SearchArgs(method=SearchMethod.MULTIAGENT) + except ValidationError: + pytest.fail("ValidationError was raised unexpectedly when score_threshold is not set.") diff --git a/src/backend/chat/tests/agent_rag/web_search/__init__.py b/src/backend/chat/tests/agent_rag/web_search/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/tests/agent_rag/web_search/test_albert_api.py b/src/backend/chat/tests/agent_rag/web_search/test_albert_api.py new file mode 100644 index 0000000..5f9a1cd --- /dev/null +++ b/src/backend/chat/tests/agent_rag/web_search/test_albert_api.py @@ -0,0 +1,105 @@ +"""Unit tests for the Albert API web search manager.""" + +import json + +import pytest +import requests +import responses + +from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage +from chat.agent_rag.web_search.albert_api import AlbertWebSearchManager + + +@pytest.fixture(autouse=True) +def albert_api_settings(settings): + """Fixture to set Albert API settings for tests.""" + settings.ALBERT_API_URL = "http://test-albert-api.com" + settings.ALBERT_API_KEY = "test-key" + + +@pytest.mark.parametrize( + "url, expected", + [ + ("http://example.com/page.html", "http://example.com/page"), + ("http://example.com/page", "http://example.com/page"), + ("http://example.com/.html", "http://example.com/"), + ], +) +def test_clean_url(url, expected): + """Test the _clean_url static method.""" + assert AlbertWebSearchManager._clean_url(url) == expected # pylint: disable=protected-access + + +@responses.activate +def test_web_search_success(settings): + """Test a successful web search.""" + settings.RAG_WEB_SEARCH_MAX_RESULTS = 20 + settings.RAG_WEB_SEARCH_CHUNK_NUMBER = 10 + + mock_albert_api = responses.post( + "http://test-albert-api.com/v1/search", + json={ + "data": [ + { + "method": "semantic", + "chunk": { + "id": 123, + "content": "This is a test chunk.", + "metadata": { + "document_name": "http://example.com/test.html", + "document_type": "html", + }, + }, + "score": 0.9, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + status=200, + content_type="application/json", + ) + + results = AlbertWebSearchManager().web_search("test query") + assert results == RAGWebResults( + data=[ + RAGWebResult(url="http://example.com/test", content="This is a test chunk.", score=0.9) + ], + usage=RAGWebUsage(prompt_tokens=10, completion_tokens=20), + ) + + # Verify the request payload + request = mock_albert_api.calls[0].request + assert json.loads(request.body) == { + "prompt": "test query", + "web_search": True, + "web_search_k": 20, # Default value from settings + "k": 10, # Default value from settings + } + + +def test_web_search_empty_query(): + """Test web_search with an empty query.""" + with pytest.raises(ValueError, match="Search query cannot be empty."): + AlbertWebSearchManager().web_search(" ") + + +@responses.activate +def test_web_search_http_error(): + """Test web_search with an HTTP error from the API.""" + responses.post("http://test-albert-api.com/v1/search", status=500) + with pytest.raises(requests.HTTPError): + AlbertWebSearchManager().web_search("test query") + + +@responses.activate +def test_web_search_json_decode_error(): + """Test web_search with a JSON decode error from the API.""" + responses.post( + "http://test-albert-api.com/v1/search", + body="invalid json", + status=200, + content_type="application/json", + ) + + with pytest.raises(requests.exceptions.JSONDecodeError): + AlbertWebSearchManager().web_search("test query")