Prevent long context

pytest.ini

@@ -26,6 +26,7 @@ env =
     USE_CACHED_SETTINGS=True
     DATA_COLLECTION_ORIGIN_PREFIX=workshop
     ENV=test
+    LANGSMITH_TRACING_ENABLED=false
 
 filterwarnings =
     ignore:.*U.*mode is deprecated:DeprecationWarning

src/app/bibliography/api/bibliography.py

@@ -1,8 +1,8 @@
 from fastapi import APIRouter
 from starlette.responses import PlainTextResponse
 
-from src.app.bibliography.models.bibliography import DocumentIDs
 from src.app.bibliography.helpers.helpers import welearn_document_to_ris
+from src.app.bibliography.models.bibliography import DocumentIDs
 from src.app.services.sql_db.queries import get_documents_by_ids
 
 router = APIRouter()

src/app/bibliography/helpers/helpers.py

@@ -1,7 +1,6 @@
 import datetime
 from typing import Any, Optional
 
-
 from welearn_database.data.models import WeLearnDocument
 
 from src.app.utils.decorators import log_time_and_error_sync

src/app/shared/infra/abst_chat.py

@@ -22,6 +22,7 @@
 
 from fastapi import BackgroundTasks, Depends, Request
 from langchain.agents import create_agent  # type: ignore
+from langchain.agents.middleware import SummarizationMiddleware  # type: ignore
 from langchain.messages import HumanMessage  # type: ignore
 from langchain_core.messages import BaseMessage  # type: ignore
 from langchain_core.runnables import RunnableConfig  # type: ignore
@@ -613,6 +614,12 @@ async def _create_agent(
             tools=[
                 get_resources_about_sustainability,
             ],
+            middleware=[
+                SummarizationMiddleware(
+                    model=agent_model,
+                    trigger=("tokens", 64000),
+                )
+            ],
             checkpointer=memory,
             system_prompt=prompts.AGENT_SYSTEM_PROMPT,
         )

src/app/tests/services/test_abst_chat.py

@@ -184,3 +184,19 @@ async def test_chat_message(self):
             mock_detect_lang.assert_called_with("this is a query")
             self.chat.chat_client.completion.assert_not_called()
             self.chat.chat_client.completion_stream.assert_called_once()
+
+    @mock.patch("src.app.shared.infra.abst_chat.create_agent")
+    @mock.patch("src.app.shared.infra.abst_chat.ChatMistralAI")
+    async def test_create_agent_adds_summarization_middleware(
+        self, mock_chat_mistral, mock_create_agent
+    ):
+        mocked_model = mock.Mock()
+        mocked_model._llm_type = "mistral-chat"  # noqa: SLF001
+        mock_chat_mistral.return_value = mocked_model
+        mock_create_agent.return_value = object()
+
+        await self.chat._create_agent(memory=None)
+
+        middleware = mock_create_agent.call_args.kwargs["middleware"]
+        assert len(middleware) == 1
+        assert middleware[0].__class__.__name__ == "SummarizationMiddleware"

src/app/tests/services/test_helpers.py

@@ -5,6 +5,13 @@
 import numpy
 from langdetect.language import Language
 
+from src.app.bibliography.helpers.helpers import (
+    compute_authors_for_ris,
+    compute_publication_date_for_ris,
+    compute_ris_doctype,
+    ris_line,
+    welearn_document_to_ris,
+)
 from src.app.models.documents import Document, DocumentPayloadModel
 from src.app.services.helpers import (
     convert_embedding_bytes,
@@ -13,13 +20,6 @@
     linkify_missing_citations,
     stringify_docs_content,
 )
-from src.app.bibliography.helpers.helpers import (
-    compute_authors_for_ris,
-    compute_publication_date_for_ris,
-    compute_ris_doctype,
-    ris_line,
-    welearn_document_to_ris,
-)
 from src.app.shared.domain.exceptions import LanguageNotSupportedError