diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f259668..9df11b57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to - ✨(back) add model fallback mechanism - ✨(back) add celery for running background tasks - 🧱(helm) add celery worker and beat deployments +- ✨(back) add Staan web search tool ### Changed diff --git a/env.d/development/kube-secret.dist b/env.d/development/kube-secret.dist index 0f470ac2..edcd304a 100644 --- a/env.d/development/kube-secret.dist +++ b/env.d/development/kube-secret.dist @@ -3,4 +3,5 @@ AI_BASE_URL=changeme AI_API_KEY=changeme ALBERT_API_URL=changeme ALBERT_API_KEY=changeme -BRAVE_API_KEY=changeme \ No newline at end of file +BRAVE_API_KEY=changeme +STAAN_API_KEY=changeme \ No newline at end of file diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index c222ed59..4d451544 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -325,6 +325,7 @@ def __init__( # pylint: disable=too-many-arguments,too-many-positional-argument user=user, session=session, web_search_enabled=self._is_web_search_enabled and self._is_smart_search_enabled, + language=self.language, ) self._web_search_tool_registered = False self._self_documentation_tool_registered = False diff --git a/src/backend/chat/clients/schema.py b/src/backend/chat/clients/schema.py index 8b1d4642..439f86d8 100644 --- a/src/backend/chat/clients/schema.py +++ b/src/backend/chat/clients/schema.py @@ -52,6 +52,7 @@ class ContextDeps: user: User session: Optional[Dict] = None web_search_enabled: bool = False + language: str | None = None @dataclasses.dataclass diff --git a/src/backend/chat/tests/tools/test_web_search_staan.py b/src/backend/chat/tests/tools/test_web_search_staan.py new file mode 100644 index 00000000..3d471778 --- /dev/null +++ b/src/backend/chat/tests/tools/test_web_search_staan.py @@ -0,0 +1,76 @@ +"""Tests for the Staan web search tool.""" + +from unittest.mock import Mock, patch + +import pytest + +from chat.tools.web_search_staan import resolve_staan_market, staan_search + +STAAN_WEB_SEARCH_URL = "https://api.staan.ai/v2/search/web" + + +@pytest.fixture(autouse=True) +def staan_settings(settings): + """Define Staan settings for tests.""" + settings.STAAN_API_KEY = "test-staan-key" + settings.STAAN_SEARCH_ENDPOINT = STAAN_WEB_SEARCH_URL + settings.STAAN_SEARCH_EXTRA_SNIPPETS = True + settings.STAAN_API_TIMEOUT = 5 + settings.STAAN_MAX_RESULTS = 10 + settings.STAAN_MAX_SNIPPET_LENGTH = 5000 + + +@pytest.mark.parametrize( + ("language", "expected_market"), + [ + ("fr-fr", "fr-fr"), + ("en-us", "en-us"), + ("de-de", "de-de"), + ("FR-FR", "fr-fr"), + ("nl-nl", "en-us"), + ], +) +def test_resolve_staan_market_from_language(settings, language, expected_market): + """Language should map to a supported Staan market.""" + settings.LANGUAGE_CODE = "en-us" + + assert resolve_staan_market(language) == expected_market + + +def test_resolve_staan_market_falls_back_to_language_code(settings): + """Missing language should fall back to Django LANGUAGE_CODE.""" + settings.LANGUAGE_CODE = "en-us" + + assert resolve_staan_market(None) == "en-us" + + +def test_resolve_staan_market_falls_back_to_english_for_unsupported_language(settings): + """Unsupported languages should fall back to English.""" + settings.LANGUAGE_CODE = "nl-nl" + + assert resolve_staan_market(None) == "en-us" + + +@patch("chat.tools.web_search_staan.requests.get") +def test_staan_search_sends_market_query_param(mock_get): + """Market must be forwarded to the Staan API as a query parameter.""" + mock_response = Mock() + mock_response.raise_for_status = Mock() + mock_response.json.return_value = { + "query": {"q": "climate tech", "market": "en-us"}, + "web": {"results": []}, + } + mock_get.return_value = mock_response + + staan_search("climate tech", "en-us") + + mock_get.assert_called_once_with( + STAAN_WEB_SEARCH_URL, + params={ + "q": "climate tech", + "market": "en-us", + "extra_snippets": "true", + }, + headers={"Authorization": "Bearer test-staan-key"}, + timeout=5, + ) diff --git a/src/backend/chat/tools/descriptions.py b/src/backend/chat/tools/descriptions.py index 1402f076..773aebae 100644 --- a/src/backend/chat/tools/descriptions.py +++ b/src/backend/chat/tools/descriptions.py @@ -108,6 +108,8 @@ - "Explique-moi comment fonctionne une boucle for" - "Écris-moi un poème sur l'automne" - "Résume ce texte" + +When using web_search tool, you can retry the search with a different query if the first one didn't return any relevant results. """ SELF_DOCUMENTATION_SYSTEM_PROMPT = ( diff --git a/src/backend/chat/tools/web_search_staan.py b/src/backend/chat/tools/web_search_staan.py new file mode 100644 index 00000000..00433505 --- /dev/null +++ b/src/backend/chat/tools/web_search_staan.py @@ -0,0 +1,203 @@ +"""Web juridique tool for the chat agent.""" + +import json +import logging + +from django.conf import settings + +import requests +from pydantic_ai import RunContext +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.messages import ToolReturn + +from chat.tools.exceptions import ModelCannotRetry +from chat.tools.utils import last_model_retry_soft_fail + +logger = logging.getLogger(__name__) + +STAAN_MARKETS = frozenset({"fr-fr", "en-us", "de-de"}) +STAAN_DEFAULT_MARKET = "en-us" + +_STAAN_MARKET_BY_LANGUAGE_PREFIX = { + "fr": "fr-fr", + "en": "en-us", + "de": "de-de", +} + + +def resolve_staan_market(language: str | None) -> str: + """Map a UI language code to a supported Staan search market.""" + user_lang = (language or settings.LANGUAGE_CODE or "").lower() + if user_lang in STAAN_MARKETS: + return user_lang + prefix = user_lang.split("-")[0] if user_lang else "" + if prefix in _STAAN_MARKET_BY_LANGUAGE_PREFIX: + return _STAAN_MARKET_BY_LANGUAGE_PREFIX[prefix] + return STAAN_DEFAULT_MARKET + + +def _resolve_staan_market(ctx: RunContext) -> str: + """Resolve the Staan market from the conversation context language.""" + return resolve_staan_market(getattr(ctx.deps, "language", None)) + + +def staan_search(query: str, market: str) -> requests.Response: + """ + Performs a search using the Staan API. + + Args: + query: User query string. + market: Staan search market (e.g. fr-fr, en-us, de-de). + + Returns: + requests.Response: Raw HTTP response from the Staan API. + + """ + if not settings.STAAN_API_KEY: + raise ValueError("Clé API Staan manquante (variable d'env STAAN_API_KEY)") + + params = { + "q": query, + "market": market, + "extra_snippets": "true" if settings.STAAN_SEARCH_EXTRA_SNIPPETS else "false", + } + + headers = { + "Authorization": f"Bearer {settings.STAAN_API_KEY}", + } + + response = requests.get( + settings.STAAN_SEARCH_ENDPOINT, + params=params, + headers=headers, + timeout=settings.STAAN_API_TIMEOUT, + ) + response.raise_for_status() + return response + + +def _collect_extra_snippets( + result: dict, + *, + max_len_snippet: int, + min_score: float, +) -> list[str]: + """Extract extra snippet chunks from a Staan result, filtering by score and length.""" + raw_snippets = result.get("extra_snippets") or [] + if not raw_snippets: + return [] + + extra_snippets: list[str] = [] + for item in raw_snippets: + if isinstance(item, dict): + chunk = item.get("chunk", "") + score = float(item.get("score", 0)) + else: + chunk = str(item) + score = min_score + + if score < min_score: + continue + + current_length = len(" ".join(extra_snippets)) + if current_length + len(chunk) >= max_len_snippet: + break + extra_snippets.append(chunk) + + return extra_snippets + + +def format_staan( + response: requests.Response, + n_results: int | None = None, + max_len_snippet: int | None = None, + min_score: float | None = None, +) -> str: + """ + Format a Staan API response to extract web results with snippets. + + Works whether or not ``extra_snippets`` was requested from the API: + - without: uses the main ``snippet`` field only + - with: adds filtered ``extra_snippets`` chunks (dict items with ``chunk`` / ``score``) + + Args: + response: requests.Response object from Staan API + n_results: Maximum number of web results to include + max_len_snippet: Maximum total length of concatenated extra snippets per result + min_score: Minimum relevance score for an extra snippet chunk to be included + + Returns: + str: JSON string of cleaned results with title, url, snippet, + published_date and extra_snippets + """ + n_results = n_results if n_results is not None else settings.STAAN_MAX_RESULTS + max_len_snippet = ( + max_len_snippet if max_len_snippet is not None else settings.STAAN_MAX_SNIPPET_LENGTH + ) + min_score = min_score if min_score is not None else 0 + + data = response.json().get("web", {}).get("results", []) + results = [] + for result in data[:n_results]: + output = { + "title": result.get("title", ""), + "url": result.get("url", ""), + "snippet": result.get("snippet", ""), + "published_date": result.get("published_date", ""), + "extra_snippets": _collect_extra_snippets( + result, + max_len_snippet=max_len_snippet, + min_score=min_score, + ), + } + results.append(output) + return json.dumps(results, ensure_ascii=False, indent=2) + + +@last_model_retry_soft_fail +async def web_search_staan(ctx: RunContext, query: str) -> ToolReturn: + """ + Search the web using the Staan API. + + Args: + ctx: Execution context used to resolve the user's search market. + query: Search query. Max 400 characters. Use site:example.com to restrict to a domain. + + Returns: + ToolReturn: Result of the search. + """ + market = _resolve_staan_market(ctx) + logger.info("Staan web search: query=%r market=%s", query, market) + try: + response = staan_search(query, market) + response_data = response.json() + api_market = response_data.get("query", {}).get("market") + logger.info("Staan API confirmed market=%s (requested=%s)", api_market, market) + sources = list({resp.get("url", "") for resp in response_data["web"]["results"]}) + return ToolReturn( + return_value=format_staan(response), + metadata={"sources": sources}, + ) + except requests.HTTPError as exc: + status_code = exc.response.status_code if exc.response is not None else None + logger.warning("Staan API HTTP error: status=%s market=%s", status_code, market) + if status_code == 429: + raise ModelRetry( + "The search API is rate limited. Please wait a moment and try again." + ) from exc + if status_code is not None and status_code >= 500: + raise ModelRetry( + "The search service is temporarily unavailable due to a server error. Retrying..." + ) from exc + raise ModelCannotRetry( + f"Web search failed with a client error (status {status_code}). " + "You must explain this to the user and not try to answer based on your knowledge." + ) from exc + except ModelCannotRetry, ModelRetry: + raise + except Exception as exc: + logger.exception("Unexpected error in web_search_staan: %s", exc) + raise ModelCannotRetry( + f"An unexpected error occurred during web search: {type(exc).__name__}. " + "You must explain this to the user and not try to answer based on your knowledge." + ) from exc diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index d12ed2d7..70934d09 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -30,6 +30,7 @@ from chat.llm_configuration import cached_load_llm_configuration, load_llm_configuration from conversations.brave_settings import BraveSettings +from conversations.staan_settings import StaanSettings # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -48,7 +49,7 @@ def get_release(): return "NA" # Default: not available -class Base(BraveSettings, Configuration): +class Base(BraveSettings, StaanSettings, Configuration): """ This is the base configuration every configuration (aka environment) should inherit from. It is recommended to configure third-party applications by creating a configuration mixins in diff --git a/src/backend/conversations/staan_settings.py b/src/backend/conversations/staan_settings.py new file mode 100644 index 00000000..80d7139a --- /dev/null +++ b/src/backend/conversations/staan_settings.py @@ -0,0 +1,44 @@ +"""Django configuration mixin for Staan settings.""" + +from configurations import values + + +class StaanSettings: + """Staan settings for web_search_staan tool.""" + + STAAN_API_KEY = values.Value( + default=None, + environ_name="STAAN_API_KEY", + environ_prefix=None, + ) + STAAN_API_TIMEOUT = values.IntegerValue( + default=20, + environ_name="STAAN_API_TIMEOUT", + environ_prefix=None, + ) + STAAN_SEARCH_ENDPOINT = values.Value( + default="https://api.staan.ai/v2/search/web", + environ_name="STAAN_SEARCH_ENDPOINT", + environ_prefix=None, + ) + STAAN_SEARCH_MARKET = values.Value( + default="fr-fr", + environ_name="STAAN_SEARCH_MARKET", + environ_prefix=None, + ) + STAAN_SEARCH_EXTRA_SNIPPETS = values.BooleanValue( + default=True, + environ_name="STAAN_SEARCH_EXTRA_SNIPPETS", + environ_prefix=None, + ) + STAAN_MAX_RESULTS = values.IntegerValue( + default=10, + environ_name="STAAN_MAX_RESULTS", + environ_prefix=None, + ) + STAAN_MAX_SNIPPET_LENGTH = values.IntegerValue( + default=5000, + help_text="Maximum length of the snippets per url to return (in characters)", + environ_name="STAAN_MAX_SNIPPET_LENGTH", + environ_prefix=None, + )