diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index a371623..16c0621 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -1,4 +1,5 @@ # tests/api/test_dependencies.py + import pytest import asyncio from unittest.mock import MagicMock, patch @@ -10,10 +11,14 @@ from app.core.services.document import DocumentService from app.core.services.rag import RAGService from app.core.services.tts import TTSService -from app.core.services.stt import STTService # Added this import +from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore from app.core.retrievers.base_retriever import Retriever + +# --- Fixtures --- + @pytest.fixture def mock_session(): """ @@ -22,6 +27,7 @@ mock = MagicMock(spec=Session) yield mock + # --- Tests for get_db dependency --- @patch('app.api.dependencies.SessionLocal') @@ -29,38 +35,37 @@ """ Tests that get_db yields a database session and ensures it's closed correctly. """ - # Arrange: Configure the mock SessionLocal to return our mock_session + # Arrange mock_session_local.return_value = mock_session - # Act: Use the generator in a context manager + # Act db_generator = get_db() db = next(db_generator) - # Assert 1: The correct session object was yielded + # Assert: Correct session yielded assert db == mock_session - # Act 2: Manually close the generator + # Act 2: Close generator with pytest.raises(StopIteration): next(db_generator) - # Assert 2: The session's close method was called + # Assert: Session closed mock_session.close.assert_called_once() + @patch('app.api.dependencies.SessionLocal') def test_get_db_closes_on_exception(mock_session_local, mock_session): """ Tests that get_db still closes the session even if an exception occurs. """ - # Arrange: Configure the mock SessionLocal to return our mock_session mock_session_local.return_value = mock_session - # Act & Assert: Call the generator and raise an exception db_generator = get_db() db = next(db_generator) + with pytest.raises(Exception): db_generator.throw(Exception("Test exception")) - - # Assert: The session's close method was still called after the exception was handled + mock_session.close.assert_called_once() @@ -70,58 +75,63 @@ """ Tests that get_current_user returns the expected user dictionary for a valid token. """ - # Act user = asyncio.run(get_current_user(token="valid_token")) - # Assert assert user == {"email": "user@example.com", "id": 1} + def test_get_current_user_with_no_token(): """ Tests that get_current_user raises an HTTPException for a missing token. """ - # Assert with pytest.raises(HTTPException) as excinfo: asyncio.run(get_current_user(token=None)) - + assert excinfo.value.status_code == 401 - assert "Unauthorized" in excinfo.value.detail + assert "Unauthorized" in str(excinfo.value.detail) + # --- Tests for ServiceContainer class --- def test_service_container_initialization(): """ - Tests that ServiceContainer initializes DocumentService and RAGService + Tests that ServiceContainer initializes DocumentService, RAGService, and other services with the correct dependencies. """ # Arrange: Create mock dependencies mock_vector_store = MagicMock(spec=FaissVectorStore) - # The DocumentService constructor needs a .embedder attribute on the vector_store mock_vector_store.embedder = MagicMock() + mock_retrievers = [MagicMock(spec=Retriever), MagicMock(spec=Retriever)] mock_tts_service = MagicMock(spec=TTSService) - - # NEW: Create a mock for STTService mock_stt_service = MagicMock(spec=STTService) + mock_workspace_service = MagicMock(spec=WorkspaceService) - # Act: Instantiate the ServiceContainer, now with all required arguments + # Act container = ServiceContainer( vector_store=mock_vector_store, retrievers=mock_retrievers, tts_service=mock_tts_service, - stt_service=mock_stt_service # Pass the new mock here + stt_service=mock_stt_service, + workspace_service=mock_workspace_service ) - # Assert: Check if the services were created and configured correctly + # Assert: DocumentService assert isinstance(container.document_service, DocumentService) assert container.document_service.vector_store == mock_vector_store - + + # Assert: RAGService assert isinstance(container.rag_service, RAGService) assert container.rag_service.retrievers == mock_retrievers - - # Assert for the tts_service and stt_service as well + + # Assert: TTSService assert isinstance(container.tts_service, TTSService) assert container.tts_service == mock_tts_service + + # Assert: STTService assert isinstance(container.stt_service, STTService) assert container.stt_service == mock_stt_service + # Assert: WorkspaceService + assert isinstance(container.workspace_service, WorkspaceService) + assert container.workspace_service == mock_workspace_service diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index a371623..16c0621 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -1,4 +1,5 @@ # tests/api/test_dependencies.py + import pytest import asyncio from unittest.mock import MagicMock, patch @@ -10,10 +11,14 @@ from app.core.services.document import DocumentService from app.core.services.rag import RAGService from app.core.services.tts import TTSService -from app.core.services.stt import STTService # Added this import +from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore from app.core.retrievers.base_retriever import Retriever + +# --- Fixtures --- + @pytest.fixture def mock_session(): """ @@ -22,6 +27,7 @@ mock = MagicMock(spec=Session) yield mock + # --- Tests for get_db dependency --- @patch('app.api.dependencies.SessionLocal') @@ -29,38 +35,37 @@ """ Tests that get_db yields a database session and ensures it's closed correctly. """ - # Arrange: Configure the mock SessionLocal to return our mock_session + # Arrange mock_session_local.return_value = mock_session - # Act: Use the generator in a context manager + # Act db_generator = get_db() db = next(db_generator) - # Assert 1: The correct session object was yielded + # Assert: Correct session yielded assert db == mock_session - # Act 2: Manually close the generator + # Act 2: Close generator with pytest.raises(StopIteration): next(db_generator) - # Assert 2: The session's close method was called + # Assert: Session closed mock_session.close.assert_called_once() + @patch('app.api.dependencies.SessionLocal') def test_get_db_closes_on_exception(mock_session_local, mock_session): """ Tests that get_db still closes the session even if an exception occurs. """ - # Arrange: Configure the mock SessionLocal to return our mock_session mock_session_local.return_value = mock_session - # Act & Assert: Call the generator and raise an exception db_generator = get_db() db = next(db_generator) + with pytest.raises(Exception): db_generator.throw(Exception("Test exception")) - - # Assert: The session's close method was still called after the exception was handled + mock_session.close.assert_called_once() @@ -70,58 +75,63 @@ """ Tests that get_current_user returns the expected user dictionary for a valid token. """ - # Act user = asyncio.run(get_current_user(token="valid_token")) - # Assert assert user == {"email": "user@example.com", "id": 1} + def test_get_current_user_with_no_token(): """ Tests that get_current_user raises an HTTPException for a missing token. """ - # Assert with pytest.raises(HTTPException) as excinfo: asyncio.run(get_current_user(token=None)) - + assert excinfo.value.status_code == 401 - assert "Unauthorized" in excinfo.value.detail + assert "Unauthorized" in str(excinfo.value.detail) + # --- Tests for ServiceContainer class --- def test_service_container_initialization(): """ - Tests that ServiceContainer initializes DocumentService and RAGService + Tests that ServiceContainer initializes DocumentService, RAGService, and other services with the correct dependencies. """ # Arrange: Create mock dependencies mock_vector_store = MagicMock(spec=FaissVectorStore) - # The DocumentService constructor needs a .embedder attribute on the vector_store mock_vector_store.embedder = MagicMock() + mock_retrievers = [MagicMock(spec=Retriever), MagicMock(spec=Retriever)] mock_tts_service = MagicMock(spec=TTSService) - - # NEW: Create a mock for STTService mock_stt_service = MagicMock(spec=STTService) + mock_workspace_service = MagicMock(spec=WorkspaceService) - # Act: Instantiate the ServiceContainer, now with all required arguments + # Act container = ServiceContainer( vector_store=mock_vector_store, retrievers=mock_retrievers, tts_service=mock_tts_service, - stt_service=mock_stt_service # Pass the new mock here + stt_service=mock_stt_service, + workspace_service=mock_workspace_service ) - # Assert: Check if the services were created and configured correctly + # Assert: DocumentService assert isinstance(container.document_service, DocumentService) assert container.document_service.vector_store == mock_vector_store - + + # Assert: RAGService assert isinstance(container.rag_service, RAGService) assert container.rag_service.retrievers == mock_retrievers - - # Assert for the tts_service and stt_service as well + + # Assert: TTSService assert isinstance(container.tts_service, TTSService) assert container.tts_service == mock_tts_service + + # Assert: STTService assert isinstance(container.stt_service, STTService) assert container.stt_service == mock_stt_service + # Assert: WorkspaceService + assert isinstance(container.workspace_service, WorkspaceService) + assert container.workspace_service == mock_workspace_service diff --git a/ui/client-app/src/components/CodeFolderAccess.js b/ui/client-app/src/components/CodeFolderAccess.js index 37d4617..e29686d 100644 --- a/ui/client-app/src/components/CodeFolderAccess.js +++ b/ui/client-app/src/components/CodeFolderAccess.js @@ -9,7 +9,7 @@ try { const directoryHandle = await window.showDirectoryPicker(); const folderPath = directoryHandle.name; // Simplified path, you might get a full path in a real app - onSelectFolder(folderPath); + onSelectFolder(directoryHandle); } catch (err) { console.error("Failed to select directory:", err); // You might want to handle user cancellation here diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index a371623..16c0621 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -1,4 +1,5 @@ # tests/api/test_dependencies.py + import pytest import asyncio from unittest.mock import MagicMock, patch @@ -10,10 +11,14 @@ from app.core.services.document import DocumentService from app.core.services.rag import RAGService from app.core.services.tts import TTSService -from app.core.services.stt import STTService # Added this import +from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore from app.core.retrievers.base_retriever import Retriever + +# --- Fixtures --- + @pytest.fixture def mock_session(): """ @@ -22,6 +27,7 @@ mock = MagicMock(spec=Session) yield mock + # --- Tests for get_db dependency --- @patch('app.api.dependencies.SessionLocal') @@ -29,38 +35,37 @@ """ Tests that get_db yields a database session and ensures it's closed correctly. """ - # Arrange: Configure the mock SessionLocal to return our mock_session + # Arrange mock_session_local.return_value = mock_session - # Act: Use the generator in a context manager + # Act db_generator = get_db() db = next(db_generator) - # Assert 1: The correct session object was yielded + # Assert: Correct session yielded assert db == mock_session - # Act 2: Manually close the generator + # Act 2: Close generator with pytest.raises(StopIteration): next(db_generator) - # Assert 2: The session's close method was called + # Assert: Session closed mock_session.close.assert_called_once() + @patch('app.api.dependencies.SessionLocal') def test_get_db_closes_on_exception(mock_session_local, mock_session): """ Tests that get_db still closes the session even if an exception occurs. """ - # Arrange: Configure the mock SessionLocal to return our mock_session mock_session_local.return_value = mock_session - # Act & Assert: Call the generator and raise an exception db_generator = get_db() db = next(db_generator) + with pytest.raises(Exception): db_generator.throw(Exception("Test exception")) - - # Assert: The session's close method was still called after the exception was handled + mock_session.close.assert_called_once() @@ -70,58 +75,63 @@ """ Tests that get_current_user returns the expected user dictionary for a valid token. """ - # Act user = asyncio.run(get_current_user(token="valid_token")) - # Assert assert user == {"email": "user@example.com", "id": 1} + def test_get_current_user_with_no_token(): """ Tests that get_current_user raises an HTTPException for a missing token. """ - # Assert with pytest.raises(HTTPException) as excinfo: asyncio.run(get_current_user(token=None)) - + assert excinfo.value.status_code == 401 - assert "Unauthorized" in excinfo.value.detail + assert "Unauthorized" in str(excinfo.value.detail) + # --- Tests for ServiceContainer class --- def test_service_container_initialization(): """ - Tests that ServiceContainer initializes DocumentService and RAGService + Tests that ServiceContainer initializes DocumentService, RAGService, and other services with the correct dependencies. """ # Arrange: Create mock dependencies mock_vector_store = MagicMock(spec=FaissVectorStore) - # The DocumentService constructor needs a .embedder attribute on the vector_store mock_vector_store.embedder = MagicMock() + mock_retrievers = [MagicMock(spec=Retriever), MagicMock(spec=Retriever)] mock_tts_service = MagicMock(spec=TTSService) - - # NEW: Create a mock for STTService mock_stt_service = MagicMock(spec=STTService) + mock_workspace_service = MagicMock(spec=WorkspaceService) - # Act: Instantiate the ServiceContainer, now with all required arguments + # Act container = ServiceContainer( vector_store=mock_vector_store, retrievers=mock_retrievers, tts_service=mock_tts_service, - stt_service=mock_stt_service # Pass the new mock here + stt_service=mock_stt_service, + workspace_service=mock_workspace_service ) - # Assert: Check if the services were created and configured correctly + # Assert: DocumentService assert isinstance(container.document_service, DocumentService) assert container.document_service.vector_store == mock_vector_store - + + # Assert: RAGService assert isinstance(container.rag_service, RAGService) assert container.rag_service.retrievers == mock_retrievers - - # Assert for the tts_service and stt_service as well + + # Assert: TTSService assert isinstance(container.tts_service, TTSService) assert container.tts_service == mock_tts_service + + # Assert: STTService assert isinstance(container.stt_service, STTService) assert container.stt_service == mock_stt_service + # Assert: WorkspaceService + assert isinstance(container.workspace_service, WorkspaceService) + assert container.workspace_service == mock_workspace_service diff --git a/ui/client-app/src/components/CodeFolderAccess.js b/ui/client-app/src/components/CodeFolderAccess.js index 37d4617..e29686d 100644 --- a/ui/client-app/src/components/CodeFolderAccess.js +++ b/ui/client-app/src/components/CodeFolderAccess.js @@ -9,7 +9,7 @@ try { const directoryHandle = await window.showDirectoryPicker(); const folderPath = directoryHandle.name; // Simplified path, you might get a full path in a real app - onSelectFolder(folderPath); + onSelectFolder(directoryHandle); } catch (err) { console.error("Failed to select directory:", err); // You might want to handle user cancellation here diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 1be2c81..2c293ab 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -1,7 +1,9 @@ // src/hooks/useCodeAssistant.js import { useState, useEffect, useRef, useCallback } from "react"; +import { connectToWebSocket } from "../services/websocket"; const useCodeAssistant = ({ pageContainerRef }) => { + // State variables for the assistant's UI and status const [chatHistory, setChatHistory] = useState([]); const [thinkingProcess, setThinkingProcess] = useState([]); const [selectedFolder, setSelectedFolder] = useState(null); @@ -10,116 +12,249 @@ const [isPaused, setIsPaused] = useState(false); const [errorMessage, setErrorMessage] = useState(""); const [showErrorModal, setShowErrorModal] = useState(false); + const [sessionId, setSessionId] = useState(null); + // Refs for the WebSocket connection and directory handle const ws = useRef(null); + const initialized = useRef(false); + const dirHandleRef = useRef(null); - // Initialize WebSocket connection - useEffect(() => { - // Replace with your WebSocket server URL - const websocketUrl = "wss://your-llm-assistant-server.com"; - ws.current = new WebSocket(websocketUrl); - - ws.current.onopen = () => { - console.log("WebSocket connected"); - setConnectionStatus("connected"); - }; - - ws.current.onmessage = (event) => { - const message = JSON.parse(event.data); - handleIncomingMessage(message); - }; - - ws.current.onclose = () => { - console.log("WebSocket disconnected"); - setConnectionStatus("disconnected"); - setIsProcessing(false); - }; - - ws.current.onerror = (error) => { - console.error("WebSocket error:", error); - setConnectionStatus("error"); - setErrorMessage("WebSocket connection failed. Please check the server."); - setShowErrorModal(true); - }; - - return () => { - ws.current.close(); - }; + // --- WebSocket Message Handlers --- + const handleChatMessage = useCallback((message) => { + setChatHistory((prev) => [...prev, { isUser: false, text: message.content }]); + setIsProcessing(false); }, []); - // Handle incoming messages from the server + const handleThinkingLog = useCallback((message) => { + setThinkingProcess((prev) => [ + ...prev, + { type: message.subtype, message: message.content }, + ]); + }, []); + + const handleError = useCallback((message) => { + setErrorMessage(message.content); + setShowErrorModal(true); + setIsProcessing(false); + }, []); + + const handleStatusUpdate = useCallback((message) => { + setIsProcessing(message.processing); + setIsPaused(message.paused); + setConnectionStatus(message.status); + }, []); + + const handleListDirectoryRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + if (!dirHandle) { + const errorMsg = "No folder selected by user."; + console.warn(errorMsg); + ws.current.send(JSON.stringify({ type: "error", content: errorMsg })); + return; + } + + try { + const files = []; + + // Recursive function to walk through directories + async function walkDirectory(handle, path = "") { + for await (const [name, entry] of handle.entries()) { + const relativePath = path ? `${path}/${name}` : name; + + if (entry.kind === "file") { + files.push(relativePath); // Store full relative path + } else if (entry.kind === "directory") { + await walkDirectory(entry, relativePath); // Recurse into subdirectory + } + } + } + + await walkDirectory(dirHandle); + + ws.current.send( + JSON.stringify({ + type: "list_directory_response", + files, + request_id: message.request_id, + }) + ); + + setThinkingProcess((prev) => [ + ...prev, + { + type: "system", + message: `Sent list of file names from folder "${dirHandle.name}" to server. Total files: ${files.length}`, + }, + ]); + } catch (error) { + console.error("Failed to list directory:", error); + ws.current.send( + JSON.stringify({ + type: "error", + content: "Failed to access folder contents.", + request_id: message.request_id, + }) + ); + } + }, []); + + + const handleReadFileRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + const { filename, request_id } = message; + + if (!dirHandle) { + ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id })); + return; + } + + try { + const fileHandle = await dirHandle.getFileHandle(filename); + const file = await fileHandle.getFile(); + const content = await file.text(); + + ws.current.send(JSON.stringify({ + type: "file_content_response", + filename, + content, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Sent content of file "${filename}" to server.` }, + ]); + } catch (error) { + console.error(`Failed to read file ${filename}:`, error); + ws.current.send(JSON.stringify({ + type: "error", + content: `Could not read file: ${filename}`, + request_id, + })); + } + }, []); + + const handleExecuteCommandRequest = useCallback((message) => { + const { command, request_id } = message; + + const output = `Simulated output for command: '${command}'`; + + ws.current.send(JSON.stringify({ + type: "command_output", + command, + output, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Simulated execution of command: '${command}'` }, + ]); + }, []); + + // Main message handler that routes messages to the correct function const handleIncomingMessage = useCallback((message) => { switch (message.type) { case "chat_message": - setChatHistory((prevHistory) => [...prevHistory, { isUser: false, text: message.content }]); - setIsProcessing(false); + handleChatMessage(message); break; case "thinking_log": - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: message.subtype, message: message.content }, - ]); - break; - case "file_lookup_request": - // In a real app, this would trigger a local file read and send the content back. - // For this hook, we just log the request. - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "system", message: `AI requested file: ${message.filename}` }, - ]); - // To send a file back: - // ws.current.send(JSON.stringify({ type: "file_content", content: "..." })); + handleThinkingLog(message); break; case "error": - setErrorMessage(message.content); - setShowErrorModal(true); - setIsProcessing(false); + handleError(message); break; case "status_update": - setIsProcessing(message.processing); - setIsPaused(message.paused); - setConnectionStatus(message.status); + handleStatusUpdate(message); + break; + case "list_directory": + handleListDirectoryRequest(message); + break; + case "read_file": + handleReadFileRequest(message); + break; + case "execute_command": + handleExecuteCommandRequest(message); break; default: console.log("Unknown message type:", message); } + }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFileRequest, handleExecuteCommandRequest]); + + // --- WebSocket Connection Setup --- + useEffect(() => { + if (initialized.current) return; + initialized.current = true; + + const setupConnection = async () => { + try { + const { ws: newWs, sessionId: newSessionId } = await connectToWebSocket( + handleIncomingMessage, + () => setConnectionStatus("connected"), + () => { + setConnectionStatus("disconnected"); + setIsProcessing(false); + }, + (error) => { + setConnectionStatus("error"); + setErrorMessage(`Failed to connect: ${error.message}`); + setShowErrorModal(true); + } + ); + ws.current = newWs; + setSessionId(newSessionId); + } catch (error) { + console.error("Setup failed:", error); + } + }; + + setupConnection(); + + return () => { + if (ws.current) { + ws.current.close(); + } + }; + }, [handleIncomingMessage]); + + // Send chat message to server + const handleSendChat = useCallback((text) => { + if (ws.current && ws.current.readyState === WebSocket.OPEN) { + setChatHistory((prev) => [...prev, { isUser: true, text }]); + setIsProcessing(true); + ws.current.send(JSON.stringify({ type: "chat_message", content: text })); + } }, []); - // Send a chat message to the server - const handleSendChat = useCallback( - (text) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setChatHistory((prevHistory) => [...prevHistory, { isUser: true, text }]); - setIsProcessing(true); - ws.current.send(JSON.stringify({ type: "chat_message", content: text })); - } - }, - [] - ); + // Open folder picker and store handle + const handleSelectFolder = useCallback(async (directoryHandle) => { + if (!window.showDirectoryPicker) { + // Don't use window.alert, use a custom modal + return; + } - // Send the selected folder path to the server - const handleSelectFolder = useCallback( - (folderPath) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setSelectedFolder(folderPath); - ws.current.send(JSON.stringify({ type: "select_folder", path: folderPath })); - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "user", message: `Selected local folder: ${folderPath}` }, - ]); - } - }, - [] - ); + try { + dirHandleRef.current = directoryHandle; + setSelectedFolder(directoryHandle.name); - // Pause the AI's processing + ws.current.send(JSON.stringify({ type: "select_folder", path: directoryHandle.name })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "user", message: `Selected local folder: ${directoryHandle.name}` }, + ]); + } catch (error) { + console.error("Folder selection canceled or failed:", error); + } + }, []); + + // Control functions const handlePause = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "pause" })); } }, []); - // Stop the AI's processing const handleStop = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "stop" })); @@ -143,4 +278,4 @@ }; }; -export default useCodeAssistant; \ No newline at end of file +export default useCodeAssistant; diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index a371623..16c0621 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -1,4 +1,5 @@ # tests/api/test_dependencies.py + import pytest import asyncio from unittest.mock import MagicMock, patch @@ -10,10 +11,14 @@ from app.core.services.document import DocumentService from app.core.services.rag import RAGService from app.core.services.tts import TTSService -from app.core.services.stt import STTService # Added this import +from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore from app.core.retrievers.base_retriever import Retriever + +# --- Fixtures --- + @pytest.fixture def mock_session(): """ @@ -22,6 +27,7 @@ mock = MagicMock(spec=Session) yield mock + # --- Tests for get_db dependency --- @patch('app.api.dependencies.SessionLocal') @@ -29,38 +35,37 @@ """ Tests that get_db yields a database session and ensures it's closed correctly. """ - # Arrange: Configure the mock SessionLocal to return our mock_session + # Arrange mock_session_local.return_value = mock_session - # Act: Use the generator in a context manager + # Act db_generator = get_db() db = next(db_generator) - # Assert 1: The correct session object was yielded + # Assert: Correct session yielded assert db == mock_session - # Act 2: Manually close the generator + # Act 2: Close generator with pytest.raises(StopIteration): next(db_generator) - # Assert 2: The session's close method was called + # Assert: Session closed mock_session.close.assert_called_once() + @patch('app.api.dependencies.SessionLocal') def test_get_db_closes_on_exception(mock_session_local, mock_session): """ Tests that get_db still closes the session even if an exception occurs. """ - # Arrange: Configure the mock SessionLocal to return our mock_session mock_session_local.return_value = mock_session - # Act & Assert: Call the generator and raise an exception db_generator = get_db() db = next(db_generator) + with pytest.raises(Exception): db_generator.throw(Exception("Test exception")) - - # Assert: The session's close method was still called after the exception was handled + mock_session.close.assert_called_once() @@ -70,58 +75,63 @@ """ Tests that get_current_user returns the expected user dictionary for a valid token. """ - # Act user = asyncio.run(get_current_user(token="valid_token")) - # Assert assert user == {"email": "user@example.com", "id": 1} + def test_get_current_user_with_no_token(): """ Tests that get_current_user raises an HTTPException for a missing token. """ - # Assert with pytest.raises(HTTPException) as excinfo: asyncio.run(get_current_user(token=None)) - + assert excinfo.value.status_code == 401 - assert "Unauthorized" in excinfo.value.detail + assert "Unauthorized" in str(excinfo.value.detail) + # --- Tests for ServiceContainer class --- def test_service_container_initialization(): """ - Tests that ServiceContainer initializes DocumentService and RAGService + Tests that ServiceContainer initializes DocumentService, RAGService, and other services with the correct dependencies. """ # Arrange: Create mock dependencies mock_vector_store = MagicMock(spec=FaissVectorStore) - # The DocumentService constructor needs a .embedder attribute on the vector_store mock_vector_store.embedder = MagicMock() + mock_retrievers = [MagicMock(spec=Retriever), MagicMock(spec=Retriever)] mock_tts_service = MagicMock(spec=TTSService) - - # NEW: Create a mock for STTService mock_stt_service = MagicMock(spec=STTService) + mock_workspace_service = MagicMock(spec=WorkspaceService) - # Act: Instantiate the ServiceContainer, now with all required arguments + # Act container = ServiceContainer( vector_store=mock_vector_store, retrievers=mock_retrievers, tts_service=mock_tts_service, - stt_service=mock_stt_service # Pass the new mock here + stt_service=mock_stt_service, + workspace_service=mock_workspace_service ) - # Assert: Check if the services were created and configured correctly + # Assert: DocumentService assert isinstance(container.document_service, DocumentService) assert container.document_service.vector_store == mock_vector_store - + + # Assert: RAGService assert isinstance(container.rag_service, RAGService) assert container.rag_service.retrievers == mock_retrievers - - # Assert for the tts_service and stt_service as well + + # Assert: TTSService assert isinstance(container.tts_service, TTSService) assert container.tts_service == mock_tts_service + + # Assert: STTService assert isinstance(container.stt_service, STTService) assert container.stt_service == mock_stt_service + # Assert: WorkspaceService + assert isinstance(container.workspace_service, WorkspaceService) + assert container.workspace_service == mock_workspace_service diff --git a/ui/client-app/src/components/CodeFolderAccess.js b/ui/client-app/src/components/CodeFolderAccess.js index 37d4617..e29686d 100644 --- a/ui/client-app/src/components/CodeFolderAccess.js +++ b/ui/client-app/src/components/CodeFolderAccess.js @@ -9,7 +9,7 @@ try { const directoryHandle = await window.showDirectoryPicker(); const folderPath = directoryHandle.name; // Simplified path, you might get a full path in a real app - onSelectFolder(folderPath); + onSelectFolder(directoryHandle); } catch (err) { console.error("Failed to select directory:", err); // You might want to handle user cancellation here diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 1be2c81..2c293ab 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -1,7 +1,9 @@ // src/hooks/useCodeAssistant.js import { useState, useEffect, useRef, useCallback } from "react"; +import { connectToWebSocket } from "../services/websocket"; const useCodeAssistant = ({ pageContainerRef }) => { + // State variables for the assistant's UI and status const [chatHistory, setChatHistory] = useState([]); const [thinkingProcess, setThinkingProcess] = useState([]); const [selectedFolder, setSelectedFolder] = useState(null); @@ -10,116 +12,249 @@ const [isPaused, setIsPaused] = useState(false); const [errorMessage, setErrorMessage] = useState(""); const [showErrorModal, setShowErrorModal] = useState(false); + const [sessionId, setSessionId] = useState(null); + // Refs for the WebSocket connection and directory handle const ws = useRef(null); + const initialized = useRef(false); + const dirHandleRef = useRef(null); - // Initialize WebSocket connection - useEffect(() => { - // Replace with your WebSocket server URL - const websocketUrl = "wss://your-llm-assistant-server.com"; - ws.current = new WebSocket(websocketUrl); - - ws.current.onopen = () => { - console.log("WebSocket connected"); - setConnectionStatus("connected"); - }; - - ws.current.onmessage = (event) => { - const message = JSON.parse(event.data); - handleIncomingMessage(message); - }; - - ws.current.onclose = () => { - console.log("WebSocket disconnected"); - setConnectionStatus("disconnected"); - setIsProcessing(false); - }; - - ws.current.onerror = (error) => { - console.error("WebSocket error:", error); - setConnectionStatus("error"); - setErrorMessage("WebSocket connection failed. Please check the server."); - setShowErrorModal(true); - }; - - return () => { - ws.current.close(); - }; + // --- WebSocket Message Handlers --- + const handleChatMessage = useCallback((message) => { + setChatHistory((prev) => [...prev, { isUser: false, text: message.content }]); + setIsProcessing(false); }, []); - // Handle incoming messages from the server + const handleThinkingLog = useCallback((message) => { + setThinkingProcess((prev) => [ + ...prev, + { type: message.subtype, message: message.content }, + ]); + }, []); + + const handleError = useCallback((message) => { + setErrorMessage(message.content); + setShowErrorModal(true); + setIsProcessing(false); + }, []); + + const handleStatusUpdate = useCallback((message) => { + setIsProcessing(message.processing); + setIsPaused(message.paused); + setConnectionStatus(message.status); + }, []); + + const handleListDirectoryRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + if (!dirHandle) { + const errorMsg = "No folder selected by user."; + console.warn(errorMsg); + ws.current.send(JSON.stringify({ type: "error", content: errorMsg })); + return; + } + + try { + const files = []; + + // Recursive function to walk through directories + async function walkDirectory(handle, path = "") { + for await (const [name, entry] of handle.entries()) { + const relativePath = path ? `${path}/${name}` : name; + + if (entry.kind === "file") { + files.push(relativePath); // Store full relative path + } else if (entry.kind === "directory") { + await walkDirectory(entry, relativePath); // Recurse into subdirectory + } + } + } + + await walkDirectory(dirHandle); + + ws.current.send( + JSON.stringify({ + type: "list_directory_response", + files, + request_id: message.request_id, + }) + ); + + setThinkingProcess((prev) => [ + ...prev, + { + type: "system", + message: `Sent list of file names from folder "${dirHandle.name}" to server. Total files: ${files.length}`, + }, + ]); + } catch (error) { + console.error("Failed to list directory:", error); + ws.current.send( + JSON.stringify({ + type: "error", + content: "Failed to access folder contents.", + request_id: message.request_id, + }) + ); + } + }, []); + + + const handleReadFileRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + const { filename, request_id } = message; + + if (!dirHandle) { + ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id })); + return; + } + + try { + const fileHandle = await dirHandle.getFileHandle(filename); + const file = await fileHandle.getFile(); + const content = await file.text(); + + ws.current.send(JSON.stringify({ + type: "file_content_response", + filename, + content, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Sent content of file "${filename}" to server.` }, + ]); + } catch (error) { + console.error(`Failed to read file ${filename}:`, error); + ws.current.send(JSON.stringify({ + type: "error", + content: `Could not read file: ${filename}`, + request_id, + })); + } + }, []); + + const handleExecuteCommandRequest = useCallback((message) => { + const { command, request_id } = message; + + const output = `Simulated output for command: '${command}'`; + + ws.current.send(JSON.stringify({ + type: "command_output", + command, + output, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Simulated execution of command: '${command}'` }, + ]); + }, []); + + // Main message handler that routes messages to the correct function const handleIncomingMessage = useCallback((message) => { switch (message.type) { case "chat_message": - setChatHistory((prevHistory) => [...prevHistory, { isUser: false, text: message.content }]); - setIsProcessing(false); + handleChatMessage(message); break; case "thinking_log": - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: message.subtype, message: message.content }, - ]); - break; - case "file_lookup_request": - // In a real app, this would trigger a local file read and send the content back. - // For this hook, we just log the request. - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "system", message: `AI requested file: ${message.filename}` }, - ]); - // To send a file back: - // ws.current.send(JSON.stringify({ type: "file_content", content: "..." })); + handleThinkingLog(message); break; case "error": - setErrorMessage(message.content); - setShowErrorModal(true); - setIsProcessing(false); + handleError(message); break; case "status_update": - setIsProcessing(message.processing); - setIsPaused(message.paused); - setConnectionStatus(message.status); + handleStatusUpdate(message); + break; + case "list_directory": + handleListDirectoryRequest(message); + break; + case "read_file": + handleReadFileRequest(message); + break; + case "execute_command": + handleExecuteCommandRequest(message); break; default: console.log("Unknown message type:", message); } + }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFileRequest, handleExecuteCommandRequest]); + + // --- WebSocket Connection Setup --- + useEffect(() => { + if (initialized.current) return; + initialized.current = true; + + const setupConnection = async () => { + try { + const { ws: newWs, sessionId: newSessionId } = await connectToWebSocket( + handleIncomingMessage, + () => setConnectionStatus("connected"), + () => { + setConnectionStatus("disconnected"); + setIsProcessing(false); + }, + (error) => { + setConnectionStatus("error"); + setErrorMessage(`Failed to connect: ${error.message}`); + setShowErrorModal(true); + } + ); + ws.current = newWs; + setSessionId(newSessionId); + } catch (error) { + console.error("Setup failed:", error); + } + }; + + setupConnection(); + + return () => { + if (ws.current) { + ws.current.close(); + } + }; + }, [handleIncomingMessage]); + + // Send chat message to server + const handleSendChat = useCallback((text) => { + if (ws.current && ws.current.readyState === WebSocket.OPEN) { + setChatHistory((prev) => [...prev, { isUser: true, text }]); + setIsProcessing(true); + ws.current.send(JSON.stringify({ type: "chat_message", content: text })); + } }, []); - // Send a chat message to the server - const handleSendChat = useCallback( - (text) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setChatHistory((prevHistory) => [...prevHistory, { isUser: true, text }]); - setIsProcessing(true); - ws.current.send(JSON.stringify({ type: "chat_message", content: text })); - } - }, - [] - ); + // Open folder picker and store handle + const handleSelectFolder = useCallback(async (directoryHandle) => { + if (!window.showDirectoryPicker) { + // Don't use window.alert, use a custom modal + return; + } - // Send the selected folder path to the server - const handleSelectFolder = useCallback( - (folderPath) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setSelectedFolder(folderPath); - ws.current.send(JSON.stringify({ type: "select_folder", path: folderPath })); - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "user", message: `Selected local folder: ${folderPath}` }, - ]); - } - }, - [] - ); + try { + dirHandleRef.current = directoryHandle; + setSelectedFolder(directoryHandle.name); - // Pause the AI's processing + ws.current.send(JSON.stringify({ type: "select_folder", path: directoryHandle.name })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "user", message: `Selected local folder: ${directoryHandle.name}` }, + ]); + } catch (error) { + console.error("Folder selection canceled or failed:", error); + } + }, []); + + // Control functions const handlePause = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "pause" })); } }, []); - // Stop the AI's processing const handleStop = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "stop" })); @@ -143,4 +278,4 @@ }; }; -export default useCodeAssistant; \ No newline at end of file +export default useCodeAssistant; diff --git a/ui/client-app/src/services/websocket.js b/ui/client-app/src/services/websocket.js new file mode 100644 index 0000000..1fe4458 --- /dev/null +++ b/ui/client-app/src/services/websocket.js @@ -0,0 +1,65 @@ +// src/services/websocket.js + +import { createSession } from "./apiService"; + +/** + * Connects to the WebSocket server, establishing a new session first. + * @param {function(Object): void} onMessageCallback - Callback for incoming messages. + * @param {function(): void} onOpenCallback - Callback when the connection is opened. + * @param {function(): void} onCloseCallback - Callback when the connection is closed. + * @param {function(Error): void} onErrorCallback - Callback when an error occurs. + * @returns {Promise<{ws: WebSocket, sessionId: string}>} The WebSocket instance and the session ID. + */ +export const connectToWebSocket = async ( + onMessageCallback, + onOpenCallback, + onCloseCallback, + onErrorCallback +) => { + try { + let sessionId = localStorage.getItem("sessionId"); + + if (!sessionId) { + // No existing session, so create one via API + const session = await createSession(); + sessionId = session.id; + + // Store it in localStorage for reuse + localStorage.setItem("sessionId", sessionId); + } + + // You now have a valid sessionId, either reused or newly created + console.log("Using session ID:", sessionId); + + // Then, use the session ID to open the WebSocket connection + const websocketUrl = `ws://localhost:8001/ws/workspace/${sessionId}`; + const ws = new WebSocket(websocketUrl); + + ws.onopen = () => { + console.log("WebSocket connected"); + if (onOpenCallback) onOpenCallback(); + }; + + ws.onmessage = (event) => { + const message = JSON.parse(event.data); + if (onMessageCallback) onMessageCallback(message); + }; + + ws.onclose = () => { + console.log("WebSocket disconnected"); + if (onCloseCallback) onCloseCallback(); + }; + + ws.onerror = (error) => { + console.error("WebSocket error:", error); + if (onErrorCallback) onErrorCallback(error); + }; + + return { ws, sessionId }; + + } catch (error) { + console.error("Failed to create session or connect to WebSocket:", error); + if (onErrorCallback) onErrorCallback(error); + throw error; + } +}; diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index a1e497f..e410b71 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -8,6 +8,7 @@ from app.core.services.rag import RAGService from app.core.services.tts import TTSService from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore @@ -28,7 +29,7 @@ class ServiceContainer: - def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService): + def __init__(self, vector_store: FaissVectorStore, retrievers: List[Retriever], tts_service: TTSService, stt_service: STTService, workspace_service: WorkspaceService): # Initialize all services within the container self.document_service = DocumentService(vector_store=vector_store) self.rag_service = RAGService( @@ -36,3 +37,4 @@ ) self.tts_service = tts_service self.stt_service = stt_service + self.workspace_service= workspace_service diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index dba14b6..c6628b3 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,9 +7,7 @@ from .documents import create_documents_router from .tts import create_tts_router from .stt import create_stt_router - -# Import the new workspace router -from .workspace import router as workspace_router +from .workspace import create_workspace_router def create_api_router(services: ServiceContainer) -> APIRouter: """ @@ -23,6 +21,6 @@ router.include_router(create_documents_router(services)) router.include_router(create_tts_router(services)) router.include_router(create_stt_router(services)) - router.include_router(workspace_router) - + router.include_router(create_workspace_router(services)) + return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index 87018cf..e85872a 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -1,8 +1,40 @@ -from fastapi import APIRouter, WebSocket +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from app.api.dependencies import ServiceContainer +import json -router = APIRouter() -@router.websocket("/ws/workspace/{session_id}") -async def websocket_endpoint(websocket: WebSocket, session_id: str): - # Your WebSocket logic here - pass \ No newline at end of file +def create_workspace_router(services: ServiceContainer) -> APIRouter: + router = APIRouter() + + @router.websocket("/ws/workspace/{session_id}") + async def websocket_endpoint(websocket: WebSocket, session_id: str): + await websocket.accept() + print(f"WebSocket connection accepted for session: {session_id}") + + # Send a welcome message to confirm the connection is active + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + try: + await websocket.send_text(json.dumps({ + "type": "connection_established", + "message": f"Connected to AI Hub. Session ID: {session_id}" + })) + + while True: + message = await websocket.receive_text() + data = json.loads(message) + + # The endpoint's only job is to dispatch the message to the service + await services.workspace_service.dispatch_message(websocket, data) + + except WebSocketDisconnect: + print(f"WebSocket connection disconnected for session: {session_id}") + except Exception as e: + print(f"An error occurred: {e}") + + finally: + print(f"Closing WebSocket for session: {session_id}") + return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index 4afaa3f..1297a9a 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -16,6 +16,7 @@ from app.api.dependencies import ServiceContainer from app.core.services.tts import TTSService from app.core.services.stt import STTService # NEW: Added the missing import for STTService +from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService # Note: The llm_clients import and initialization are removed as they # are not used in RAGService's constructor based on your services.py # from app.core.llm_clients import DeepSeekClient, GeminiClient @@ -99,13 +100,17 @@ # 7. Initialize the STTService stt_service = STTService(stt_provider=stt_provider) - # 8. Initialize the Service Container with all services + # 8. Initialize the WorkspaceService + workspace_service = WorkspaceService() + + # 9. Initialize the Service Container with all services # This replaces the previous, redundant initialization services = ServiceContainer( vector_store=vector_store, retrievers=retrievers, tts_service=tts_service, - stt_service=stt_service # NEW: Pass the new STT service instance + stt_service=stt_service, + workspace_service=workspace_service ) # Create and include the API router, injecting the service diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py new file mode 100644 index 0000000..907d9a8 --- /dev/null +++ b/ai-hub/app/core/services/workspace.py @@ -0,0 +1,88 @@ +import asyncio +import json +from typing import Dict, Any, Callable, Awaitable +from fastapi import WebSocket +from sqlalchemy.orm import Session + +# A type hint for our handler functions +MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]] + +class WorkspaceService: + """ + Manages the full lifecycle of an AI workspace session, including + handling various message types and dispatching them to the correct handlers. + """ + def __init__(self): + # The dispatcher map: keys are message types, values are handler functions + self.message_handlers: Dict[str, MessageHandler] = { + "select_folder": self.handle_select_folder_response, + "list_directory_response": self.handle_list_directory_response, + "file_content_response": self.handle_file_content_response, + "execute_command_response": self.handle_command_output, + # Add more message types here as needed + } + + async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): + """ + Routes an incoming message to the appropriate handler based on its 'type'. + """ + message_type = message.get("type") + handler = self.message_handlers.get(message_type) + if handler: + await handler(websocket, message) + else: + print(f"Warning: No handler found for message type: {message_type}") + await websocket.send_text(json.dumps({"type": "error", "content": f"Unknown message type: {message_type}"})) + + async def handle_select_folder_response(self, websocket:WebSocket, data: Dict[str, Any]): + """Handles the client's response to a select folder response.""" + path = data.get("path") + request_id = data.get("request_id") + print(f"Received folder selected (request_id: {request_id}): Path: {path}") + # After the server received the request that folder selected, we immediately ask for the file lists in the folder. + await websocket.send_text(json.dumps({ + "type": "list_directory", + "request_id": request_id + })) + + + async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the client's response to a list_directory request.""" + # This is where the AI logic would pick up after getting the file list + files = data.get("files", []) + folders = data.get("folders", []) + request_id = data.get("request_id") + await websocket.send_text(json.dumps({ + "type": "list", + "content": f"Analyzing the content of file: {files[0]}" + })) + print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}") + + + async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the content of a file sent by the client.""" + filename = data.get("filename") + content = data.get("content") + request_id = data.get("request_id") + + print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}") + + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Analyzing the content of file: {filename}" + })) + + async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]): + """Handles the output from a command executed by the client.""" + command = data.get("command") + output = data.get("output") + request_id = data.get("request_id") + + print(f"Received output for command '{command}' (request_id: {request_id}). Output: {output}") + + # The AI would process the command output to determine the next step + await websocket.send_text(json.dumps({ + "type": "thinking_log", + "content": f"Command '{command}' completed. Analyzing output." + })) + diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index a371623..16c0621 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -1,4 +1,5 @@ # tests/api/test_dependencies.py + import pytest import asyncio from unittest.mock import MagicMock, patch @@ -10,10 +11,14 @@ from app.core.services.document import DocumentService from app.core.services.rag import RAGService from app.core.services.tts import TTSService -from app.core.services.stt import STTService # Added this import +from app.core.services.stt import STTService +from app.core.services.workspace import WorkspaceService from app.core.vector_store.faiss_store import FaissVectorStore from app.core.retrievers.base_retriever import Retriever + +# --- Fixtures --- + @pytest.fixture def mock_session(): """ @@ -22,6 +27,7 @@ mock = MagicMock(spec=Session) yield mock + # --- Tests for get_db dependency --- @patch('app.api.dependencies.SessionLocal') @@ -29,38 +35,37 @@ """ Tests that get_db yields a database session and ensures it's closed correctly. """ - # Arrange: Configure the mock SessionLocal to return our mock_session + # Arrange mock_session_local.return_value = mock_session - # Act: Use the generator in a context manager + # Act db_generator = get_db() db = next(db_generator) - # Assert 1: The correct session object was yielded + # Assert: Correct session yielded assert db == mock_session - # Act 2: Manually close the generator + # Act 2: Close generator with pytest.raises(StopIteration): next(db_generator) - # Assert 2: The session's close method was called + # Assert: Session closed mock_session.close.assert_called_once() + @patch('app.api.dependencies.SessionLocal') def test_get_db_closes_on_exception(mock_session_local, mock_session): """ Tests that get_db still closes the session even if an exception occurs. """ - # Arrange: Configure the mock SessionLocal to return our mock_session mock_session_local.return_value = mock_session - # Act & Assert: Call the generator and raise an exception db_generator = get_db() db = next(db_generator) + with pytest.raises(Exception): db_generator.throw(Exception("Test exception")) - - # Assert: The session's close method was still called after the exception was handled + mock_session.close.assert_called_once() @@ -70,58 +75,63 @@ """ Tests that get_current_user returns the expected user dictionary for a valid token. """ - # Act user = asyncio.run(get_current_user(token="valid_token")) - # Assert assert user == {"email": "user@example.com", "id": 1} + def test_get_current_user_with_no_token(): """ Tests that get_current_user raises an HTTPException for a missing token. """ - # Assert with pytest.raises(HTTPException) as excinfo: asyncio.run(get_current_user(token=None)) - + assert excinfo.value.status_code == 401 - assert "Unauthorized" in excinfo.value.detail + assert "Unauthorized" in str(excinfo.value.detail) + # --- Tests for ServiceContainer class --- def test_service_container_initialization(): """ - Tests that ServiceContainer initializes DocumentService and RAGService + Tests that ServiceContainer initializes DocumentService, RAGService, and other services with the correct dependencies. """ # Arrange: Create mock dependencies mock_vector_store = MagicMock(spec=FaissVectorStore) - # The DocumentService constructor needs a .embedder attribute on the vector_store mock_vector_store.embedder = MagicMock() + mock_retrievers = [MagicMock(spec=Retriever), MagicMock(spec=Retriever)] mock_tts_service = MagicMock(spec=TTSService) - - # NEW: Create a mock for STTService mock_stt_service = MagicMock(spec=STTService) + mock_workspace_service = MagicMock(spec=WorkspaceService) - # Act: Instantiate the ServiceContainer, now with all required arguments + # Act container = ServiceContainer( vector_store=mock_vector_store, retrievers=mock_retrievers, tts_service=mock_tts_service, - stt_service=mock_stt_service # Pass the new mock here + stt_service=mock_stt_service, + workspace_service=mock_workspace_service ) - # Assert: Check if the services were created and configured correctly + # Assert: DocumentService assert isinstance(container.document_service, DocumentService) assert container.document_service.vector_store == mock_vector_store - + + # Assert: RAGService assert isinstance(container.rag_service, RAGService) assert container.rag_service.retrievers == mock_retrievers - - # Assert for the tts_service and stt_service as well + + # Assert: TTSService assert isinstance(container.tts_service, TTSService) assert container.tts_service == mock_tts_service + + # Assert: STTService assert isinstance(container.stt_service, STTService) assert container.stt_service == mock_stt_service + # Assert: WorkspaceService + assert isinstance(container.workspace_service, WorkspaceService) + assert container.workspace_service == mock_workspace_service diff --git a/ui/client-app/src/components/CodeFolderAccess.js b/ui/client-app/src/components/CodeFolderAccess.js index 37d4617..e29686d 100644 --- a/ui/client-app/src/components/CodeFolderAccess.js +++ b/ui/client-app/src/components/CodeFolderAccess.js @@ -9,7 +9,7 @@ try { const directoryHandle = await window.showDirectoryPicker(); const folderPath = directoryHandle.name; // Simplified path, you might get a full path in a real app - onSelectFolder(folderPath); + onSelectFolder(directoryHandle); } catch (err) { console.error("Failed to select directory:", err); // You might want to handle user cancellation here diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 1be2c81..2c293ab 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -1,7 +1,9 @@ // src/hooks/useCodeAssistant.js import { useState, useEffect, useRef, useCallback } from "react"; +import { connectToWebSocket } from "../services/websocket"; const useCodeAssistant = ({ pageContainerRef }) => { + // State variables for the assistant's UI and status const [chatHistory, setChatHistory] = useState([]); const [thinkingProcess, setThinkingProcess] = useState([]); const [selectedFolder, setSelectedFolder] = useState(null); @@ -10,116 +12,249 @@ const [isPaused, setIsPaused] = useState(false); const [errorMessage, setErrorMessage] = useState(""); const [showErrorModal, setShowErrorModal] = useState(false); + const [sessionId, setSessionId] = useState(null); + // Refs for the WebSocket connection and directory handle const ws = useRef(null); + const initialized = useRef(false); + const dirHandleRef = useRef(null); - // Initialize WebSocket connection - useEffect(() => { - // Replace with your WebSocket server URL - const websocketUrl = "wss://your-llm-assistant-server.com"; - ws.current = new WebSocket(websocketUrl); - - ws.current.onopen = () => { - console.log("WebSocket connected"); - setConnectionStatus("connected"); - }; - - ws.current.onmessage = (event) => { - const message = JSON.parse(event.data); - handleIncomingMessage(message); - }; - - ws.current.onclose = () => { - console.log("WebSocket disconnected"); - setConnectionStatus("disconnected"); - setIsProcessing(false); - }; - - ws.current.onerror = (error) => { - console.error("WebSocket error:", error); - setConnectionStatus("error"); - setErrorMessage("WebSocket connection failed. Please check the server."); - setShowErrorModal(true); - }; - - return () => { - ws.current.close(); - }; + // --- WebSocket Message Handlers --- + const handleChatMessage = useCallback((message) => { + setChatHistory((prev) => [...prev, { isUser: false, text: message.content }]); + setIsProcessing(false); }, []); - // Handle incoming messages from the server + const handleThinkingLog = useCallback((message) => { + setThinkingProcess((prev) => [ + ...prev, + { type: message.subtype, message: message.content }, + ]); + }, []); + + const handleError = useCallback((message) => { + setErrorMessage(message.content); + setShowErrorModal(true); + setIsProcessing(false); + }, []); + + const handleStatusUpdate = useCallback((message) => { + setIsProcessing(message.processing); + setIsPaused(message.paused); + setConnectionStatus(message.status); + }, []); + + const handleListDirectoryRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + if (!dirHandle) { + const errorMsg = "No folder selected by user."; + console.warn(errorMsg); + ws.current.send(JSON.stringify({ type: "error", content: errorMsg })); + return; + } + + try { + const files = []; + + // Recursive function to walk through directories + async function walkDirectory(handle, path = "") { + for await (const [name, entry] of handle.entries()) { + const relativePath = path ? `${path}/${name}` : name; + + if (entry.kind === "file") { + files.push(relativePath); // Store full relative path + } else if (entry.kind === "directory") { + await walkDirectory(entry, relativePath); // Recurse into subdirectory + } + } + } + + await walkDirectory(dirHandle); + + ws.current.send( + JSON.stringify({ + type: "list_directory_response", + files, + request_id: message.request_id, + }) + ); + + setThinkingProcess((prev) => [ + ...prev, + { + type: "system", + message: `Sent list of file names from folder "${dirHandle.name}" to server. Total files: ${files.length}`, + }, + ]); + } catch (error) { + console.error("Failed to list directory:", error); + ws.current.send( + JSON.stringify({ + type: "error", + content: "Failed to access folder contents.", + request_id: message.request_id, + }) + ); + } + }, []); + + + const handleReadFileRequest = useCallback(async (message) => { + const dirHandle = dirHandleRef.current; + const { filename, request_id } = message; + + if (!dirHandle) { + ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id })); + return; + } + + try { + const fileHandle = await dirHandle.getFileHandle(filename); + const file = await fileHandle.getFile(); + const content = await file.text(); + + ws.current.send(JSON.stringify({ + type: "file_content_response", + filename, + content, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Sent content of file "${filename}" to server.` }, + ]); + } catch (error) { + console.error(`Failed to read file ${filename}:`, error); + ws.current.send(JSON.stringify({ + type: "error", + content: `Could not read file: ${filename}`, + request_id, + })); + } + }, []); + + const handleExecuteCommandRequest = useCallback((message) => { + const { command, request_id } = message; + + const output = `Simulated output for command: '${command}'`; + + ws.current.send(JSON.stringify({ + type: "command_output", + command, + output, + request_id, + })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "system", message: `Simulated execution of command: '${command}'` }, + ]); + }, []); + + // Main message handler that routes messages to the correct function const handleIncomingMessage = useCallback((message) => { switch (message.type) { case "chat_message": - setChatHistory((prevHistory) => [...prevHistory, { isUser: false, text: message.content }]); - setIsProcessing(false); + handleChatMessage(message); break; case "thinking_log": - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: message.subtype, message: message.content }, - ]); - break; - case "file_lookup_request": - // In a real app, this would trigger a local file read and send the content back. - // For this hook, we just log the request. - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "system", message: `AI requested file: ${message.filename}` }, - ]); - // To send a file back: - // ws.current.send(JSON.stringify({ type: "file_content", content: "..." })); + handleThinkingLog(message); break; case "error": - setErrorMessage(message.content); - setShowErrorModal(true); - setIsProcessing(false); + handleError(message); break; case "status_update": - setIsProcessing(message.processing); - setIsPaused(message.paused); - setConnectionStatus(message.status); + handleStatusUpdate(message); + break; + case "list_directory": + handleListDirectoryRequest(message); + break; + case "read_file": + handleReadFileRequest(message); + break; + case "execute_command": + handleExecuteCommandRequest(message); break; default: console.log("Unknown message type:", message); } + }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFileRequest, handleExecuteCommandRequest]); + + // --- WebSocket Connection Setup --- + useEffect(() => { + if (initialized.current) return; + initialized.current = true; + + const setupConnection = async () => { + try { + const { ws: newWs, sessionId: newSessionId } = await connectToWebSocket( + handleIncomingMessage, + () => setConnectionStatus("connected"), + () => { + setConnectionStatus("disconnected"); + setIsProcessing(false); + }, + (error) => { + setConnectionStatus("error"); + setErrorMessage(`Failed to connect: ${error.message}`); + setShowErrorModal(true); + } + ); + ws.current = newWs; + setSessionId(newSessionId); + } catch (error) { + console.error("Setup failed:", error); + } + }; + + setupConnection(); + + return () => { + if (ws.current) { + ws.current.close(); + } + }; + }, [handleIncomingMessage]); + + // Send chat message to server + const handleSendChat = useCallback((text) => { + if (ws.current && ws.current.readyState === WebSocket.OPEN) { + setChatHistory((prev) => [...prev, { isUser: true, text }]); + setIsProcessing(true); + ws.current.send(JSON.stringify({ type: "chat_message", content: text })); + } }, []); - // Send a chat message to the server - const handleSendChat = useCallback( - (text) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setChatHistory((prevHistory) => [...prevHistory, { isUser: true, text }]); - setIsProcessing(true); - ws.current.send(JSON.stringify({ type: "chat_message", content: text })); - } - }, - [] - ); + // Open folder picker and store handle + const handleSelectFolder = useCallback(async (directoryHandle) => { + if (!window.showDirectoryPicker) { + // Don't use window.alert, use a custom modal + return; + } - // Send the selected folder path to the server - const handleSelectFolder = useCallback( - (folderPath) => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - setSelectedFolder(folderPath); - ws.current.send(JSON.stringify({ type: "select_folder", path: folderPath })); - setThinkingProcess((prevProcess) => [ - ...prevProcess, - { type: "user", message: `Selected local folder: ${folderPath}` }, - ]); - } - }, - [] - ); + try { + dirHandleRef.current = directoryHandle; + setSelectedFolder(directoryHandle.name); - // Pause the AI's processing + ws.current.send(JSON.stringify({ type: "select_folder", path: directoryHandle.name })); + + setThinkingProcess((prev) => [ + ...prev, + { type: "user", message: `Selected local folder: ${directoryHandle.name}` }, + ]); + } catch (error) { + console.error("Folder selection canceled or failed:", error); + } + }, []); + + // Control functions const handlePause = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "pause" })); } }, []); - // Stop the AI's processing const handleStop = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { ws.current.send(JSON.stringify({ type: "control", command: "stop" })); @@ -143,4 +278,4 @@ }; }; -export default useCodeAssistant; \ No newline at end of file +export default useCodeAssistant; diff --git a/ui/client-app/src/services/websocket.js b/ui/client-app/src/services/websocket.js new file mode 100644 index 0000000..1fe4458 --- /dev/null +++ b/ui/client-app/src/services/websocket.js @@ -0,0 +1,65 @@ +// src/services/websocket.js + +import { createSession } from "./apiService"; + +/** + * Connects to the WebSocket server, establishing a new session first. + * @param {function(Object): void} onMessageCallback - Callback for incoming messages. + * @param {function(): void} onOpenCallback - Callback when the connection is opened. + * @param {function(): void} onCloseCallback - Callback when the connection is closed. + * @param {function(Error): void} onErrorCallback - Callback when an error occurs. + * @returns {Promise<{ws: WebSocket, sessionId: string}>} The WebSocket instance and the session ID. + */ +export const connectToWebSocket = async ( + onMessageCallback, + onOpenCallback, + onCloseCallback, + onErrorCallback +) => { + try { + let sessionId = localStorage.getItem("sessionId"); + + if (!sessionId) { + // No existing session, so create one via API + const session = await createSession(); + sessionId = session.id; + + // Store it in localStorage for reuse + localStorage.setItem("sessionId", sessionId); + } + + // You now have a valid sessionId, either reused or newly created + console.log("Using session ID:", sessionId); + + // Then, use the session ID to open the WebSocket connection + const websocketUrl = `ws://localhost:8001/ws/workspace/${sessionId}`; + const ws = new WebSocket(websocketUrl); + + ws.onopen = () => { + console.log("WebSocket connected"); + if (onOpenCallback) onOpenCallback(); + }; + + ws.onmessage = (event) => { + const message = JSON.parse(event.data); + if (onMessageCallback) onMessageCallback(message); + }; + + ws.onclose = () => { + console.log("WebSocket disconnected"); + if (onCloseCallback) onCloseCallback(); + }; + + ws.onerror = (error) => { + console.error("WebSocket error:", error); + if (onErrorCallback) onErrorCallback(error); + }; + + return { ws, sessionId }; + + } catch (error) { + console.error("Failed to create session or connect to WebSocket:", error); + if (onErrorCallback) onErrorCallback(error); + throw error; + } +}; diff --git a/ui/run_web.sh b/ui/run_web.sh index d4e2a17..896415e 100644 --- a/ui/run_web.sh +++ b/ui/run_web.sh @@ -18,7 +18,7 @@ # Resolve script directory SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" AI_HUB_DIR="$(realpath "$SCRIPT_DIR/../ai-hub")" -TTS_CLIENT_DIR="$SCRIPT_DIR/tts-client-app" +TTS_CLIENT_DIR="$SCRIPT_DIR/client-app" AI_HUB_HOST="0.0.0.0" AI_HUB_PORT="8001"