From 5006271abb25119a81a2d6003c5ef0b0ab668916 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 11 Mar 2026 11:43:27 +0000 Subject: [PATCH] fix stream stuff (#2293) --- .../api/answer/services/stream_processor.py | 30 ++++- .../answer/services/test_stream_processor.py | 118 +++++++++++++++++- 2 files changed, 145 insertions(+), 3 deletions(-) diff --git a/application/api/answer/services/stream_processor.py b/application/api/answer/services/stream_processor.py index 588bb642..0cf47094 100644 --- a/application/api/answer/services/stream_processor.py +++ b/application/api/answer/services/stream_processor.py @@ -350,9 +350,35 @@ class StreamProcessor: self.source = {} self.all_sources = [] + def _resolve_agent_id(self) -> Optional[str]: + """Resolve agent_id from request, then fall back to conversation context.""" + request_agent_id = self.data.get("agent_id") + if request_agent_id: + return str(request_agent_id) + + if not self.conversation_id or not self.initial_user_id: + return None + + try: + conversation = self.conversation_service.get_conversation( + self.conversation_id, self.initial_user_id + ) + except Exception: + return None + + if not conversation: + return None + + conversation_agent_id = conversation.get("agent_id") + if conversation_agent_id: + return str(conversation_agent_id) + + return None + def _configure_agent(self): """Configure the agent based on request data""" - agent_id = self.data.get("agent_id") + agent_id = self._resolve_agent_id() + self.agent_key, self.is_shared_usage, self.shared_token = self._get_agent_key( agent_id, self.initial_user_id ) @@ -471,7 +497,7 @@ class StreamProcessor: def pre_fetch_docs(self, question: str) -> tuple[Optional[str], Optional[list]]: """Pre-fetch documents for template rendering before agent creation""" - if self.data.get("isNoneDoc", False): + if self.data.get("isNoneDoc", False) and not self.agent_id: logger.info("Pre-fetch skipped: isNoneDoc=True") return None, None try: diff --git a/tests/api/answer/services/test_stream_processor.py b/tests/api/answer/services/test_stream_processor.py index 45cbbb67..18f282b7 100644 --- a/tests/api/answer/services/test_stream_processor.py +++ b/tests/api/answer/services/test_stream_processor.py @@ -1,5 +1,5 @@ import pytest -from bson import ObjectId +from bson import DBRef, ObjectId @pytest.mark.unit @@ -214,6 +214,122 @@ class TestStreamProcessorAgentConfiguration: assert isinstance(processor.agent_config, dict) assert processor.agent_id is None + def test_conversation_agent_overrides_request_active_docs(self, mock_mongo_db): + from application.api.answer.services.stream_processor import StreamProcessor + from application.core.settings import settings + + db = mock_mongo_db[settings.MONGO_DB_NAME] + agents_collection = db["agents"] + sources_collection = db["sources"] + conversations_collection = db["conversations"] + + agent_id = ObjectId() + conversation_id = ObjectId() + agent_source_id = ObjectId() + request_source_id = ObjectId() + + sources_collection.insert_many( + [ + {"_id": agent_source_id, "name": "Agent source", "retriever": "classic"}, + { + "_id": request_source_id, + "name": "Request source", + "retriever": "hybrid", + }, + ] + ) + + agents_collection.insert_one( + { + "_id": agent_id, + "key": "agent_key_2", + "user": "user_123", + "prompt_id": "default", + "agent_type": "classic", + "source": DBRef("sources", agent_source_id), + } + ) + + conversations_collection.insert_one( + { + "_id": conversation_id, + "user": "user_123", + "agent_id": str(agent_id), + "queries": [], + } + ) + + processor = StreamProcessor( + { + "question": "Test", + "conversation_id": str(conversation_id), + "active_docs": str(request_source_id), + }, + {"sub": "user_123"}, + ) + + processor._configure_agent() + processor._configure_source() + + assert processor.agent_id == str(agent_id) + assert processor.source["active_docs"] == str(agent_source_id) + + +@pytest.mark.unit +class TestStreamProcessorDocPrefetch: + + def test_prefetch_not_skipped_for_agent_when_isNoneDoc_true(self, mock_mongo_db): + from unittest.mock import MagicMock + + from application.api.answer.services.stream_processor import StreamProcessor + from application.core.settings import settings + + db = mock_mongo_db[settings.MONGO_DB_NAME] + agents_collection = db["agents"] + sources_collection = db["sources"] + + agent_id = ObjectId() + source_id = ObjectId() + + sources_collection.insert_one( + {"_id": source_id, "name": "Agent source", "retriever": "classic"} + ) + agents_collection.insert_one( + { + "_id": agent_id, + "key": "agent_prefetch_key", + "user": "user_123", + "prompt_id": "default", + "agent_type": "classic", + "source": DBRef("sources", source_id), + } + ) + + processor = StreamProcessor( + { + "question": "Summarize context", + "agent_id": str(agent_id), + "isNoneDoc": True, + }, + {"sub": "user_123"}, + ) + processor.initialize() + + mock_retriever = MagicMock() + mock_retriever.chunks = 2 + mock_retriever.doc_token_limit = 50000 + mock_retriever.search.return_value = [ + {"text": "Agent doc content", "source": "agent.pdf"} + ] + processor.create_retriever = MagicMock(return_value=mock_retriever) + + docs_together, docs = processor.pre_fetch_docs("Summarize context") + + processor.create_retriever.assert_called_once() + assert docs is not None + assert docs_together is not None + assert "Agent doc content" in docs_together + @pytest.mark.unit class TestStreamProcessorAttachments: