diff --git a/pytest.ini b/pytest.ini index 94e738e..6e6b0f9 100644 --- a/pytest.ini +++ b/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 diff --git a/src/app/bibliography/api/bibliography.py b/src/app/bibliography/api/bibliography.py index 87c4285..1fe919a 100644 --- a/src/app/bibliography/api/bibliography.py +++ b/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() diff --git a/src/app/bibliography/helpers/helpers.py b/src/app/bibliography/helpers/helpers.py index a064ee6..0246030 100644 --- a/src/app/bibliography/helpers/helpers.py +++ b/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 diff --git a/src/app/shared/infra/abst_chat.py b/src/app/shared/infra/abst_chat.py index c356ea4..e0e3e74 100644 --- a/src/app/shared/infra/abst_chat.py +++ b/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, ) diff --git a/src/app/tests/services/test_abst_chat.py b/src/app/tests/services/test_abst_chat.py index dc6f857..7f38bc5 100644 --- a/src/app/tests/services/test_abst_chat.py +++ b/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" diff --git a/src/app/tests/services/test_helpers.py b/src/app/tests/services/test_helpers.py index 0af167c..ce414f0 100644 --- a/src/app/tests/services/test_helpers.py +++ b/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