diff --git a/ai-hub/app/core/pipelines/question_decider.py b/ai-hub/app/core/pipelines/question_decider.py index 56b5d4b..9df1c88 100644 --- a/ai-hub/app/core/pipelines/question_decider.py +++ b/ai-hub/app/core/pipelines/question_decider.py @@ -44,6 +44,7 @@ 4. **Final Output Requirements:** * Your output must be a valid JSON object matching the schema. + * **Crucial New Rule:** Do not mention internal system variables or the DSPy signature fields (`retrieved_paths_with_content`, `retrieved_paths_without_content`) in the `reasoning`, `answer`, or `code_diff` fields. The output should be user-friendly and not leak implementation details. * Be helpful, precise, and adhere strictly to these rules. Do not hallucinate file paths or content. --- """ diff --git a/ai-hub/app/core/providers/llm/general.py b/ai-hub/app/core/providers/llm/general.py index 1cc9d31..fdbf1a9 100644 --- a/ai-hub/app/core/providers/llm/general.py +++ b/ai-hub/app/core/providers/llm/general.py @@ -7,7 +7,10 @@ self.api_key = api_key self.system_prompt = system_prompt # Call the parent constructor - super().__init__(model=model_name, max_tokens=10000000, **kwargs) + max_tokens = 1000 + if model_name == "gemini": + max_tokens = 10000000 + super().__init__(model=model_name, max_tokens=max_tokens, **kwargs) def _prepare_messages(self, prompt=None, messages=None): """Helper to prepare the messages list, including the system prompt.""" diff --git a/ai-hub/tests/core/providers/llm/test_llm_general.py b/ai-hub/tests/core/providers/llm/test_llm_general.py index 978a4dd..f331689 100644 --- a/ai-hub/tests/core/providers/llm/test_llm_general.py +++ b/ai-hub/tests/core/providers/llm/test_llm_general.py @@ -50,7 +50,7 @@ messages=[{"role": "user", "content": self.prompt}], api_key=self.api_key, temperature=0.0, # <-- Add this line - max_tokens=1000 # <-- Add this line + max_tokens=10000000 # <-- Add this line ) # Run the async test function diff --git a/ai-hub/tests/core/services/test_workspace.py b/ai-hub/tests/core/services/test_workspace.py deleted file mode 100644 index 2177e11..0000000 --- a/ai-hub/tests/core/services/test_workspace.py +++ /dev/null @@ -1,267 +0,0 @@ -import unittest -import json -from unittest.mock import MagicMock, AsyncMock, patch -from app.core.services.workspace import WorkspaceService -from fastapi import WebSocket -import asyncio -from app.db import models -from dspy.teleprompt import LabeledFewShot -from app.core.pipelines.dspy_rag import DspyRagPipeline -from app.core.pipelines.file_selector import CodeRagFileSelector -from app.core.providers.factory import get_llm_provider - -# Use IsolatedAsyncioTestCase for cleaner async testing -class TestWorkspaceService(unittest.IsolatedAsyncioTestCase): - """ - Unit tests for the WorkspaceService class. - """ - - def setUp(self): - """ - Set up a new WorkspaceService instance for each test. - """ - # Patch the database session to prevent real database connections - self.db_patcher = patch('app.core.services.workspace.SessionLocal') - self.mock_session_local = self.db_patcher.start() - self.mock_db = MagicMock() - self.mock_session_local.return_value = self.mock_db - - # Create a WebSocket mock object - self.mock_websocket = AsyncMock() - self.mock_websocket.scope = {"client": ("127.0.0.1", 8000)} - - self.service = WorkspaceService() - - def tearDown(self): - """ - Clean up patches after each test. - """ - self.db_patcher.stop() - - async def test_generate_request_id(self): - """ - Test that generate_request_id returns a valid UUID string. - """ - request_id = self.service.generate_request_id() - self.assertIsInstance(request_id, str) - self.assertTrue(len(request_id) > 0) - - @patch('app.core.services.workspace.WorkspaceService.generate_request_id', return_value="test-uuid") - async def test_send_command_success(self, mock_generate_id): - """ - Test that send_command sends a correctly formatted message. - """ - data = {"path": "/test/path"} - await self.service.send_command(self.mock_websocket, "list_directory", data) - - # Verify the correct message was sent - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "list_directory") - self.assertEqual(sent_message["request_id"], "test-uuid") - self.assertEqual(sent_message["round"], 1) - self.assertEqual(sent_message["path"], "/test/path") - - # Verify session state was updated - self.assertEqual(self.service.sessions[self.mock_websocket.scope["client"]]["round"], 1) - - async def test_send_command_unknown_command(self): - """ - Test that send_command raises an error for an unknown command. - """ - with self.assertRaises(ValueError) as context: - await self.service.send_command(self.mock_websocket, "unknown_command") - self.assertIn("Unknown command: unknown_command", str(context.exception)) - - @patch('app.core.services.workspace.WorkspaceService.handle_select_folder_response', new_callable=AsyncMock) - async def test_dispatch_message_valid_type(self, mock_handler): - """ - Test that dispatch_message calls the correct handler for a known message type. - """ - message = {"type": "select_folder_response", "path": "/test/path"} - - # Recreate the service after patching so the patched method is used - self.service = WorkspaceService() - - await self.service.dispatch_message(self.mock_websocket, message) - mock_handler.assert_called_once_with(self.mock_websocket, message) - - async def test_dispatch_message_unknown_type(self): - """ - Test that dispatch_message handles an unknown message type gracefully. - """ - message = {"type": "unknown_type"} - await self.service.dispatch_message(self.mock_websocket, message) - - # Verify an error message was sent back to the client - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "error") - self.assertIn("Unknown message type", sent_message["content"]) - - @patch('app.core.services.workspace.WorkspaceService.send_command', new_callable=AsyncMock) - async def test_handle_select_folder_response(self, mock_send_command): - """ - Test that the folder response handler correctly triggers a `list_directory` command. - """ - data = {"path": "/user/project", "request_id": "req-123"} - await self.service.handle_select_folder_response(self.mock_websocket, data) - - mock_send_command.assert_called_once_with(self.mock_websocket, "list_directory", data={"path": "/user/project"}) - - @patch('app.core.services.workspace.ast.literal_eval') - @patch('app.core.services.workspace.CodeRagFileSelector') - @patch('app.core.services.workspace.get_llm_provider') - async def test_handle_list_directory_response_success(self, mock_get_llm, mock_cfs, mock_literal_eval): - """ - Test the list directory handler's success path, ensuring it requests file content. - """ - # Mock LLM and dspy pipeline behavior - mock_llm_provider = MagicMock() - mock_get_llm.return_value = mock_llm_provider - mock_pipeline_instance = AsyncMock() - mock_cfs.return_value = mock_pipeline_instance - mock_pipeline_instance.return_value = '["file1.py", "file2.js"]' # LLM output - mock_literal_eval.return_value = ["file1.py", "file2.js"] # Parsed list - - # Mock send_command to verify it's called - self.service.send_command = AsyncMock() - - data = {"files": ["file1.py", "file2.js", "README.md"], "provider_name": "gemini"} - await self.service.handle_list_directory_response(self.mock_websocket, data) - - # Verify the thinking log message was sent - self.mock_websocket.send_text.assert_called() - log_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(log_message["type"], "thinking_log") - self.assertIn("AI selected files: ['file1.py', 'file2.js']", log_message["content"]) - - # Verify the get_file_content command was sent - self.service.send_command.assert_called_once_with( - self.mock_websocket, "get_file_content", data={"filenames": ["file1.py", "file2.js"]} - ) - - @patch('app.core.services.workspace.ast.literal_eval', side_effect=ValueError("Syntax error")) - @patch('app.core.services.workspace.CodeRagFileSelector') - @patch('app.core.services.workspace.get_llm_provider') - async def test_handle_list_directory_response_eval_error(self, mock_get_llm, mock_cfs, mock_literal_eval): - """ - Test that the handler gracefully handles an error from `ast.literal_eval`. - """ - mock_llm_provider = MagicMock() - mock_get_llm.return_value = mock_llm_provider - mock_pipeline_instance = AsyncMock() - mock_cfs.return_value = mock_pipeline_instance - mock_pipeline_instance.return_value = '{"not a list": "!"}' - - data = {"files": ["file1.py"], "provider_name": "gemini"} - await self.service.handle_list_directory_response(self.mock_websocket, data) - - # Verify that an error log was sent to the client - self.mock_websocket.send_text.assert_called() - error_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(error_message["type"], "thinking_log") - self.assertIn("Warning: AI's file list could not be parsed.", error_message["content"]) - - async def test_handle_files_content_response_success(self): - """ - Test that the file content handler processes files correctly. - """ - data = { - "request_id": "req-123", - "files": [ - {"filename": "file1.py", "content": "print('hello world')"}, - {"filename": "file2.js", "content": "console.log('hi')"} - ] - } - await self.service.handle_files_content_response(self.mock_websocket, data) - - # Verify that thinking logs were sent for each file - self.assertEqual(self.mock_websocket.send_text.call_count, 2) - sent_message = json.loads(self.mock_websocket.send_text.call_args_list[0][0][0]) - self.assertEqual(sent_message["content"], "Analyzing the content of file: file1.py") - - sent_message = json.loads(self.mock_websocket.send_text.call_args_list[1][0][0]) - self.assertEqual(sent_message["content"], "Analyzing the content of file: file2.js") - - async def test_handle_files_content_response_empty_data(self): - """ - Test that the handler does nothing if no files are provided. - """ - data = {"request_id": "req-123", "files": []} - await self.service.handle_files_content_response(self.mock_websocket, data) - self.mock_websocket.send_text.assert_not_called() - - async def test_handle_command_output(self): - """ - Test that the command output handler sends a thinking log. - """ - data = {"request_id": "req-123", "command": "ls -l", "output": "file.txt"} - await self.service.handle_command_output(self.mock_websocket, data) - - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "thinking_log") - self.assertIn("Command 'ls -l' completed. Analyzing output.", sent_message["content"]) - - @patch('app.core.services.workspace.DspyRagPipeline') - @patch('app.core.services.workspace.get_llm_provider') - async def test_handle_chat_message_success(self, mock_get_llm, mock_dspy_rag): - """ - Test that a chat message is handled, saved, and a response is sent. - """ - # Mock database session and query behavior - mock_session = MagicMock() - mock_message_instance = models.Message(session_id="test-session-id", sender="user", content="hello") - mock_session.messages = [mock_message_instance] - self.mock_db.query.return_value.options.return_value.filter.return_value.first.return_value = mock_session - - # Mock LLM and dspy pipeline - mock_llm_provider = MagicMock() - mock_get_llm.return_value = mock_llm_provider - mock_dspy_pipeline = AsyncMock() - mock_dspy_rag.return_value = mock_dspy_pipeline - mock_dspy_pipeline.return_value = "Hello! How can I help you?" - - data = {"content": "Hello", "session_id": "test-session-id", "provider_name": "gemini"} - await self.service.handle_chat_message(self.mock_websocket, data) - - # Verify database calls - self.mock_db.add.assert_called() - self.mock_db.commit.assert_called() - self.mock_db.refresh.assert_called() - - # Verify the response was sent to the client - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "chat_message") - self.assertEqual(sent_message["content"], "Hello! How can I help you?") - - async def test_handle_chat_message_no_session_id(self): - """ - Test that the handler returns an error if session_id is missing. - """ - data = {"content": "Hello"} - await self.service.handle_chat_message(self.mock_websocket, data) - - self.mock_db.query.assert_not_called() - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "error") - self.assertIn("session_id is required", sent_message["content"]) - - async def test_handle_chat_message_session_not_found(self): - """ - Test that the handler returns an error if the session is not found in the DB. - """ - # Mock database query to return None - self.mock_db.query.return_value.options.return_value.filter.return_value.first.return_value = None - - data = {"content": "Hello", "session_id": "non-existent-id"} - await self.service.handle_chat_message(self.mock_websocket, data) - - self.mock_db.add.assert_not_called() - self.mock_websocket.send_text.assert_called_once() - sent_message = json.loads(self.mock_websocket.send_text.call_args[0][0]) - self.assertEqual(sent_message["type"], "error") - self.assertIn("Session with ID non-existent-id not found", sent_message["content"]) diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 63c98a4..6a121ab 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -156,7 +156,7 @@ filesData.push({ filepath, content }); readFiles.push(filepath); } catch (error) { - console.error(`Failed to read file ${filepath}:`, error); + console.warn(`Failed to read file: ${filepath}`, error); ws.current.send(JSON.stringify({ type: "error", content: `Could not read file: ${filepath}`,