From a1cf87909dfb1aecb2c526b8758d534131ae2d91 Mon Sep 17 00:00:00 2001 From: "Billy D." Date: Wed, 18 Feb 2026 07:37:13 -0500 Subject: [PATCH] fixing ruff suggestions and tests needed updating. --- chat_handler.py | 19 +++++++++++-------- tests/conftest.py | 18 +++++++++++++++--- tests/test_chat_handler.py | 20 ++++++++++++++------ 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/chat_handler.py b/chat_handler.py index 20a75e6..7116a29 100644 --- a/chat_handler.py +++ b/chat_handler.py @@ -134,9 +134,7 @@ class ChatHandler(Handler): enable_tts = data.get("enable_tts", self.chat_settings.enable_tts) system_prompt = data.get("system_prompt") # companions-frontend may set a custom response subject - response_subject = data.get( - "response_subject", f"ai.chat.response.{request_id}" - ) + response_subject = data.get("response_subject", f"ai.chat.response.{request_id}") logger.info(f"Processing request {request_id}: {query[:50]}...") @@ -159,7 +157,9 @@ class ChatHandler(Handler): # 2. Search Milvus for context documents = await self._search_context( - embedding, collection, top_k=top_k, + embedding, + collection, + top_k=top_k, ) # 3. Optionally rerank documents @@ -172,21 +172,24 @@ class ChatHandler(Handler): if reranked: context = self._build_context(reranked) rag_sources = [ - d.get("source", d.get("document", "")[:80]) - for d in reranked[:3] + d.get("source", d.get("document", "")[:80]) for d in reranked[:3] ] used_rag = True # 5. Generate LLM response (with or without RAG context) response_text = await self._generate_response( - query, context or None, system_prompt, + query, + context or None, + system_prompt, ) # 6. Stream response chunks if requested if enable_streaming: stream_subject = f"ai.chat.response.stream.{request_id}" await self._publish_streaming_chunks( - stream_subject, request_id, response_text, + stream_subject, + request_id, + response_text, ) # 7. Optionally synthesize speech diff --git a/tests/conftest.py b/tests/conftest.py index 7a29ae7..b29ffc7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -53,7 +53,7 @@ def sample_reranked(): def mock_nats_message(): """Create a mock NATS message.""" msg = MagicMock() - msg.subject = "ai.chat.request" + msg.subject = "ai.chat.user.test-user-1.message" msg.reply = "ai.chat.response.test-123" return msg @@ -63,7 +63,13 @@ def mock_chat_request(): """Sample chat request payload.""" return { "request_id": "test-request-123", - "query": "What is machine learning?", + "user_id": "test-user-1", + "username": "testuser", + "message": "What is machine learning?", + "premium": True, + "enable_rag": True, + "enable_reranker": True, + "enable_streaming": False, "collection": "test_collection", "enable_tts": False, "system_prompt": None, @@ -75,7 +81,13 @@ def mock_chat_request_with_tts(): """Sample chat request with TTS enabled.""" return { "request_id": "test-request-456", - "query": "Tell me about AI", + "user_id": "test-user-2", + "username": "testuser2", + "message": "Tell me about AI", + "premium": True, + "enable_rag": True, + "enable_reranker": True, + "enable_streaming": False, "collection": "documents", "enable_tts": True, "system_prompt": "You are a helpful assistant.", diff --git a/tests/test_chat_handler.py b/tests/test_chat_handler.py index d857283..3e56b22 100644 --- a/tests/test_chat_handler.py +++ b/tests/test_chat_handler.py @@ -86,7 +86,7 @@ class TestChatHandler: def test_init(self, handler): """Test handler initialization.""" - assert handler.subject == "ai.chat.request" + assert handler.subject == "ai.chat.user.*.message" assert handler.queue_group == "chat-handlers" assert handler.chat_settings.service_name == "chat-handler" @@ -111,12 +111,15 @@ class TestChatHandler: result = await handler.handle_message(mock_nats_message, mock_chat_request) # Verify - assert result["request_id"] == "test-request-123" + assert result["user_id"] == "test-user-1" + assert result["success"] is True assert "response" in result assert result["response"] == "Machine learning is a subset of AI that..." - assert "sources" in result # include_sources is True by default + assert result["response_text"] == result["response"] + assert result["used_rag"] is True + assert isinstance(result["rag_sources"], list) - # Verify pipeline was called + # Verify RAG pipeline was called (enable_rag=True in fixture) handler.embeddings.embed_single.assert_called_once() handler.milvus.search_with_texts.assert_called_once() handler.reranker.rerank.assert_called_once() @@ -142,7 +145,9 @@ class TestChatHandler: result = await handler.handle_message(mock_nats_message, mock_chat_request) - assert "sources" not in result + # New response format doesn't have a separate "sources" key; + # rag_sources is always present (may be empty) + assert "rag_sources" in result @pytest.mark.asyncio async def test_handle_message_with_tts( @@ -180,7 +185,10 @@ class TestChatHandler: """Test LLM is called with custom system prompt.""" request = { "request_id": "test-123", - "query": "Hello", + "user_id": "user-42", + "message": "Hello", + "premium": True, + "enable_rag": True, "system_prompt": "You are a pirate. Respond like one.", }