diff --git a/ai-hub/app/core/pipelines/question_decider.py b/ai-hub/app/core/pipelines/question_decider.py index 9ec1e2b..f5db8e2 100644 --- a/ai-hub/app/core/pipelines/question_decider.py +++ b/ai-hub/app/core/pipelines/question_decider.py @@ -1,8 +1,9 @@ import dspy import json import os +from app.db import models from typing import List, Dict, Any, Tuple, Optional -from datetime import datetime +from typing import List, Callable, Optional class QuestionDecider(dspy.Signature): """ @@ -76,16 +77,24 @@ class CodeRagQuestionDecider(dspy.Module): - def __init__(self, log_dir: str = "ai_payloads"): + def __init__(self, log_dir: str = "ai_payloads", history_formatter: Optional[Callable[[List[models.Message]], str]] = None): super().__init__() self.log_dir = log_dir # Initializes the dspy Predict module with the refined system prompt self.decider = dspy.ChainOfThought(QuestionDecider) + self.history_formatter = history_formatter or self._default_history_formatter + + def _default_history_formatter(self, history: List[models.Message]) -> str: + return "\n".join( + f"{'Human' if msg.sender == 'user' else 'Assistant'}: {msg.content}" + for msg in history + ) + async def forward( self, question: str, - history: List[str], + history: List[models.Message], retrieved_data: Dict[str, Any] ) -> Tuple[str, str, str]: """ @@ -99,7 +108,6 @@ Returns: A tuple of (answer, decision, code_diff). """ - history_text = "\n".join(history) # --- INTERNAL LOGIC TO SPLIT DATA, WITH NULL/POINTER CHECKS --- with_content = [] @@ -128,6 +136,7 @@ retrieved_with_content_json = json.dumps(with_content, indent=2) retrieved_without_content_json = json.dumps(without_content, indent=2) + history_text = self.history_formatter(history) input_payload = { "question": question, "chat_history": history_text, diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py index e3be4d9..2a2cf23 100644 --- a/ai-hub/app/core/services/workspace.py +++ b/ai-hub/app/core/services/workspace.py @@ -268,7 +268,8 @@ if retrievedFile and retrievedFile.content: return retrievedFile.content else: - raise ValueError(f"File with path {file_path} not found for request ID {request_id} or has no content.") + 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]]: """ @@ -351,6 +352,7 @@ 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. @@ -362,6 +364,15 @@ Returns: The new content with the diff applied. """ + # Handle the case where the original content is empty. + if not original_content: + # If the original content is empty, just add the new lines from the diff. + new_content: List[str] = [] + for line in file_diff.splitlines(keepends=True): + if 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) @@ -526,13 +537,28 @@ """ files_data: List[Dict[str, str]] = data.get("files", []) request_id = data.get("request_id") + session_id = data.get("session_id") 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) - + + 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 request_id: + await websocket.send_text(json.dumps({ + "type": "error", + "content": "Error: request_id is required to process file content." + })) + return + # Retrieve the updated context from the database context_data = await self._retrieve_by_request_id(self.db, request_id=request_id) @@ -547,15 +573,26 @@ "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 with dspy.context(lm=get_llm_provider(provider_name="gemini")): crqd = CodeRagQuestionDecider() raw_answer_text, reasoning, decision, code_diff = await crqd( question=context_data.get("question", ""), - history="", + history=session.messages, retrieved_data=context_data ) dspy.inspect_history(n=1) # Inspect the last DSPy operation for debugging + if decision in [ "code_change", "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) + if decision == "files": await websocket.send_text(json.dumps({ "type": "thinking_log", diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index f8fcb70..2b1db06 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -3,7 +3,6 @@ import { v4 as uuidv4 } from 'uuid'; const useCodeAssistant = ({ pageContainerRef }) => { - // State variables for the assistant's UI and status const [chatHistory, setChatHistory] = useState([]); const [thinkingProcess, setThinkingProcess] = useState([]); const [selectedFolder, setSelectedFolder] = useState(null); @@ -14,23 +13,29 @@ const [showErrorModal, setShowErrorModal] = useState(false); const [sessionId, setSessionId] = useState(null); - // Refs for the WebSocket connection and directory handle + const sessionIdRef = useRef(null); // ✅ Always current sessionId const ws = useRef(null); const initialized = useRef(false); const dirHandleRef = useRef(null); const handleChatMessage = useCallback((message) => { console.log("Received chat message:", message); - // Update chat history with the formatted content - setChatHistory((prev) => [...prev, { isUser: false, text: message.content, dicision: message.dicision, code_diff: message.code_diff, reasoning: message.reasoning }]); + setChatHistory((prev) => [...prev, { + isUser: false, + text: message.content, + dicision: message.dicision, + code_diff: message.code_diff, + reasoning: message.reasoning + }]); setIsProcessing(false); }, []); const handleThinkingLog = useCallback((message) => { - setThinkingProcess((prev) => [ - ...prev, - { type: "remote", message: message.content, round: message.round }, - ]); + setThinkingProcess((prev) => [...prev, { + type: "remote", + message: message.content, + round: message.round + }]); }, []); const handleError = useCallback((message) => { @@ -49,6 +54,7 @@ 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); @@ -56,15 +62,9 @@ return; } - // setThinkingProcess((prev) => [ - // ...prev, - // { type: "system", message: `Scanning directory...`, round }, - // ]); - try { const files = []; - // Recursive function to walk through directories async function walkDirectory(handle, path = '') { for await (const entry of handle.values()) { const entryPath = `${path}/${entry.name}`; @@ -78,111 +78,101 @@ created: file.created, }); } else if (entry.kind === "directory") { - await walkDirectory(entry, entryPath); // Recurse into subdirectory + await walkDirectory(entry, entryPath); } } } await walkDirectory(dirHandle); - ws.current.send( - JSON.stringify({ - type: "list_directory_response", - files, - request_id: request_id, - session_id: sessionId, - }) - ); + 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.`, - }, - ]); + 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: sessionId, - }) - ); + ws.current.send(JSON.stringify({ + type: "error", + content: "Failed to access folder contents.", + request_id, + session_id: sessionIdRef.current, + })); } - }, [sessionId]); + }, []); - // Helper function to recursively get a file handle from a full path const getFileHandleFromPath = async (dirHandle, filePath) => { - // Split the path and filter out any empty parts to prevent the error 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) { - // Last part is the file name - return await currentHandle.getFileHandle(part); - } else { - // Part is a directory - currentHandle = await currentHandle.getDirectoryHandle(part); - } - } catch (error) { - console.error(`Error navigating to path part '${part}':`, error); - // Return null or re-throw based on desired error handling - return null; + 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(`Error navigating to path part '${part}':`, 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: sessionId })); + ws.current.send(JSON.stringify({ + type: "error", + content: "No folder selected.", + request_id, + session_id: sessionIdRef.current, + })); return; } - + const filesData = []; - const readFiles = []; // Array to store names of successfully read files - - // setThinkingProcess((prev) => [...prev, { type: "local", message: `Reading content of ${filepaths.length} files...` }]); - + const readFiles = []; + for (const filepath of filepaths) { try { const fileHandle = await getFileHandleFromPath(dirHandle, filepath); const file = await fileHandle.getFile(); const content = await file.text(); filesData.push({ filepath, content }); - readFiles.push(filepath); // Add successfully read file to the list + readFiles.push(filepath); } catch (error) { console.error(`Failed to read file ${filepath}:`, error); ws.current.send(JSON.stringify({ type: "error", content: `Could not read file: ${filepath}`, - request_id: request_id, - session_id: sessionId, + request_id, + session_id: sessionIdRef.current, })); } } - + ws.current.send(JSON.stringify({ type: "file_content_response", files: filesData, - request_id: request_id, - session_id: sessionId, + request_id, + session_id: sessionIdRef.current, })); - - // Consolidate all thinking process updates into a single call + setThinkingProcess((prev) => { const newMessages = []; - - // Add a message summarizing the files that were read + if (readFiles.length > 0) { const displayMessage = readFiles.length > 10 ? `Read ${readFiles.length} files successfully.` @@ -191,7 +181,7 @@ } return [...prev, ...newMessages]; }); - }, [sessionId]); + }, []); const handleExecuteCommandRequest = useCallback((message) => { const { command, request_id } = message; @@ -202,16 +192,15 @@ command, output, request_id, - session_id: sessionId, + session_id: sessionIdRef.current, })); - setThinkingProcess((prev) => [ - ...prev, - { type: "system", message: `Simulated execution of command: '${command}'` }, - ]); - }, [sessionId]); + setThinkingProcess((prev) => [...prev, { + type: "system", + message: `Simulated execution of command: '${command}'` + }]); + }, []); - // Main message handler that routes messages to the correct function const handleIncomingMessage = useCallback((message) => { switch (message.type) { case "chat_message": @@ -238,9 +227,17 @@ default: console.log("Unknown message type:", message); } - }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFilesRequest, handleExecuteCommandRequest]); + }, [ + handleChatMessage, + handleThinkingLog, + handleError, + handleStatusUpdate, + handleListDirectoryRequest, + handleReadFilesRequest, + handleExecuteCommandRequest + ]); - // --- WebSocket Connection Setup --- + // WebSocket Setup useEffect(() => { if (initialized.current) return; initialized.current = true; @@ -262,6 +259,7 @@ ); ws.current = newWs; setSessionId(newSessionId); + sessionIdRef.current = newSessionId; // ✅ Keep ref in sync } catch (error) { console.error("Setup failed:", error); } @@ -276,51 +274,54 @@ }; }, [handleIncomingMessage]); - // Send chat message to server const handleSendChat = useCallback(async (text) => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { setChatHistory((prev) => [...prev, { isUser: true, text }]); setIsProcessing(true); - // Removed the extra call to getSessionId as it is already in state - ws.current.send(JSON.stringify({ type: "chat_message",content: text, session_id: sessionId, path: dirHandleRef .current ? dirHandleRef.current.name : null })); + ws.current.send(JSON.stringify({ + type: "chat_message", + content: text, + session_id: sessionIdRef.current, + path: dirHandleRef.current ? dirHandleRef.current.name : null + })); } - }, [sessionId]); + }, []); - // Open folder picker and store handle const handleSelectFolder = useCallback(async (directoryHandle) => { - if (!window.showDirectoryPicker) { - return; - } + if (!window.showDirectoryPicker) return; try { dirHandleRef.current = directoryHandle; setSelectedFolder(directoryHandle.name); - // Send the initial message to the server with a unique request_id - // const request_id = uuidv4(); - // ws.current.send(JSON.stringify({ type: "select_folder_response", path: directoryHandle.name, request_id, session_id: sessionId })); - - setThinkingProcess((prev) => [ - ...prev, - { type: "user", message: `Selected local folder: ${directoryHandle.name}` }, - ]); + setThinkingProcess((prev) => [...prev, { + type: "user", + message: `Selected local folder: ${directoryHandle.name}` + }]); } catch (error) { console.error("Folder selection canceled or failed:", error); } - }, [sessionId]); + }, []); - // Control functions const handlePause = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { - ws.current.send(JSON.stringify({ type: "control", command: "pause", session_id: sessionId })); + ws.current.send(JSON.stringify({ + type: "control", + command: "pause", + session_id: sessionIdRef.current + })); } - }, [sessionId]); + }, []); const handleStop = useCallback(() => { if (ws.current && ws.current.readyState === WebSocket.OPEN) { - ws.current.send(JSON.stringify({ type: "control", command: "stop", session_id: sessionId })); + ws.current.send(JSON.stringify({ + type: "control", + command: "stop", + session_id: sessionIdRef.current + })); } - }, [sessionId]); + }, []); return { chatHistory, @@ -339,4 +340,4 @@ }; }; -export default useCodeAssistant; \ No newline at end of file +export default useCodeAssistant;