diff --git a/ai-hub/app/api/dependencies.py b/ai-hub/app/api/dependencies.py index 8a0a3bd..24efc59 100644 --- a/ai-hub/app/api/dependencies.py +++ b/ai-hub/app/api/dependencies.py @@ -54,11 +54,11 @@ self.document_service = DocumentService(vector_store=vector_store) return self - def with_rag_service(self, retrievers: List[Retriever]): + def with_rag_service(self, retrievers: List[Retriever], prompt_service = None): """ Adds a RAGService instance to the container. """ - self.rag_service = RAGService(retrievers=retrievers) + self.rag_service = RAGService(retrievers=retrievers, prompt_service=prompt_service) return self def __getattr__(self, name: str) -> Any: diff --git a/ai-hub/app/api/routes/api.py b/ai-hub/app/api/routes/api.py index 5515e75..161df72 100644 --- a/ai-hub/app/api/routes/api.py +++ b/ai-hub/app/api/routes/api.py @@ -7,7 +7,6 @@ from .tts import create_tts_router from .general import create_general_router from .stt import create_stt_router -from .workspace import create_workspace_router from .user import create_users_router def create_api_router(services: ServiceContainer) -> APIRouter: @@ -23,7 +22,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(create_workspace_router(services)) router.include_router(create_users_router(services)) return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/workspace.md b/ai-hub/app/api/routes/workspace.md deleted file mode 100644 index b2b6dd7..0000000 --- a/ai-hub/app/api/routes/workspace.md +++ /dev/null @@ -1,74 +0,0 @@ -# AI Hub: Interactive Coding Assistant - -The goal of this project is to extend the existing AI Hub with an interactive coding assistant. This assistant will enable an AI to directly "work" on a user's local codebase, allowing it to understand, debug, and modify a project's files by communicating with a lightweight client running on the user's machine. - -This is a significant step beyond a traditional RESTful API, as it allows for real-time, bi-directional communication, which is essential for the dynamic, back-and-forth nature of a coding session. - -### Core Components - -The project is built on two main components that communicate over a **WebSocket (WSS)** connection: - -* **AI Hub (Server):** The central brain of the operation. This is your existing FastAPI application, which will be extended to include WebSocket endpoints. The server's primary responsibilities are: - * Maintaining the WSS connection with the client. - * Using a **Large Language Model (LLM)** to reason about a user's coding problem. - * Sending structured commands to the client (e.g., "read this file," "execute this command"). - * Receiving and processing command output from the client. - * Generating file changes or new commands based on the current project state. - -* **Local Client:** A new, lightweight application that runs on the user's machine. Its sole purpose is to act as a secure, local intermediary. Its responsibilities are: - * Establishing and maintaining the WSS connection to the AI Hub. - * **Security:** Enforcing a strict sandbox, ensuring the AI can only interact with the user-selected project directory. - * Reading file contents and sending them to the server. - * Executing shell commands and sending the output back to the server. - * Receiving and applying file modifications from the server. - ---- - -### Communication Protocol - -A clear, robust communication schema is the foundation of this project. All communication over the WebSocket will be a JSON object with a consistent structure, using a `type` to identify the command and a `request_id` to link a request to its response. - -#### Server-to-Client Commands - -These are the commands the AI Hub sends to the local client. - -| `type` | Description | Payload | -| :---------------- | :--------------------------------------------------------------------------- | :---------------------------------------------------- | -| `list_files` | Requests a list of files and folders in a specified directory. | `{"path": "./"}` | -| `read_files` | Requests the content of one or more specific files. | `{"files": ["src/main.py", "tests/test.py"]}` | -| `execute_command` | Instructs the client to execute a shell command in the project root directory. | `{"command": "npm test"}` | -| `write_file` | Instructs the client to overwrite a specific file with new content. | `{"path": "src/main.py", "content": "new code..."}` | - -#### Client-to-Server Responses - -These are the responses the local client sends back to the AI Hub. All responses must include the original `request_id`. - -| `type` | Description | Payload | -| :-------------------- | :------------------------------------------------------------------- | :----------------------------------------------------------------------- | -| `list_files_response` | Returns a list of files and folders. | `{"status": "success", "files": [{"name": "...", "is_dir": true}]}` | -| `read_files_response` | Returns the content of the requested files. | `{"status": "success", "files": {"path": "content"}}` | -| `command_output` | Returns the `stdout` and `stderr` from the executed command. | `{"status": "success", "stdout": "...", "stderr": "..."}` | -| `write_file_response` | Confirms if a file was successfully written or returns an error. | `{"status": "success"}` | - -In the event of an error (e.g., a file not found), the client will always return a `status: "error"` in the payload, along with a human-readable `error_message`. - ---- - -### Kickoff Plan - -We can tackle this project in three phases to ensure a solid foundation. - -#### Phase 1: Establish the Connection -* **Server:** Create the new `workspace.py` route file and define a basic `websocket` endpoint. -* **Client:** Build a minimal proof-of-concept client (e.g., a Python or Node.js script) that can connect to the server's WebSocket endpoint. -* **Goal:** Successfully send a "ping" from the client and receive a "pong" from the server. - -#### Phase 2: Implement Core Functionality -* **Server:** Define the logic for handling incoming client messages (`command_output`, etc.) and for sending commands to the client. -* **Client:** Implement the message handlers for `list_files`, `read_files`, and `execute_command`. This client must include robust file system checks to prevent access outside the designated directory. -* **Goal:** The server can successfully send a `list_files` command to the client, receive a list of files, and then send a `read_files` command for one of those files. - -#### Phase 3: Integrate and Secure -* **Server:** Integrate the new WebSocket communication layer with your existing LLM and RAG pipelines. The LLM's output should be formatted as one of the structured commands. -* **Client:** Implement security features like command validation and a user confirmation UI for potentially destructive commands (e.g., `rm`). -* **Goal:** A user can initiate a conversation with the AI, the AI can ask for a file list and file content, and the client can provide it, all without manual intervention. diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py deleted file mode 100644 index 4583c68..0000000 --- a/ai-hub/app/api/routes/workspace.py +++ /dev/null @@ -1,40 +0,0 @@ -from fastapi import APIRouter, WebSocket, WebSocketDisconnect -from app.api.dependencies import ServiceContainer -import json - - -def create_workspace_router(services: ServiceContainer) -> APIRouter: - router = APIRouter() - - @router.websocket("/ws/workspace/{client_id}") - async def websocket_endpoint(websocket: WebSocket, client_id: str): - await websocket.accept() - print(f"WebSocket connection accepted for client: {client_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. Client ID: {client_id}" - })) - - try: - await websocket.send_text(json.dumps({ - "type": "connection_established", - "message": f"Connected to AI Hub. Client ID: {client_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 client: {client_id}") - except Exception as e: - print(f"An error occurred: {e}") - - finally: - print(f"Closing WebSocket for client: {client_id}") - return router \ No newline at end of file diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index f86a5ef..99c775b 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -19,11 +19,11 @@ from app.utils import print_config from app.api.dependencies import ServiceContainer, get_db from app.core.services.session import SessionService -from app.core.services import SessionService 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.user import UserService -from app.core.services.workspace import WorkspaceService # NEW: Added the missing import for STTService +from app.core.services.prompt import PromptService +from app.core.services.tool import ToolService # 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 @@ -114,17 +114,19 @@ except Exception as e: logger.warning(f"Failed to initialize TTS/STT: {e}") + prompt_service = PromptService() # 9. Initialize the Service Container with all initialized services services = ServiceContainer() - services.with_rag_service(retrievers=retrievers) + services.with_rag_service(retrievers=retrievers, prompt_service=prompt_service) services.with_document_service(vector_store=vector_store) services.with_service("stt_service", service=STTService(stt_provider=stt_provider)) services.with_service("tts_service", service=TTSService(tts_provider=tts_provider)) - services.with_service("workspace_service", service=WorkspaceService()) + services.with_service("prompt_service", service=prompt_service) services.with_service("session_service", service=SessionService()) services.with_service("user_service", service=UserService()) + services.with_service("tool_service", service=ToolService()) # Create and include the API router, injecting the service api_router = create_api_router(services=services) diff --git a/ai-hub/app/core/pipelines/code_changer.py b/ai-hub/app/core/pipelines/code_changer.py index 163c9a0..bdebed0 100644 --- a/ai-hub/app/core/pipelines/code_changer.py +++ b/ai-hub/app/core/pipelines/code_changer.py @@ -52,7 +52,11 @@ filepath: str, original_files: List[Dict[str, Any]], updated_files: List[Dict[str, Any]], - llm_provider = None + llm_provider = None, + prompt_service = None, + db: Optional[Session] = None, + user_id: Optional[str] = None, + prompt_slug: str = "code-changer" ) -> Tuple[str, str]: if not llm_provider: raise ValueError("LLM Provider is required.") @@ -60,7 +64,13 @@ original_json = json.dumps(original_files) updated_json = json.dumps(updated_files) - prompt = PROMPT_TEMPLATE.format( + template = PROMPT_TEMPLATE + if prompt_service and db and user_id: + db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id) + if db_prompt: + template = db_prompt.content + + prompt = template.format( overall_plan=overall_plan, instruction=instruction, filepath=filepath, diff --git a/ai-hub/app/core/pipelines/code_reviewer.py b/ai-hub/app/core/pipelines/code_reviewer.py index ef609c7..095ac0e 100644 --- a/ai-hub/app/core/pipelines/code_reviewer.py +++ b/ai-hub/app/core/pipelines/code_reviewer.py @@ -1,5 +1,6 @@ import json from typing import List, Dict, Any, Tuple, Optional, Callable +from sqlalchemy.orm import Session PROMPT_TEMPLATE = """ ### 🧠 Core Directives @@ -36,7 +37,11 @@ execution_plan: str, final_code_changes: List[Dict[str, Any]], original_files: List[Dict[str, Any]], - llm_provider = None + llm_provider = None, + prompt_service = None, + db: Optional[Session] = None, + user_id: Optional[str] = None, + prompt_slug: str = "code-reviewer" ) -> Tuple[str, str, str]: if not llm_provider: raise ValueError("LLM Provider is required.") @@ -44,7 +49,13 @@ final_code_changes_json = json.dumps(final_code_changes) original_files_json = json.dumps(original_files) - prompt = PROMPT_TEMPLATE.format( + template = PROMPT_TEMPLATE + if prompt_service and db and user_id: + db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id) + if db_prompt: + template = db_prompt.content + + prompt = template.format( original_question=original_question, execution_plan=execution_plan, final_code_changes=final_code_changes_json, diff --git a/ai-hub/app/core/pipelines/file_selector.py b/ai-hub/app/core/pipelines/file_selector.py index 9c689cf..e03ca5d 100644 --- a/ai-hub/app/core/pipelines/file_selector.py +++ b/ai-hub/app/core/pipelines/file_selector.py @@ -1,6 +1,7 @@ import json from app.db import models -from typing import List, Dict, Any, Tuple +from typing import List, Dict, Any, Tuple, Optional +from sqlalchemy.orm import Session PROMPT_TEMPLATE = """ You're an **expert file navigator** for a large codebase. Your task is to select the most critical and relevant file paths to answer a user's question. All file paths you select must exist within the provided `retrieved_files` list. @@ -40,14 +41,30 @@ for msg in history ) - async def forward(self, question: str, retrieved_data: List[str], history: List[models.Message], llm_provider = None) -> Tuple[List[str], str]: + async def forward( + self, + question: str, + retrieved_data: List[str], + history: List[models.Message], + llm_provider = None, + prompt_service = None, + db: Optional[Session] = None, + user_id: Optional[str] = None, + prompt_slug: str = "file-selector" + ) -> Tuple[List[str], str]: if not llm_provider: raise ValueError("LLM Provider is required.") retrieved_json = json.dumps(retrieved_data) history_text = self._default_history_formatter(history) - prompt = PROMPT_TEMPLATE.format( + template = PROMPT_TEMPLATE + if prompt_service and db and user_id: + db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id) + if db_prompt: + template = db_prompt.content + + prompt = template.format( question=question, chat_history=history_text, retrieved_files=retrieved_json diff --git a/ai-hub/app/core/pipelines/question_decider.py b/ai-hub/app/core/pipelines/question_decider.py index 70f35ad..4b0ffbe 100644 --- a/ai-hub/app/core/pipelines/question_decider.py +++ b/ai-hub/app/core/pipelines/question_decider.py @@ -3,6 +3,7 @@ from app.core.pipelines.validator import Validator, TokenLimitExceededError from app.db import models from typing import List, Dict, Any, Tuple, Optional, Callable +from sqlalchemy.orm import Session PROMPT_TEMPLATE = """ ### 🧠 **Core Directives** @@ -71,7 +72,11 @@ question: str, history: List[models.Message], retrieved_data: Dict[str, Any], - llm_provider = None + llm_provider = None, + prompt_service = None, + db: Optional[Session] = None, + user_id: Optional[str] = None, + prompt_slug: str = "question-decider" ) -> Tuple[str, str, str, Optional[List[Dict]]]: if not llm_provider: raise ValueError("LLM Provider is required.") @@ -91,7 +96,14 @@ history_text = self.history_formatter(history) - prompt = PROMPT_TEMPLATE.format( + # Decide which prompt template to use + template = PROMPT_TEMPLATE + if prompt_service and db and user_id: + db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id) + if db_prompt: + template = db_prompt.content + + prompt = template.format( question=question, chat_history=history_text, retrieved_paths_with_content=json.dumps(with_content, indent=2), diff --git a/ai-hub/app/core/pipelines/rag_pipeline.py b/ai-hub/app/core/pipelines/rag_pipeline.py index f362356..fcd93ea 100644 --- a/ai-hub/app/core/pipelines/rag_pipeline.py +++ b/ai-hub/app/core/pipelines/rag_pipeline.py @@ -1,9 +1,23 @@ import logging -from typing import List, Callable, Optional +from typing import List, Dict, Any, Optional, Callable from sqlalchemy.orm import Session from app.db import models +# Define a default prompt template outside the class or as a class constant +# This is inferred from the usage in the provided diff. +PROMPT_TEMPLATE = """Generate a natural and context-aware answer to the user's question using the provided knowledge and conversation history. + +Relevant excerpts from the knowledge base: +{context} + +Conversation History: +{chat_history} + +User Question: {question} + +Answer:""" + class RagPipeline: """ A flexible and extensible RAG pipeline updated to remove DSPy dependency. @@ -19,18 +33,37 @@ self.history_formatter = history_formatter or self._default_history_formatter self.response_postprocessor = response_postprocessor - async def forward(self, question: str, history: List[models.Message], context_chunks: List[str], llm_provider=None) -> str: + async def forward( + self, + question: str, + context_chunks: List[Dict[str, Any]], + history: List[models.Message], + llm_provider = None, + prompt_service = None, + db: Optional[Session] = None, + user_id: Optional[str] = None, + prompt_slug: str = "rag-pipeline" + ) -> str: logging.debug(f"[RagPipeline.forward] Received question: '{question}'") - context_text = self.context_postprocessor(context_chunks) - history_text = self.history_formatter(history) - - # Step 3: Generate response using manual prompt - prompt = self._build_prompt(context_text, history_text, question) - if not llm_provider: - raise ValueError("LLM Provider is required for RAG pipeline.") + raise ValueError("LLM Provider is required.") + history_text = self.history_formatter(history) + context_text = self.context_postprocessor(context_chunks) + + template = PROMPT_TEMPLATE + if prompt_service and db and user_id: + db_prompt = prompt_service.get_prompt_by_slug(db, prompt_slug, user_id) + if db_prompt: + template = db_prompt.content + + prompt = template.format( + question=question, + context=context_text, + chat_history=history_text + ) + prediction = await llm_provider.acompletion(prompt=prompt) raw_response = prediction.choices[0].message.content diff --git a/ai-hub/app/core/services/__init__.py b/ai-hub/app/core/services/__init__.py index e6075ea..e95bf8d 100644 --- a/ai-hub/app/core/services/__init__.py +++ b/ai-hub/app/core/services/__init__.py @@ -2,7 +2,6 @@ from .session import SessionService from .stt import STTService from .tts import TTSService -from .workspace import WorkspaceService from .user import UserService from .rag import RAGService from .document import DocumentService \ No newline at end of file diff --git a/ai-hub/app/core/services/prompt.py b/ai-hub/app/core/services/prompt.py new file mode 100644 index 0000000..d0bcf7a --- /dev/null +++ b/ai-hub/app/core/services/prompt.py @@ -0,0 +1,70 @@ +from typing import Dict, Any, Optional +from sqlalchemy.orm import Session +from app.db import models +import logging + +logger = logging.getLogger(__name__) + +class PromptService: + """ + Manages centralized system prompts, including fetching, rendering, + and permission validation. + """ + + def get_prompt_by_slug(self, db: Session, slug: str, user_id: str) -> Optional[models.PromptTemplate]: + """ + Fetches a prompt template by its unique slug, verifying if the user has access. + """ + prompt = db.query(models.PromptTemplate).filter(models.PromptTemplate.slug == slug).first() + if not prompt: + logger.warning(f"Prompt with slug '{slug}' not found.") + return None + + if self._can_access_prompt(db, prompt, user_id): + return prompt + + logger.warning(f"User '{user_id}' denied access to prompt '{slug}'.") + return None + + def _can_access_prompt(self, db: Session, prompt: models.PromptTemplate, user_id: str) -> bool: + """ + Internal logic to check if a user can access a specific prompt template. + """ + # 1. Is it public? + if prompt.is_public: + return True + + # 2. Is the user the owner? + if prompt.owner_id == user_id: + return True + + # 3. Does the user belong to the allowed group? + user = db.query(models.User).filter(models.User.id == user_id).first() + if user and prompt.group_id and user.group_id == prompt.group_id: + return True + + # 4. Check explicit granular permissions + permission = db.query(models.AssetPermission).filter( + models.AssetPermission.resource_type == 'prompt', + models.AssetPermission.resource_id == prompt.id, + (models.AssetPermission.user_id == user_id) | + ((models.AssetPermission.group_id == user_id) if user and user.group_id else False) + ).first() + + if permission: + return True + + return False + + def render_prompt(self, prompt: models.PromptTemplate, variables: Dict[str, Any]) -> str: + """ + Renders a prompt template string using the provided variables. + """ + try: + return prompt.content.format(**variables) + except KeyError as e: + logger.error(f"Missing variable for prompt rendering: {e}") + return prompt.content # Fallback to raw if rendering fails partially + except Exception as e: + logger.error(f"Error rendering prompt: {e}") + return prompt.content diff --git a/ai-hub/app/core/services/rag.py b/ai-hub/app/core/services/rag.py index 75d9fec..6d11545 100644 --- a/ai-hub/app/core/services/rag.py +++ b/ai-hub/app/core/services/rag.py @@ -1,4 +1,4 @@ -from typing import List, Tuple +from typing import List, Tuple, Optional from sqlalchemy.orm import Session, joinedload from app.db import models @@ -12,8 +12,9 @@ Service for orchestrating conversational RAG pipelines. Manages chat interactions and message history for a session. """ - def __init__(self, retrievers: List[Retriever]): + def __init__(self, retrievers: List[Retriever], prompt_service = None): self.retrievers = retrievers + self.prompt_service = prompt_service self.faiss_retriever = next((r for r in retrievers if isinstance(r, FaissDBRetriever)), None) async def chat_with_rag( @@ -23,7 +24,8 @@ prompt: str, provider_name: str, load_faiss_retriever: bool = False, - user_service = None + user_service = None, + user_id: Optional[str] = None ) -> Tuple[str, str, int]: """ Processes a user prompt within a session, saves the chat history, and returns a response. @@ -92,7 +94,11 @@ question=prompt, history=session.messages, context_chunks = context_chunks, - llm_provider = llm_provider + llm_provider = llm_provider, + prompt_service = self.prompt_service, + db = db, + user_id = user_id or session.user_id, + prompt_slug = "rag-pipeline" ) # Save assistant's response diff --git a/ai-hub/app/core/services/tool.py b/ai-hub/app/core/services/tool.py new file mode 100644 index 0000000..e05f104 --- /dev/null +++ b/ai-hub/app/core/services/tool.py @@ -0,0 +1,59 @@ +from typing import List, Dict, Any, Optional +from sqlalchemy.orm import Session +from app.db import models +from app.core.skills.base import BaseSkill +import logging + +logger = logging.getLogger(__name__) + +class ToolService: + """ + Orchestrates AI tools (Skills) available to users. + Handles discovery, permission checks, and execution routing. + """ + + def __init__(self, local_skills: List[BaseSkill] = []): + self._local_skills = {s.name: s for s in local_skills} + + def get_available_tools(self, db: Session, user_id: str) -> List[Dict[str, Any]]: + """ + Retrieves all tools the user is authorized to use. + """ + # 1. Start with system/local skills + tools = [s.to_tool_definition() for s in self._local_skills.values()] + + # 2. Add DB-defined skills with permission checks + db_skills = db.query(models.Skill).filter( + (models.Skill.is_system == True) | + (models.Skill.owner_id == user_id) + ).all() + + # TODO: Implement more complex group-based permission logic + for ds in db_skills: + # Prevent duplicates if name overlaps with local + if any(t["function"]["name"] == ds.name for t in tools): + continue + + tools.append({ + "type": "function", + "function": { + "name": ds.name, + "description": ds.description, + "parameters": ds.config.get("parameters", {}) + } + }) + + return tools + + async def call_tool(self, tool_name: str, arguments: Dict[str, Any], **context) -> Any: + """ + Executes a registered skill. + """ + if tool_name in self._local_skills: + skill = self._local_skills[tool_name] + result = await skill.execute(**arguments) + return result.dict() + + # TODO: Handle remote/gRPC skills or MCP skills here + logger.error(f"Tool '{tool_name}' not found or handled yet.") + return {"success": False, "error": "Tool not found"} diff --git a/ai-hub/app/core/services/utils/code_change.py b/ai-hub/app/core/services/utils/code_change.py deleted file mode 100644 index bf3acc0..0000000 --- a/ai-hub/app/core/services/utils/code_change.py +++ /dev/null @@ -1,308 +0,0 @@ -import logging -import json -import re -import uuid -import asyncio -from sqlalchemy.orm import Session -from app.db import file_retriever_models -from typing import Dict, List, Any, Optional, Tuple -from app.core.providers.factory import get_llm_provider -from app.core.pipelines.code_changer import CodeRagCodeChanger -from app.core.pipelines.code_reviewer import CodeReviewer -from fastapi import WebSocket -logger = logging.getLogger(__name__) - -class CodeChangeHelper: - """ - A helper class to process and manage a sequence of code change instructions. - """ - - def __init__(self, db: Session, provider_name: str, original_question: str,input_data: str, reasoning: str,request_id: uuid.UUID): - """ - Initializes the CodeChangeHelper, parsing the input and setting up dependencies. - - Args: - db (Session): The database session. - provider_name (str): The name of the LLM provider. - input_data (str): A JSON string representing a list of code change steps. - """ - self.db = db - self.input_data = input_data - self.reasoning = reasoning - self.llm_provider = get_llm_provider(provider_name) - self.original_question = original_question - self.code_changer = CodeRagCodeChanger() - self.code_reviewer = CodeReviewer() - - self.current_execution_plans: List[Dict[str, Any]] = [] - self.history_plans: List[Dict[str, Any]] = [] - self.updated_files: Dict[str, Dict[str, str]] = {} - self.original_files: Dict[str, str] = {} - self.last_step_index: int = -1 - self.max_round = 3 - try: - self._parse_input_data() - self._preload_original_files(request_id=request_id) - except (json.JSONDecodeError, ValueError) as e: - logger.error(f"Initialization failed due to invalid input: {e}") - raise - - def _parse_input_data(self) -> None: - """ - Parses the input JSON string and validates its structure. - """ - cleaned_input = re.sub(r"^```json\s*|\s*```$", "", self.input_data.strip(), flags=re.DOTALL) - - current_execution_plans = json.loads(cleaned_input) - if not isinstance(current_execution_plans, list): - raise ValueError("Input is not a JSON array.") - - required_keys = ["file_path", "change_instruction", "original_files", "updated_files", "action"] - for item in current_execution_plans: - if not all(key in item for key in required_keys): - raise ValueError(f"An item is missing required keys. Found: {list(item.keys())}, Required: {required_keys}") - - self.current_execution_plans = current_execution_plans - self.history_plans.append(current_execution_plans) - - def _preload_original_files(self, request_id: uuid.UUID) -> None: - """ - Fetches and caches the content of all required original files. - """ - unique_file_paths = set() - for item in self.current_execution_plans: - file_paths = item.get("original_files", []) - for path in file_paths: - unique_file_paths.add(path) - - for file_path in unique_file_paths: - content = self._fetch_file_content(file_path, request_id=request_id) - if content is not None: - self.original_files[file_path] = content - - def _fetch_file_content(self, file_path: str, request_id: uuid.UUID) -> Optional[str]: - """ - Fetches the content of a file from the database. - - Returns None if the file is not found or an error occurs. - """ - try: - # Assuming a single request_id for simplicity; adjust if needed - retrieved_file = self.db.query(file_retriever_models.RetrievedFile).filter_by(file_path=file_path, request_id=request_id).first() - if retrieved_file: - return retrieved_file.content - else: - logger.warning(f"File not found in the database: {file_path}") - return None - except Exception as e: - logger.error(f"Error fetching file content for '{file_path}': {e}") - return None - - async def _process_ai_question(self, item: Dict[str, Any]) -> Dict[str,str]: - """ - Processes a single code change instruction using the AI. - """ - # Prepare the list of original files from the preloaded cache - original_files_list = [ - {"file_path": path, "content": self.original_files.get(path, "")} - for path in item.get("original_files", []) - ] - - # Prepare the list of updated files from the current state - updated_files_list = [ - {"file_path": path, "content": data["content"], "reasoning": data["reasoning"]} - for path, data in self.updated_files.items() - ] - - instruction = item.get("change_instruction") - file_path_to_change = item.get("file_path") - overall_plan = self.input_data - - content, reasoning = await self.code_changer.forward( - overall_plan=overall_plan, - instruction=instruction, - filepath=file_path_to_change, - original_files=original_files_list, - updated_files=updated_files_list, - llm_provider=self.llm_provider - ) - return {"content": content, "reasoning": reasoning} - - async def _handle_thinking_log(self, websocket: WebSocket, reasoning :str): - client_log :Dict[str, Any] = { - "type": "thinking_log", - "content": reasoning - } - await websocket.send_text(json.dumps(client_log)) - - async def _handle_intermediate_chat_message(self, websocket: WebSocket, title: str = "**AI-Generated Execution Plan:**"): - steps_content = [ - ] - - # # Add each change instruction as a numbered list item. - # for i, data in enumerate(self.current_execution_plans): - # # Use f-string to create numbered list items with proper indentation. - # steps_content.append(f"{i+1}. {data['change_instruction']}") - - - client_log: Dict[str, Any] = { - "type": "code_change", - "content": title, - "steps": self.current_execution_plans, - "reasoning": self.reasoning, - "done": False, - } - await websocket.send_text(json.dumps(client_log)) - - async def _post_process(self, websocket: WebSocket)->None: - result = {} - - # Regex to find and extract content from a Markdown code block - # The `re.DOTALL` flag allows `.` to match newlines - code_block_pattern = re.compile(r'```[a-zA-Z]*\n(.*)```', re.DOTALL) - - for file_path, detail in self.updated_files.items(): - original_content = self.original_files.get(file_path, "") - updated_content = detail.get("content", "") - - # Check if the content is wrapped in a code block - match = code_block_pattern.search(updated_content) - if match: - # If a match is found, use the captured group as the new content, - # and strip any leading/trailing whitespace - cleaned_content = match.group(1).strip() - else: - # If no code block is found, use the content as-is - cleaned_content = updated_content - - result[file_path] = { - "old": original_content, - "new": cleaned_content, - "reasoning": detail.get("reasoning", "") - } - - - # Send the final processed changes to the client - payload = json.dumps({ - "type": "code_change", - "code_changes": result, - "content": "Completed all requested code changes.", - "done": True - }) - logger.info(f"Sending code change response to client: {payload}") - await websocket.send_text(payload) - - async def _review_changes(self, final_code_changes: List[Dict[str, Any]]) -> Tuple[str, str, str]: - decision, reasoning, answer = await self.code_reviewer.forward( - original_question=self.original_question, - execution_plan= json.dumps(self.history_plans), - final_code_changes=final_code_changes, - original_files=[{"file_path": k, "content": v} for k, v in self.original_files.items()], - llm_provider=self.llm_provider - ) - return decision, reasoning, answer - - async def _inline_code_replacement(self, input_text: str) -> str: - """ - Replaces placeholders with content from the self.original_files dictionary. - """ - # Define the regex pattern to match the specified format - pattern = re.compile(r'#\[unchanged_section\]\|(.*?)\|(\d+)\|(\d+)') - - # Split the input text into lines - lines = input_text.splitlines() - updated_lines = [] - - for line in lines: - match = pattern.search(line) - if match: - file_path, start_line_str, end_line_str = match.groups() - start_line = int(start_line_str) - end_line = int(end_line_str) - - # Ensure start and end lines are in the correct order - if start_line > end_line: - start_line, end_line = end_line, start_line - - file_content_from_dict = self.original_files.get(file_path) - - if file_content_from_dict is not None: - content_lines = file_content_from_dict.splitlines(keepends=True) - - # Adjust indices for 0-based list - start_idx = start_line - 1 - end_idx = end_line - - if 0 <= start_idx < len(content_lines) and 0 <= end_idx <= len(content_lines): - section = content_lines[start_idx:end_idx] - updated_lines.append("".join(section)) - else: - updated_lines.append(f"# Error: Lines {start_line}-{end_line} are out of bounds for {file_path}\n") - else: - updated_lines.append(f"# Error: File path {file_path} not found in original_files dictionary\n") - else: - updated_lines.append(line) - - return "\n".join(updated_lines) - - async def process(self, websocket: WebSocket, round :int = 0 ,title:str = "**AI-Generated Execution Plan:**") -> Dict[str, Dict[str, str]]: - """ - Executes all code change instructions in sequence. - - Returns: - A dictionary of all updated files with their content and reasoning. - """ - await self._handle_intermediate_chat_message(websocket, title) - for item in self.current_execution_plans: - action = item.get("action") - filepath = item.get("file_path") - reasoning = "" - if action == "delete": - reasoning = f"File deleted: {filepath}" - self.updated_files[filepath] = {"content": "", "reasoning": reasoning} - elif action == "move": - change_instruction = item.get("change_instruction","") - # Use regex to find all strings wrapped in backticks - matches = re.findall(r'`(.*?)`', change_instruction) - if matches: - # Get the last matched string, which is the target path - targetpath = matches[-1] - reasoning =f"File moved from {filepath} to {targetpath}" - self.updated_files[targetpath]= {"content": self.updated_files[filepath].get("content", ""), "reasoning": reasoning} - self.updated_files[filepath]= {"content": "", "reasoning":reasoning} - else: - # extract the target path. - response = await self._process_ai_question(item) - response["content"] = await self._inline_code_replacement(response.pop("content")) - self.updated_files[filepath] = response - reasoning = response.get("reasoning", "") - - if reasoning: - await self._handle_thinking_log(websocket, reasoning) - - if round file_retriever_models.FileRetrievalRequest: - """ - Retrieves an existing FileRetrievalRequest or creates a new one if it doesn't exist. - """ - file_request = self.db.query(file_retriever_models.FileRetrievalRequest).filter_by( - session_id=session_id, directory_path=path - ).first() - - if not file_request: - file_request = file_retriever_models.FileRetrievalRequest( - session_id=session_id, - question=prompt, - directory_path=path - ) - self.db.add(file_request) - self.db.commit() - self.db.refresh(file_request) - else: - # If file_request is found, update it with the latest prompt - file_request.question = prompt - self.db.commit() - self.db.refresh(file_request) - return file_request - - async def _get_file_request_by_id(self, request_id: uuid.UUID) -> file_retriever_models.FileRetrievalRequest: - """ - Retrieves a FileRetrievalRequest by its ID. - """ - return self.db.query(file_retriever_models.FileRetrievalRequest).filter_by(id=request_id).first() - - async def _store_retrieved_files(self, request_id: uuid.UUID, files: List[Dict[str, Any]]): - """ - Synchronizes the database's retrieved files with the client's file list. - - This function compares existing files against new ones and performs - updates, additions, or deletions as necessary. - """ - # 1. Get existing files from the database for this request - existing_files = self.db.query(file_retriever_models.RetrievedFile).filter_by(request_id=request_id).all() - existing_files_map = {file.file_path: file for file in existing_files} - - # Keep track of which existing files are also in the incoming list - incoming_file_paths = set() - - # 2. Iterate through incoming files to handle updates and additions - for file_info in files: - file_path = file_info.get("path") - last_modified_timestamp_ms = file_info.get("lastModified") - - if not file_path or last_modified_timestamp_ms is None: - logger.warning("Skipping file with missing path or timestamp.") - continue - - last_modified_datetime = datetime.fromtimestamp(last_modified_timestamp_ms / 1000.0) - incoming_file_paths.add(file_path) - - # Check if the file already exists in the database - if file_path in existing_files_map: - db_file = existing_files_map[file_path] - # Compare the last modified timestamps - if last_modified_datetime > db_file.last_updated: - # Case: File has been updated, so override the existing record. - logger.info(f"Updating file {file_path}. New timestamp: {last_modified_datetime}") - db_file.last_updated = last_modified_datetime - # The content remains empty for now, as it will be fetched later. - else: - # Case: File is identical or older, do nothing. - # logger.debug(f"File {file_path} is identical or older, skipping.") - pass - - else: - # Case: This is a newly introduced file. - logger.info(f"Adding new file: {file_path}") - new_file = file_retriever_models.RetrievedFile( - request_id=request_id, - file_path=file_path, - file_name=file_info.get("name", ""), - content="", # Content is deliberately left empty. - type="original", - last_updated=last_modified_datetime, - ) - self.db.add(new_file) - - # 3. Purge non-existing files - # Find files in the database that were not in the incoming list - files_to_purge = [ - file for file in existing_files if file.file_path not in incoming_file_paths - ] - if files_to_purge: - logger.info(f"Purging {len(files_to_purge)} non-existing files.") - for file in files_to_purge: - self.db.delete(file) - - # 4. Commit all changes (updates, additions, and deletions) in a single transaction - self.db.commit() - logger.info("File synchronization complete.") - - # def generate_request_id(self) -> str: - # """Generates a unique request ID.""" - # return str(uuid.uuid4()) - - async def _retrieve_by_request_id(self, db: Session, request_id: str) -> Optional[Dict[str, Any]]: - """ - Retrieves a FileRetrievalRequest and all its associated files from the database, - returning the data in a well-formatted JSON-like dictionary. - - Args: - db: The SQLAlchemy database session. - request_id: The UUID of the FileRetrievalRequest. - - Returns: - A dictionary containing the request and file data, or None if the request is not found. - """ - try: - # Convert string request_id to UUID object for the query - request_uuid = uuid.UUID(request_id) - except ValueError: - print(f"Invalid UUID format for request_id: {request_id}") - return None - - # Fetch the request and its related files in a single query using join - request = db.query(file_retriever_models.FileRetrievalRequest).filter( - file_retriever_models.FileRetrievalRequest.id == request_uuid - ).options( - # Eagerly load the retrieved_files to avoid N+1 query problem - joinedload(file_retriever_models.FileRetrievalRequest.retrieved_files) - ).first() - - if not request: - return None - - # Build the dictionary to represent the JSON structure - retrieved_data = { - "request_id": str(request.id), - "question": request.question, - "directory_path": request.directory_path, - "session_id": request.session_id, - "created_at": request.created_at.isoformat() if request.created_at else None, - "retrieved_files": [] - } - - for file in request.retrieved_files: - if file.content: - # For files with content, show the full detailed structure - file_data = { - "file_path": file.file_path, - "content": file.content, - "id": str(file.id), - "name": file.file_name, - "type": file.type, - "last_updated": file.last_updated.isoformat() if file.last_updated else None, - "created_at": file.created_at.isoformat() if file.created_at else None, - } - else: - # For empty files, use a compact representation - file_data = { - "file_path": file.file_path, - "type": file.type - } - retrieved_data["retrieved_files"].append(file_data) - - return retrieved_data - - - - async def get_file_content_by_request_id_and_path(self, db: Session, request_id: uuid.UUID, file_path: str) ->str: - """ - Retrieves a FileRetrievalRequest by its ID. - """ - retrievedFile = db.query(file_retriever_models.RetrievedFile).filter_by(request_id = request_id , file_path=file_path).first() - if retrievedFile and retrievedFile.content: - return retrievedFile.content - else: - logger.warning(f"File with path {file_path} not found for request ID {request_id} or has no content.") - return "" - - # async def _handle_code_change_response(self, db: Session, request_id: str, code_diff: str) -> List[Dict[str, Any]]: - # """ - # Parses the diff, retrieves original file content, and returns a structured, - # per-file dictionary for the client. - # """ - # # Normalize the diff string to ensure consistent splitting, handling cases where - # # the separator may be missing a leading newline. - # normalized_diff = re.sub(r'(? Optional[List[str]]: - """ - Retrieves all files associated with a FileRetrievalRequest from the database, - returning the data in a list of JSON-like dictionaries. - - Args: - db: The SQLAlchemy database session. - request_id: The UUID of the FileRetrievalRequest. - - Returns: - A list of dictionaries containing file data, or None if the request is not found. - """ - try: - request_uuid = uuid.UUID(request_id) - except ValueError: - print(f"Invalid UUID format for request_id: {request_id}") - return None - - request = db.query(file_retriever_models.FileRetrievalRequest).filter( - file_retriever_models.FileRetrievalRequest.id == request_uuid - ).options( - joinedload(file_retriever_models.FileRetrievalRequest.retrieved_files) - ).first() - - if not request: - return None - - retrieved_files = [] - - for file in request.retrieved_files: - retrieved_files.append(file.file_path) - - return retrieved_files - - # def _format_diff(self, raw_diff: str) -> str: - # # Remove Markdown-style code block markers - # content = re.sub(r'^```diff\n|```$', '', raw_diff.strip(), flags=re.MULTILINE) - - # # Unescape common sequences - # content = content.encode('utf-8').decode('unicode_escape') - - # return content - - # def _apply_diff(self, original_content: str, file_diff: str) -> str: - # """ - # Applies a unified diff to the original content and returns the new content. - - # Args: - # original_content: The original file content as a single string. - # file_diff: The unified diff string. - - # Returns: - # The new content with the diff applied. - # """ - # # Handle the case where the original content is empty. - # if not original_content: - # new_content: List[str] = [] - # for line in file_diff.splitlines(keepends=True): - # if line.startswith('+') and not line.startswith('+++'): - # new_content.append(line[1:]) - # return ''.join(new_content) - - # original_lines = original_content.splitlines(keepends=True) - # diff_lines = file_diff.splitlines(keepends=True) - - # i = 0 - # new_content: List[str] = [] - # orig_idx = 0 - - # while i < len(diff_lines): - # # Skip diff headers like --- and +++ - # if diff_lines[i].startswith('---') or diff_lines[i].startswith('+++'): - # i += 1 - # continue - - # # Hunk header - # if not diff_lines[i].startswith('@@'): - # i += 1 - # continue - - # hunk_header = diff_lines[i] - # m = re.match(r'^@@ -(\d+),?(\d*) \+(\d+),?(\d*) @@', hunk_header) - # if not m: - # raise ValueError(f"Invalid hunk header: {hunk_header.strip()}") - - # orig_start = int(m.group(1)) - 1 # convert from 1-based to 0-based index - # i += 1 - - # # Copy unchanged lines before this hunk - # while orig_idx < orig_start: - # new_content.append(original_lines[orig_idx]) - # orig_idx += 1 - - # # Process lines in hunk - # while i < len(diff_lines): - # line = diff_lines[i] - - # if line.startswith('@@'): - # # Start of next hunk - # break - # elif line.startswith(' '): - # # Context line - # if orig_idx < len(original_lines): - # new_content.append(original_lines[orig_idx]) - # orig_idx += 1 - # elif line.startswith('-'): - # # Removed line - # orig_idx += 1 - # elif line.startswith('+'): - # # Added line - # new_content.append(line[1:]) - # i += 1 - - # # Add remaining lines from original - # new_content.extend(original_lines[orig_idx:]) - - # return ''.join(new_content) - - - async def send_command(self, websocket: WebSocket, command_name: str, data: Dict[str, Any] = {}): - if command_name not in self.command_map: - raise ValueError(f"Unknown command: {command_name}") - - message_to_send = { - "type": self.command_map[command_name]["type"], - **data, - } - - await websocket.send_text(json.dumps(message_to_send)) - - async def dispatch_message(self, websocket: WebSocket, message: Dict[str, Any]): - """ - Routes an incoming message to the appropriate handler based on its 'type'. - Retrieves session state to maintain context. - """ - message_type = message.get("type") - # request_id = message.get("request_id") - # round_num = message.get("round") - - # In a real-world app, you'd retrieve historical data based on request_id or session_id - # For this example, we'll just print it. - # print(f"Received message of type '{message_type}' (request_id: {request_id}, round: {round_num})") - # logger.info(f"Received message: {message}") - - handler = self.message_handlers.get(message_type) - if handler: - logger.debug(f"Dispatching to handler for message type: {message_type}") - 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") - # if request_id is None: - # await websocket.send_text(json.dumps({ - # "type": "error", - # "content": "Error: request_id is missing in the response." - # })) - # return - # print(f"Received folder selected (request_id: {request_id}): Path: {path}") - - # After a folder is selected, the next step is to list its contents. - # This now uses the send_command helper. - # await self.send_command(websocket, "list_directory", data={"path": path, "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.""" - logger.debug(f"List directory response data: {data}") - files = data.get("files", []) - if not files: - await websocket.send_text(json.dumps({ - "type": "error", - "content": "Error: No files found in the directory." - })) - return - request_id = data.get("request_id") - if not request_id: - await websocket.send_text(json.dumps({ - "type": "error", - "content": "Error: request_id is missing in the response." - })) - return - try: - Validator().precheck_tokensize(files) - except TokenLimitExceededError as e: - await websocket.send_text(json.dumps({ - "type": "error", - "content": f"Error: {e}" - })) - return - file_request = await self._get_file_request_by_id(uuid.UUID(request_id)) - if not file_request: - await websocket.send_text(json.dumps({ - "type": "error", - "content": f"Error: No matching file request found for request_id {request_id}." - })) - return - await self._store_retrieved_files(request_id=uuid.UUID(request_id), files=files) - await self.handle_files_content_response(websocket, {"files": [], "request_id": request_id, "session_id": file_request.session_id}) - - async def handle_files_content_response(self, websocket: WebSocket, data: Dict[str, Any]): - """ - Handles the content of a list of files sent by the client. - """ - files_data: List[Dict[str, str]] = data.get("files", []) - request_id = data.get("request_id") - session_id = data.get("session_id") - - if not request_id: - await websocket.send_text(json.dumps({ - "type": "error", - "content": "Error: request_id is required to process file content." - })) - return - - if not session_id: - await websocket.send_text(json.dumps({ - "type": "error", - "content": "Error: session_id is required to process file content." - })) - return - - if not files_data: - print(f"Warning: No files data received for request_id: {request_id}") - else: - print(f"Received content for {len(files_data)} files (request_id: {request_id}).") - await self._update_file_content(request_id=uuid.UUID(request_id), files_with_content=files_data) - - # Retrieve the updated context from the database - context_data = await self._retrieve_by_request_id(self.db, request_id=request_id) - - if not context_data: - print(f"Error: Context not found for request_id: {request_id}") - await websocket.send_text(json.dumps({ - "type": "error", - "content": "An internal error occurred. Please try again." - })) - return - - await websocket.send_text(json.dumps({ - "type": "thinking_log", - "content": f"AI is analyzing the retrieved files to determine next steps." - })) - - session = self.db.query(models.Session).options( - joinedload(models.Session.messages) - ).filter(models.Session.id == session_id).first() - - # Use the LLM to make a decision - llm_provider = get_llm_provider(provider_name="gemini") - crqd = CodeRagQuestionDecider() - original_question = context_data.get("question", "") - try: - raw_answer_text, reasoning, decision, instructions = await crqd.forward( - question=original_question, - history=session.messages, - retrieved_data=context_data, - llm_provider=llm_provider - ) - except ValueError as e: - await websocket.send_text(json.dumps({ - "type": "error", - "content": f"Failed to process AI decision request. Error: {e}" - })) - return - - if decision == "answer": - # Handle regular answer - assistant_message = models.Message(session_id=session_id, sender="assistant", content=raw_answer_text) - self.db.add(assistant_message) - self.db.commit() - self.db.refresh(assistant_message) - await websocket.send_text(json.dumps({ - "type": "chat_message", - "content": raw_answer_text, - "reasoning": reasoning - })) - - elif decision == "files": - # Handle file retrieval request - await websocket.send_text(json.dumps({ - "type": "thinking_log", - "content": f"AI decided files are needed: {raw_answer_text}." - })) - try: - # raw_answer_text might be a JSON string or a list from the LLM - if isinstance(raw_answer_text, list): - answer_text = raw_answer_text - else: - json_match = re.search(r'\[.*\]', raw_answer_text, re.DOTALL) - if json_match: - json_string = json_match.group(0) - answer_text = ast.literal_eval(json_string) - else: - answer_text = ast.literal_eval(raw_answer_text) - - if not isinstance(answer_text, list): - raise ValueError("Parsed result is not a list.") - except (ValueError, SyntaxError) as e: - print(f"Error parsing LLM output: {e}") - answer_text = [] - await websocket.send_text(json.dumps({ - "type": "thinking_log", - "content": f"Warning: AI's file list could not be parsed. Error: {e}" - })) - return - assistant_message = models.Message(session_id=session_id, sender="assistant", content=f"{reasoning}\n Request Files: {answer_text} ") - self.db.add(assistant_message) - self.db.commit() - self.db.refresh(assistant_message) - await self.send_command(websocket, "get_file_content", data={"filepaths": answer_text, "request_id": request_id}) - - elif decision == "code_change": - # Handle code change request - await websocket.send_text(json.dumps({ - "type": "thinking_log", - "content": "AI is generating the necessary code changes. This may take a moment." - })) - - try: - # The input_data should be a JSON string of code change instructions - input_data = json.dumps(instructions) if instructions else "[]" - cch = CodeChangeHelper( - db=self.db, - provider_name="gemini", - original_question=original_question, - input_data=input_data, - reasoning=reasoning, - request_id=uuid.UUID(request_id) - ) - - # Use the CodeChangeHelper to process all code changes - await cch.process(websocket=websocket) - - except (json.JSONDecodeError, ValueError) as e: - logger.error(f"Error processing code changes: {e}") - await websocket.send_text(json.dumps({ - "type": "error", - "content": f"Failed to process code change request. Error: {e}" - })) - - else: # Fallback for any other decision - await websocket.send_text(json.dumps({ - "type": "thinking_log", - "content": f"Answering user's question directly." - })) - await websocket.send_text(json.dumps({ - "type": "chat_message", - "content": raw_answer_text, - "reasoning": reasoning - })) - - - 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." - })) - - async def handle_chat_message(self, websocket: WebSocket, data: Dict[str, Any]): - """Handles incoming chat messages from the client.""" - # TODO: Enhance this function to process the chat message and determine the next action. - prompt = data.get("content") - provider_name = data.get("provider_name", "gemini") - session_id = data.get("session_id") - if session_id is None: - await websocket.send_text(json.dumps({ - "type": "error", - "content": "Error: session_id is required for chat messages." - })) - return - session = self.db.query(models.Session).options( - joinedload(models.Session.messages) - ).filter(models.Session.id == session_id).first() - - if not session: - await websocket.send_text(json.dumps({ - "type": "error", - "content": f"Error: Session with ID {session_id} not found." - })) - return - user_message = models.Message(session_id=session_id, sender="user", content=prompt) - self.db.add(user_message) - self.db.commit() - self.db.refresh(user_message) - - # Auto-title the session from the very first user message - if session.title in (None, "New Chat Session", ""): - session.title = prompt[:60].strip() + ("..." if len(prompt) > 60 else "") - - # Keep provider_name in sync with the model actually being used - if session.provider_name != provider_name: - session.provider_name = provider_name - - self.db.commit() - - path = data.get("path", "") - if path: - # If file path is provided, initiate file retrieval process. - file_request = await self._get_or_create_file_request(session_id, path, prompt) - await self.send_command(websocket, "list_directory", data={"request_id": str(file_request.id)}) - return - # Fetch user preferences for overrides - api_key_override = None - model_name_override = "" - user = session.user - if user and user.preferences: - llm_prefs = user.preferences.get("llm", {}).get("providers", {}).get(provider_name, {}) - api_key_override = llm_prefs.get("api_key") - model_name_override = llm_prefs.get("model", "") - - # Get the appropriate LLM provider - llm_provider = get_llm_provider( - provider_name, - model_name=model_name_override, - api_key_override=api_key_override - ) - chat = RagPipeline() - answer_text = await chat.forward( - question=prompt, - history=session.messages, - context_chunks=[], - llm_provider=llm_provider - ) - # Save assistant's response - assistant_message = models.Message(session_id=session_id, sender="assistant", content=answer_text) - self.db.add(assistant_message) - self.db.commit() - self.db.refresh(assistant_message) - - # 📝 Add this section to send the response back to the client - # The client-side `handleChatMessage` handler will process this message - await websocket.send_text(json.dumps({ - "type": "chat_message", - "content": answer_text - })) - logger.info(f"Sent chat response to client: {answer_text}") \ No newline at end of file diff --git a/ai-hub/app/core/skills/base.py b/ai-hub/app/core/skills/base.py new file mode 100644 index 0000000..6670e81 --- /dev/null +++ b/ai-hub/app/core/skills/base.py @@ -0,0 +1,47 @@ +from abc import ABC, abstractmethod +from typing import Dict, Any, Optional +from pydantic import BaseModel + +class SkillResult(BaseModel): + success: bool + output: Any + error: Optional[str] = None + +class BaseSkill(ABC): + """ + Abstract Base Class for all AI-Hub Skills (Tools). + """ + + @property + @abstractmethod + def name(self) -> str: + """Unique name of the skill.""" + pass + + @property + @abstractmethod + def description(self) -> str: + """Description of what the skill does, used by LLM.""" + pass + + @property + @abstractmethod + def parameters_schema(self) -> Dict[str, Any]: + """JSON Schema of the input parameters.""" + pass + + @abstractmethod + async def execute(self, **kwargs) -> SkillResult: + """The actual logic of the skill.""" + pass + + def to_tool_definition(self) -> Dict[str, Any]: + """Converts the skill to an LLM tool definition.""" + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters_schema + } + } diff --git a/ai-hub/app/db/file_retriever_models.py b/ai-hub/app/db/file_retriever_models.py deleted file mode 100644 index 9b562e1..0000000 --- a/ai-hub/app/db/file_retriever_models.py +++ /dev/null @@ -1,74 +0,0 @@ -from datetime import datetime -from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON -from sqlalchemy.dialects.postgresql import UUID as PG_UUID -from sqlalchemy.orm import relationship -import uuid - -# Assuming Base is imported from your database.py -from .database import Base - -class FileRetrievalRequest(Base): - """ - SQLAlchemy model for the 'file_retrieval_requests' table. - - Each entry represents a single user request to process a directory - for the AI coding assistant. - """ - __tablename__ = 'file_retrieval_requests' - - # Primary key, a UUID for unique and secure identification. - id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True) - # The client's main question or instruction for the AI. - question = Column(Text, nullable=False) - # The path or name of the directory provided by the client. - directory_path = Column(String, nullable=False) - # Foreign key to link this request to a specific chat session. - session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False) - # Timestamp for when the request was made. - created_at = Column(DateTime, default=datetime.now, nullable=False) - - # Defines a one-to-many relationship with the RetrievedFile table. - # 'cascade' ensures that all associated files are deleted when the request is deleted. - retrieved_files = relationship( - "RetrievedFile", - back_populates="retrieval_request", - cascade="all, delete-orphan" - ) - - def __repr__(self): - return f"" - -class RetrievedFile(Base): - """ - SQLAlchemy model for the 'retrieved_files' table. - - This table stores the content and metadata for each file retrieved - during a file processing request. - """ - __tablename__ = 'retrieved_files' - - # Primary key for the file entry. - id = Column(Integer, primary_key=True, index=True) - # Foreign key linking this file back to its parent retrieval request. - request_id = Column(PG_UUID(as_uuid=True), ForeignKey('file_retrieval_requests.id'), nullable=False) - # The full path to the file. - file_path = Column(String, nullable=False) - # The name of the file. - file_name = Column(String, nullable=False) - # The actual content of the file. - content = Column(Text, nullable=False) - # The type of the file content (e.g., 'original', 'updated'). - type = Column(String, nullable=False, default='original') - # Timestamp for when the file was last modified. - last_updated = Column(DateTime, default=datetime.now) - # Timestamp for when this entry was created in the database. - created_at = Column(DateTime, default=datetime.now, nullable=False) - - # Defines a many-to-one relationship back to the FileRetrievalRequest. - retrieval_request = relationship( - "FileRetrievalRequest", - back_populates="retrieved_files" - ) - - def __repr__(self): - return f"" \ No newline at end of file diff --git a/ai-hub/app/db/models.py b/ai-hub/app/db/models.py index 08cc446..9f5b439 100644 --- a/ai-hub/app/db/models.py +++ b/ai-hub/app/db/models.py @@ -225,3 +225,105 @@ Provides a helpful string representation of the object for debugging. """ return f"" + + +# --- New Asset Management Models --- + +class PromptTemplate(Base): + """ + SQLAlchemy model for centralized system prompts. + """ + __tablename__ = 'prompt_templates' + + id = Column(Integer, primary_key=True, index=True) + slug = Column(String, unique=True, index=True, nullable=False) + title = Column(String, nullable=False) + content = Column(Text, nullable=False) + version = Column(Integer, default=1) + + owner_id = Column(String, ForeignKey('users.id'), nullable=False) + group_id = Column(String, ForeignKey('groups.id'), nullable=True) + is_public = Column(Boolean, default=False) + + created_at = Column(DateTime, default=datetime.utcnow) + updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) + + owner = relationship("User") + group = relationship("Group") + + def __repr__(self): + return f"" + +class Skill(Base): + """ + SQLAlchemy model for AI capabilities (Skills/Tools). + """ + __tablename__ = 'skills' + + id = Column(Integer, primary_key=True, index=True) + name = Column(String, unique=True, index=True, nullable=False) + description = Column(String, nullable=True) + # type: 'local', 'remote_grpc', 'mcp' + skill_type = Column(String, default="local", nullable=False) + # Stores tool definition, parameters, or endpoint config + config = Column(JSON, default={}, nullable=True) + + owner_id = Column(String, ForeignKey('users.id'), nullable=False) + group_id = Column(String, ForeignKey('groups.id'), nullable=True) + is_system = Column(Boolean, default=False) + + created_at = Column(DateTime, default=datetime.utcnow) + + owner = relationship("User") + group = relationship("Group") + + def __repr__(self): + return f"" + +class MCPServer(Base): + """ + SQLAlchemy model for Model Context Protocol (MCP) server configurations. + """ + __tablename__ = 'mcp_servers' + + id = Column(Integer, primary_key=True, index=True) + name = Column(String, nullable=False) + url = Column(String, nullable=False) + auth_config = Column(JSON, default={}, nullable=True) + + owner_id = Column(String, ForeignKey('users.id'), nullable=False) + group_id = Column(String, ForeignKey('groups.id'), nullable=True) + + created_at = Column(DateTime, default=datetime.utcnow) + + owner = relationship("User") + group = relationship("Group") + + def __repr__(self): + return f"" + +class AssetPermission(Base): + """ + SQLAlchemy model for granular permission control on assets. + """ + __tablename__ = 'asset_permissions' + + id = Column(Integer, primary_key=True, index=True) + # resource_type: 'prompt', 'skill', 'mcp_server' + resource_type = Column(String, nullable=False, index=True) + resource_id = Column(Integer, nullable=False, index=True) + + # Grant to a specific user OR a specific group + user_id = Column(String, ForeignKey('users.id'), nullable=True) + group_id = Column(String, ForeignKey('groups.id'), nullable=True) + + # access_level: 'view', 'execute', 'admin' + access_level = Column(String, default="execute", nullable=False) + + created_at = Column(DateTime, default=datetime.utcnow) + + user = relationship("User") + group = relationship("Group") + + def __repr__(self): + return f"" diff --git a/ai-hub/integration_tests/conftest.py b/ai-hub/integration_tests/conftest.py index c878803..97d1da2 100644 --- a/ai-hub/integration_tests/conftest.py +++ b/ai-hub/integration_tests/conftest.py @@ -59,30 +59,6 @@ assert delete_response.status_code == 200 -@pytest_asyncio.fixture(scope="function") -async def websocket_client(base_url, session_id): - """ - Fixture to provide an active, connected WebSocket client for testing the - /ws/workspace/{session_id} endpoint. - - The client will be disconnected and closed at the end of the test. - """ - # Replace 'http' with 'ws' for the WebSocket URL scheme - ws_url = base_url.replace("http", "ws") - - # We use httpx.AsyncClient to handle the WebSocket connection. - client = httpx.AsyncClient() - - # Context manager handles the connection lifecycle (connect, disconnect, close) - async with client.websocket_connect(f"{ws_url}/ws/workspace/{session_id}") as websocket: - # The first message from the server should be the 'connection_established' - # message after the initial accept(). We read it to ensure the connection - # is fully established before yielding. - initial_message = await websocket.receive_text() - print(f"\nReceived initial WS message: {initial_message}") - - # The fixture yields the active WebSocket object - yield websocket - - # When the context manager exits, the websocket connection is automatically closed. - await client.aclose() \ No newline at end of file + # Teardown: Delete the document after the test + delete_response = await http_client.delete(f"/documents/{document_id}") + assert delete_response.status_code == 200 \ No newline at end of file diff --git a/ai-hub/scripts/__init__.py b/ai-hub/scripts/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/ai-hub/scripts/__init__.py diff --git a/ai-hub/scripts/seed_prompts.py b/ai-hub/scripts/seed_prompts.py new file mode 100644 index 0000000..74bac06 --- /dev/null +++ b/ai-hub/scripts/seed_prompts.py @@ -0,0 +1,206 @@ +import sys +import os + +# Ensure the app directory is in the path +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +from app.db.session import SessionLocal +from app.db.models import PromptTemplate, User +import logging + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# Prompts extracted from code +PROMPTS = { + "question-decider": { + "title": "Code RAG Question Decider", + "content": """### 🧠 **Core Directives** + +You are a specialized AI assistant for software engineering tasks. Your responses—providing an answer, suggesting a code change, or requesting more files—must be based **exclusively** on the provided codebase content. Your primary goal is to be helpful and accurate while adhering strictly to the following directives. + +----- + +## 1. Data Analysis and Availability + +* **Analyze the User's Request:** Carefully examine the **`question`** and **`chat_history`** to understand what the user wants. +* **Source of Information:** The only information you can use to generate a code-related answer comes from the files provided in the **`retrieved_paths_with_content`** list. + +* **File Data & Availability** +* **`retrieved_paths_with_content`**: Files with content available. +* **`retrieved_paths_without_content`**: Files that exist but content is not loaded. + +----- + +## 2. Decision Logic + +You must choose one of three mutually exclusive decisions: `answer`, `code_change`, or `files`. + +### `decision='answer'` +* Choose this if you have all necessary info to explain a non-code-modification question. + +### `decision='code_change'` +* Choose this for any code manipulation (modify, create, delete). +* Provide a high-level strategy plan in the `answer` field as a numbered list. +* Provide the actual code instructions in a valid JSON list format. + +### `decision='files'` +* Request more files from `retrieved_paths_without_content`. + +----- + +## 3. Output Format + +You MUST respond in valid JSON format with the following fields: +- `reasoning`: Your step-by-step logic. +- `decision`: Either 'answer', 'files', or 'code_change'. +- `answer`: Depending on decision (Markdown text, file list, or high-level plan). +- `instructions`: (Only for 'code_change') The JSON list of file operations. + +User Question: {question} +Chat History: {chat_history} +Available Content: {retrieved_paths_with_content} +Missing Content: {retrieved_paths_without_content} + +Strict JSON Output:""" + }, + "code-changer": { + "title": "Code RAG Code Changer", + "content": """### 🧠 Core Directives + +You are a code generation assistant specialized in producing **one precise and complete code change** per instruction. Your output must be a strict JSON object containing: + +- `reasoning`: A concise explanation of the change. +- `content`: The **full content of the file** (or an empty string for deletions). + +--- + +### 1. Input Structure + +- `overall_plan`: {overall_plan} +- `instruction`: {instruction} +- `filepath`: {filepath} +- `original_files`: {original_files} +- `updated_files`: {updated_files} + +----- + +### 2. 💻 Code Generation Rules + +Please provide **one complete and functional code file** per request, for the specified `file_path`. You must output the **entire, modified file**. + +* **Identical Code Sections:** Use the `#[unchanged_section]|||` syntax for large, sequential blocks of code that are not being modified. +* **Complete File Output:** Always provide the **full file contents** in the `content` block. Do not use placeholders like `...`. +* **Imports:** Ensure all required imports are included. + +--- + +### 3. Output Format + +Return exactly one JSON object: +{{ + "reasoning": "Brief explanation.", + "content": "Full file content" +}}""" + }, + "code-reviewer": { + "title": "Code RAG Reviewer", + "content": """### 🧠 Core Directives + +### **Code Review Directives** +Your role is a specialized code review AI. Your primary task is to review a set of code changes and confirm they **fully and accurately address the user's original request**. + +--- +### **Critical Constraints** +Your review is strictly limited to **code content completeness**. Do not suggest or perform any file splits, moves, or large-scale refactoring. +Identify and resolve any missing logic, placeholders (like "unchanged," "same as original," "to-do"), or incomplete code. + +Return exactly one JSON object: +{{ + "reasoning": "A detailed explanation of why the decision was made.", + "decision": "Either 'complete' or 'modify'.", + "answer": "If 'complete', an empty string. If 'modify', the new execution plan instructions in JSON." +}} + +Input: +- `original_question`: {original_question} +- `execution_plan`: {execution_plan} +- `final_code_changes`: {final_code_changes} +- `original_files`: {original_files}""" + }, + "file-selector": { + "title": "Code RAG File Selector", + "content": """You're an **expert file navigator** for a large codebase. Your task is to select the most critical and relevant file paths to answer a user's question. All file paths you select must exist within the provided `retrieved_files` list. + +--- + +### File Selection Criteria + +1. **Prioritize Core Files:** Identify files that contain the central logic. +2. **Be Selective:** Aim for **2 to 4 files**. +3. **Exclude Irrelevant and Unreadable Files:** Ignore binaries, images, etc. +4. **Infer User Intent:** Return only file paths that exist in the `retrieved_files` list. + +--- + +### Output Format + +Return exactly one JSON array of strings: +[ + "/path/to/file1", + "/path/to/file2" +] + +Input: +- `question`: {question} +- `chat_history`: {chat_history} +- `retrieved_files`: {retrieved_files}""" + }, + "rag-pipeline": { + "title": "Default RAG Pipeline", + "content": """Generate a natural and context-aware answer to the user's question using the provided knowledge and conversation history. + +Relevant excerpts from the knowledge base: +{context} + +Conversation History: +{chat_history} + +User Question: {question} + +Answer:""" + } +} + +def seed(): + db = SessionLocal() + try: + user = db.query(User).first() + if not user: + logger.error("No users found in database. Please run migrations and create a user first.") + return + + for slug, data in PROMPTS.items(): + existing = db.query(PromptTemplate).filter(PromptTemplate.slug == slug).first() + if not existing: + logger.info(f"Seeding prompt: {slug}") + prompt = PromptTemplate( + slug=slug, + title=data["title"], + content=data["content"], + owner_id=user.id, + is_public=True + ) + db.add(prompt) + else: + logger.info(f"Prompt '{slug}' already exists, skipping.") + + db.commit() + except Exception as e: + db.rollback() + logger.error(f"Error seeding prompts: {e}") + finally: + db.close() + +if __name__ == "__main__": + seed() diff --git a/ai-hub/tests/api/test_dependencies.py b/ai-hub/tests/api/test_dependencies.py index 454a2fd..7edf391 100644 --- a/ai-hub/tests/api/test_dependencies.py +++ b/ai-hub/tests/api/test_dependencies.py @@ -12,7 +12,6 @@ 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 from app.core.retrievers.base_retriever import Retriever diff --git a/ui/client-app/src/App.js b/ui/client-app/src/App.js index 8745c91..8e60cd7 100644 --- a/ui/client-app/src/App.js +++ b/ui/client-app/src/App.js @@ -144,7 +144,7 @@ }; return ( -
+
{currentPage !== "login" && ( )} -
+
{renderPage()}
diff --git a/ui/client-app/src/components/CodeFolderAccess.js b/ui/client-app/src/components/CodeFolderAccess.js deleted file mode 100644 index 07e07e5..0000000 --- a/ui/client-app/src/components/CodeFolderAccess.js +++ /dev/null @@ -1,42 +0,0 @@ -// src/components/CodeFolderAccess.js -import React from "react"; - -const CodeFolderAccess = ({ selectedFolder, onSelectFolder }) => { - const handleFolderChange = async () => { - // This is a placeholder. In a real application, you would use a browser API - // or a desktop framework like Electron to access the file system. - // Example using the File System Access API (requires user gesture) - try { - const directoryHandle = await window.showDirectoryPicker(); - const folderPath = directoryHandle.name; // Simplified path, you might get a full path in a real app - onSelectFolder(directoryHandle); - } catch (err) { - console.warn("Failed to select directory:", err); - // You might want to handle user cancellation here - } - }; - - return ( -
- {/* Button to trigger the folder selection dialog */} - - - {/* Display the selected folder path */} -
-

- Selected Folder: -

-

- {selectedFolder || "No folder selected"} -

-
-
- ); -}; - -export default CodeFolderAccess; \ No newline at end of file diff --git a/ui/client-app/src/components/InteractionLog.js b/ui/client-app/src/components/InteractionLog.js deleted file mode 100644 index 9ecb6ea..0000000 --- a/ui/client-app/src/components/InteractionLog.js +++ /dev/null @@ -1,91 +0,0 @@ -import React, { useState, useEffect, useRef } from "react"; -import ReactMarkdown from 'react-markdown'; - -const InteractionLog = ({ logs }) => { - const [expandedLogs, setExpandedLogs] = useState({}); - const [containerHeight, setContainerHeight] = useState("500px"); - - const scrollRef = useRef(null); - - // Auto-size based on window height - useEffect(() => { - const updateHeight = () => { - const windowHeight = window.innerHeight; - const offset = 300; // adjust this based on layout - setContainerHeight(`${windowHeight - offset}px`); - }; - - updateHeight(); - window.addEventListener("resize", updateHeight); - return () => window.removeEventListener("resize", updateHeight); - }, []); - - // Auto-scroll to bottom on new logs - useEffect(() => { - if (scrollRef.current) { - scrollRef.current.scrollTop = scrollRef.current.scrollHeight; - } - }, [logs]); - - const toggleExpand = (index) => { - setExpandedLogs((prev) => ({ - ...prev, - [index]: !prev[index], - })); - }; - - const getPreviewText = (message, maxLength = 200) => { - if (message.length <= maxLength) return message; - return message.slice(0, maxLength) + "…"; - }; - - return ( -
-
- {logs.length > 0 ? ( - logs.map((log, index) => { - const isExpanded = expandedLogs[index]; - - return ( -
toggleExpand(index)} - > -

- {log.type.charAt(0).toUpperCase() + log.type.slice(1)}: -

- {isExpanded ? log.message : getPreviewText(log.message)} - {!isExpanded && log.message.length > 200 && ( -

- Click to expand -

- )} -
- ); - }) - ) : ( -

- No interactions yet. -

- )} -
-
- ); -}; - -export default InteractionLog; diff --git a/ui/client-app/src/components/VoiceControls.js b/ui/client-app/src/components/VoiceControls.js index 23d960b..984b079 100644 --- a/ui/client-app/src/components/VoiceControls.js +++ b/ui/client-app/src/components/VoiceControls.js @@ -1,4 +1,4 @@ -// src/components/Controls.js +// src/components/VoiceControls.js import React from "react"; import { FaMicrophone, FaRegStopCircle } from "react-icons/fa"; @@ -12,43 +12,53 @@ onToggleAutoMode, }) => { const micButtonColorClass = isRecording - ? "bg-red-600 hover:bg-red-700 active:bg-red-800" - : "bg-indigo-600 hover:bg-indigo-700 active:bg-indigo-800"; + ? "bg-red-500 hover:bg-red-600 active:bg-red-700 shadow-red-500/20" + : "bg-indigo-600 hover:bg-indigo-700 active:bg-indigo-800 shadow-indigo-500/20"; const micButtonState = isAutoMode && isAutoListening ? isAutoListening : isRecording; return ( -
-
- {status} +
+ {/* Status indicator */} +
+
+
+ {status || (isBusy ? "Thinking..." : "Ready")} +
-
+ +
+ {/* Mic Toggle Button */} -
- -
+ {/* Auto Mode Toggle */} +
); diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 56a781d..ec8bcde 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -1,15 +1,10 @@ import { useState, useEffect, useRef, useCallback } from "react"; -import { connectToWebSocket, getSessionId } from "../services/websocket"; -import { getSessionTokenStatus, getSessionMessages } from "../services/apiService"; -import { v4 as uuidv4 } from 'uuid'; +import { getSessionId } from "../services/websocket"; +import { getSessionTokenStatus, getSessionMessages, chatWithAI, getUserConfig, getSession } from "../services/apiService"; const useCodeAssistant = ({ pageContainerRef }) => { const [chatHistory, setChatHistory] = useState([]); - const [thinkingProcess, setThinkingProcess] = useState([]); - const [selectedFolder, setSelectedFolder] = useState(null); - const [connectionStatus, setConnectionStatus] = useState("disconnected"); const [isProcessing, setIsProcessing] = useState(false); - const [isPaused, setIsPaused] = useState(false); const [errorMessage, setErrorMessage] = useState(""); const [showErrorModal, setShowErrorModal] = useState(false); const [isConfigured, setIsConfigured] = useState(true); @@ -19,10 +14,8 @@ const [userConfigData, setUserConfigData] = useState(null); const [localActiveLLM, setLocalActiveLLM] = useState(''); - const sessionIdRef = useRef(null); // ✅ Always current sessionId - const ws = useRef(null); + const sessionIdRef = useRef(null); const initialized = useRef(false); - const dirHandleRef = useRef(null); const fetchTokenUsage = useCallback(async () => { if (!sessionIdRef.current) return; @@ -34,359 +27,78 @@ } }, []); - const handleChatMessage = useCallback((message) => { - console.log("Received chat message:", message); - setThinkingProcess((prev) => [...prev, { - type: "system", - message: "AI processing is complete" - }]) - setChatHistory((prev) => [...prev, { - isUser: false, - isPureAnswer: true, - text: message.content, - dicision: message.dicision, - reasoning: message.reasoning - }]); - setIsProcessing(false); - }, []); - - const handleCodeChange = useCallback((message) => { - console.log("Received code change:", message); - setChatHistory((prev) => [...prev, { - isUser: false, - isPureAnswer: false, - text: message.content, - code_changes: message.code_changes, - steps: message.steps, - reasoning: message.reasoning - }]); - if (message.done === true) { - setThinkingProcess((prev) => [...prev, { - type: "system", - message: "AI processing is complete" - }]) - setIsProcessing(false); - } else { - setIsProcessing(true); - } - }, []); - - - const handleThinkingLog = useCallback((message) => { - setThinkingProcess((prev) => [...prev, { - type: "remote", - 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 { request_id } = message; - console.log("Received list directory request:", 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, request_id })); - return; - } - + const fetchSessionHistory = useCallback(async (sid) => { try { - const files = []; - setThinkingProcess((prev) => [...prev, { - type: "local", - message: `Scanning the directory: ${dirHandle.name}...`, - }]); - async function walkDirectory(handle, path = '') { - for await (const entry of handle.values()) { - const entryPath = `${path}/${entry.name}`; - if (entry.kind === "file") { - const file = await entry.getFile(); - files.push({ - name: file.name, - path: entryPath, - size: file.size, - lastModified: file.lastModified, - created: file.created, - }); - } else if (entry.kind === "directory") { - await walkDirectory(entry, entryPath); - } - } + const messagesData = await getSessionMessages(sid); + if (messagesData && messagesData.messages) { + const formattedHistory = messagesData.messages.map((msg) => ({ + isUser: msg.sender === "user", + isPureAnswer: true, + text: msg.content, + })); + setChatHistory(formattedHistory); } - - await walkDirectory(dirHandle); - - ws.current.send(JSON.stringify({ - type: "list_directory_response", - files, - request_id, - session_id: sessionIdRef.current, // ✅ Always current - })); - - setThinkingProcess((prev) => [...prev, { - type: "local", - message: `Sent ${files.length} files metadata information to server.`, - }]); - } catch (error) { - console.error("Failed to list directory:", error); - ws.current.send(JSON.stringify({ - type: "error", - content: "Failed to access folder contents.", - request_id, - session_id: sessionIdRef.current, - })); + } catch (err) { + console.warn("Failed to load chat history", err); } }, []); - const getFileHandleFromPath = async (dirHandle, filePath) => { - const pathParts = filePath.split('/').filter(Boolean); - let currentHandle = dirHandle; - - for (let i = 0; i < pathParts.length; i++) { - const part = pathParts[i]; - try { - if (i === pathParts.length - 1) { - return await currentHandle.getFileHandle(part); - } else { - currentHandle = await currentHandle.getDirectoryHandle(part); - } - } catch (error) { - console.error(`Path not found: ${filePath}`, error); - return null; - } - } - return null; - }; - - const handleReadFilesRequest = useCallback(async (message) => { - const { filepaths, request_id } = message; - console.log("Received read files request:", message); - const dirHandle = dirHandleRef.current; - - if (!dirHandle) { - ws.current.send(JSON.stringify({ - type: "error", - content: "No folder selected.", - request_id, - session_id: sessionIdRef.current, - })); - return; - } - - const filesData = []; - const readFiles = []; - - for (const filepath of filepaths) { - try { - const fileHandle = await getFileHandleFromPath(dirHandle, filepath); - if (!fileHandle) { - filesData.push({ filepath, content: null }); - } - const file = await fileHandle.getFile(); - const content = await file.text(); - filesData.push({ filepath, content }); - readFiles.push(filepath); - } catch (error) { - console.warn(`Failed to read file: ${filepath}`, error); - // ws.current.send(JSON.stringify({ - // type: "error", - // content: `Could not read file: ${filepath}`, - // request_id, - // session_id: sessionIdRef.current, - // })); - } - } - - ws.current.send(JSON.stringify({ - type: "file_content_response", - files: filesData, - request_id, - session_id: sessionIdRef.current, - })); - - setThinkingProcess((prev) => { - const newMessages = []; - - if (readFiles.length > 0) { - const displayMessage = readFiles.length > 10 - ? `Read ${readFiles.length} files successfully.` - : `Read files and send successfully: [${readFiles.map(f => `"${f}"`).join(', ')}]`; - newMessages.push({ type: "local", message: displayMessage }); - } - return [...prev, ...newMessages]; - }); - }, []); - - const handleExecuteCommandRequest = useCallback((message) => { - const { command, request_id } = message; - const output = `Simulated output for command: '${command}'`; - - ws.current.send(JSON.stringify({ - type: "execute_command_response", - command, - output, - request_id, - session_id: sessionIdRef.current, - })); - - setThinkingProcess((prev) => [...prev, { - type: "system", - message: `Simulated execution of command: '${command}'` - }]); - }, []); - - const handleIncomingMessage = useCallback((message) => { - switch (message.type) { - case "chat_message": - handleChatMessage(message); - break; - case "code_change": - handleCodeChange(message); - break; - case "thinking_log": - handleThinkingLog(message); - break; - case "error": - handleError(message); - break; - case "status_update": - handleStatusUpdate(message); - break; - case "list_directory": - handleListDirectoryRequest(message); - break; - case "get_file_content": - handleReadFilesRequest(message); - break; - case "execute_command": - handleExecuteCommandRequest(message); - break; - default: - console.log("Unknown message type:", message); - } - }, [ - handleChatMessage, - handleCodeChange, - handleThinkingLog, - handleError, - handleStatusUpdate, - handleListDirectoryRequest, - handleReadFilesRequest, - handleExecuteCommandRequest - ]); - - // WebSocket Setup + // Setup useEffect(() => { if (initialized.current) return; initialized.current = true; - const setupConnection = async () => { + const setup = async () => { try { - let configDataToUse = null; - let providerToUse = "deepseek"; + let configData = null; + let provider = "gemini"; try { - const { getUserConfig } = require("../services/apiService"); - configDataToUse = await getUserConfig(); - setUserConfigData(configDataToUse); - if (configDataToUse.effective?.llm?.active_provider) { - providerToUse = configDataToUse.effective.llm.active_provider; + configData = await getUserConfig(); + setUserConfigData(configData); + if (configData.effective?.llm?.active_provider) { + provider = configData.effective.llm.active_provider; } } catch (e) { - console.warn("Could not load user config for Assistant", e); + console.warn("Could not load user config", e); } - // Ensure we have a valid session before connecting or handling messages - const currentSessionId = await getSessionId("coding_assistant", providerToUse); - setSessionId(currentSessionId); - sessionIdRef.current = currentSessionId; + const sid = await getSessionId("coding_assistant", provider); + setSessionId(sid); + sessionIdRef.current = sid; - let llmToSet = providerToUse; + let llm = provider; try { - const { getSession } = require("../services/apiService"); - const sessionInfo = await getSession(currentSessionId); + const sessionInfo = await getSession(sid); if (sessionInfo && sessionInfo.provider_name) { - llmToSet = sessionInfo.provider_name; + llm = sessionInfo.provider_name; } - } catch (e) { console.warn("Could not check session provider config", e); } + } catch (e) { console.warn("Could not check session provider", e); } + setLocalActiveLLM(llm); - setLocalActiveLLM(llmToSet); + // Config check + const eff = configData?.effective || {}; + const missing = []; + const llmProviders = eff.llm?.providers || {}; + const hasLLMKey = Object.values(llmProviders).some(p => p.api_key && p.api_key !== 'None'); + if (!hasLLMKey) missing.push("Language Model (LLM) API Key"); - // Check if configuration is fully populated - try { - const eff = configDataToUse?.effective || {}; - const missing = []; - - const llmProviders = eff.llm?.providers || {}; - const hasLLMKey = Object.values(llmProviders).some(p => p.api_key && p.api_key !== 'None'); - if (!hasLLMKey) missing.push("Language Model (LLM) API Key"); - - if (missing.length > 0) { - setIsConfigured(false); - setMissingConfigs(missing); - } else { - setIsConfigured(true); - setMissingConfigs([]); - } - } catch (configErr) { - console.warn("Failed to check LLM configuration", configErr); + if (missing.length > 0) { + setIsConfigured(false); + setMissingConfigs(missing); + } else { + setIsConfigured(true); + setMissingConfigs([]); } - try { - const messagesData = await getSessionMessages(currentSessionId); - if (messagesData && messagesData.messages) { - const formattedHistory = messagesData.messages.map((msg) => ({ - isUser: msg.sender === "user", - isPureAnswer: true, - text: msg.content, - })); - setChatHistory(formattedHistory); - } - } catch (historyErr) { - console.warn("Failed to load chat history", historyErr); - } - - fetchTokenUsage(); - - const { ws: newWs, clientId } = await connectToWebSocket( - handleIncomingMessage, - () => setConnectionStatus("connected"), - () => { - setConnectionStatus("disconnected"); - setIsProcessing(false); - }, - (error) => { - setConnectionStatus("error"); - setErrorMessage(`Failed to connect: ${error.message}`); - setShowErrorModal(true); - } - ); - ws.current = newWs; + await fetchSessionHistory(sid); + await fetchTokenUsage(); } catch (error) { console.error("Setup failed:", error); } }; - setupConnection(); - - return () => { - if (ws.current) { - ws.current.close(); - } - }; - }, [handleIncomingMessage]); + setup(); + }, [fetchSessionHistory, fetchTokenUsage]); const handleSendChat = useCallback(async (text) => { if (!isConfigured && text.trim().toLowerCase() !== "/new") { @@ -397,120 +109,61 @@ if (text.trim().toLowerCase() === "/new") { setChatHistory([]); - setThinkingProcess([{ - type: "system", - message: "Started a new session." - }]); localStorage.removeItem("sessionId_coding_assistant"); - const newSessionId = await getSessionId("coding_assistant", localActiveLLM || "deepseek"); - setSessionId(newSessionId); - sessionIdRef.current = newSessionId; + const newSid = await getSessionId("coding_assistant", localActiveLLM || "gemini"); + setSessionId(newSid); + sessionIdRef.current = newSid; fetchTokenUsage(); return; } - 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, - session_id: sessionIdRef.current, - provider_name: localActiveLLM || "deepseek", - path: dirHandleRef.current ? dirHandleRef.current.name : null - })); - } - }, []); - - const handleSelectFolder = useCallback(async (directoryHandle) => { - if (!window.showDirectoryPicker) return; + setIsProcessing(true); + setChatHistory((prev) => [...prev, { isUser: true, text }]); try { - dirHandleRef.current = directoryHandle; - setSelectedFolder(directoryHandle.name); - - setThinkingProcess((prev) => [...prev, { - type: "user", - message: `Selected local folder: ${directoryHandle.name}` + const response = await chatWithAI(sessionIdRef.current, text, localActiveLLM || "gemini"); + setChatHistory((prev) => [...prev, { + isUser: false, + isPureAnswer: true, + text: response.answer, + provider: response.provider_used }]); - setThinkingProcess((prev) => [...prev, { - type: "system", - message: `I'm ready to help with your project! What would you like to do next? (e.g., refactor code, get an overview, etc.)` - }]); + fetchTokenUsage(); } catch (error) { - console.error("Folder selection canceled or failed:", error); + setErrorMessage(error.message); + setShowErrorModal(true); + } finally { + setIsProcessing(false); } - }, []); + }, [isConfigured, localActiveLLM, fetchTokenUsage]); const handleSwitchSession = useCallback(async (targetSessionId) => { localStorage.setItem("sessionId_coding_assistant", targetSessionId); setSessionId(targetSessionId); sessionIdRef.current = targetSessionId; - setChatHistory([]); - setThinkingProcess([{ type: "system", message: `Loading session #${targetSessionId}...` }]); try { - try { - const { getSession } = require("../services/apiService"); - const sessionInfo = await getSession(targetSessionId); - if (sessionInfo && sessionInfo.provider_name) { - setLocalActiveLLM(sessionInfo.provider_name); - } - } catch (e) { console.warn("Failed to get switched session provider info", e); } - - const messagesData = await getSessionMessages(targetSessionId); - if (messagesData && messagesData.messages) { - const mappedHistory = messagesData.messages.map(msg => ({ - text: msg.content, - isUser: msg.sender === 'user' - })); - setChatHistory(mappedHistory); + const sessionInfo = await getSession(targetSessionId); + if (sessionInfo && sessionInfo.provider_name) { + setLocalActiveLLM(sessionInfo.provider_name); } - fetchTokenUsage(); - setThinkingProcess([{ type: "system", message: `Successfully switched to session #${targetSessionId}.` }]); + await fetchSessionHistory(targetSessionId); + await fetchTokenUsage(); } catch (error) { console.error("Failed to switch session:", error); - setThinkingProcess([{ type: "system", message: "Failed to load session history." }]); } - }, [fetchTokenUsage]); - - const handlePause = useCallback(() => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - ws.current.send(JSON.stringify({ - type: "control", - command: "pause", - session_id: sessionIdRef.current - })); - } - }, []); - - const handleStop = useCallback(() => { - if (ws.current && ws.current.readyState === WebSocket.OPEN) { - ws.current.send(JSON.stringify({ - type: "control", - command: "stop", - session_id: sessionIdRef.current - })); - } - }, []); + }, [fetchSessionHistory, fetchTokenUsage]); return { chatHistory, - thinkingProcess, - selectedFolder, - connectionStatus, isProcessing, - isPaused, errorMessage, showErrorModal, tokenUsage, isConfigured, missingConfigs, handleSendChat, - handleSelectFolder, - handlePause, - handleStop, setShowErrorModal, handleSwitchSession, sessionId, diff --git a/ui/client-app/src/index.js b/ui/client-app/src/index.js index d563c0f..111004f 100644 --- a/ui/client-app/src/index.js +++ b/ui/client-app/src/index.js @@ -1,4 +1,4 @@ -import React from 'react'; +import React from 'react'; // version 1.0.1 import ReactDOM from 'react-dom/client'; import './index.css'; import App from './App'; diff --git a/ui/client-app/src/pages/CodingAssistantPage.js b/ui/client-app/src/pages/CodingAssistantPage.js index f114309..eed9113 100644 --- a/ui/client-app/src/pages/CodingAssistantPage.js +++ b/ui/client-app/src/pages/CodingAssistantPage.js @@ -1,34 +1,19 @@ -import React, { useState, useRef, useEffect, useCallback } from "react"; -// Import the components you'll create for each section +import React, { useState, useRef, useEffect } from "react"; import ChatArea from "../components/ChatArea"; -import CodeFolderAccess from "../components/CodeFolderAccess"; -import InteractionLog from "../components/InteractionLog"; -// import Controls from "../components/Controls"; import SessionSidebar from "../components/SessionSidebar"; - -// A custom hook to manage WebSocket connection and state import useCodeAssistant from "../hooks/useCodeAssistant"; import { updateSession } from "../services/apiService"; const CodeAssistantPage = () => { - // Reference for the main container to manage scrolling const pageContainerRef = useRef(null); - // Use a custom hook to handle the core logic const { chatHistory, - thinkingProcess, - connectionStatus, - selectedFolder, isProcessing, - isPaused, errorMessage, showErrorModal, tokenUsage, handleSendChat, - handleSelectFolder, - handlePause, - handleStop, setShowErrorModal, handleSwitchSession, sessionId, @@ -39,7 +24,6 @@ missingConfigs } = useCodeAssistant({ pageContainerRef }); - const [isPanelExpanded, setIsPanelExpanded] = useState(false); const [showConfigModal, setShowConfigModal] = useState(false); const [sidebarRefreshTick, setSidebarRefreshTick] = useState(0); @@ -55,12 +39,11 @@ } }; - // Scroll to the bottom of the page when new content is added useEffect(() => { if (pageContainerRef.current) { pageContainerRef.current.scrollTop = pageContainerRef.current.scrollHeight; } - }, [chatHistory, thinkingProcess]); + }, [chatHistory]); return (
@@ -74,128 +57,94 @@ {/* Main content area */}
- {/* Adjusted grid for narrow right panel */} -
- {/* Left Column: Chat Area */} -
- {/* Area 1: Chat with LLM */} -
-

-
- Chat with the LLM -
- {!isConfigured && ( -
- - - -
-

Missing Key

-
    - {missingConfigs?.map((m, i) =>
  • {m}
    • )} -
-
-
- )} - +
+ {/* Chat Area */} +
+
+

+
+
+ + +
+
+ Coding Assistant + Enhanced with Skills & Prompts +
+ {!isConfigured && ( +
+ + + +
+

Missing Key

+
    + {missingConfigs?.map((m, i) =>
  • {m}
    • )} +
+
+
+ )} +
-
+
-
- Context Usage +
+ Token Usage
-
+
80 ? 'bg-red-500' : 'bg-green-500'}`} + className={`h-full transition-all duration-700 ease-out ${tokenUsage?.percentage > 80 ? 'bg-red-500' : 'bg-indigo-500'}`} style={{ width: `${Math.min(tokenUsage?.percentage || 0, 100)}%` }} >
-
80 ? 'text-red-500' : 'text-gray-400'}`}> + 80 ? 'text-red-500' : 'text-gray-400'}`}> {tokenUsage?.percentage || 0}% -
+
- - Unlock powerful coding assistance with RAG! - -

- {/* Note: ChatArea component needs to be implemented with a