diff --git a/.gitignore b/.gitignore
index 176d2c0..d7df9e9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,4 @@
**.bin
**.db
ai-hub/data/*
-.vscode/
+
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000..011873d
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,31 @@
+{
+ // Use IntelliSense to learn about possible attributes.
+ // Hover to view descriptions of existing attributes.
+ // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+ "version": "0.2.0",
+ "configurations": [
+ {
+ "name": "Python Debugger: Current File",
+ "type": "debugpy",
+ "request": "launch",
+ "program": "${file}",
+ "console": "integratedTerminal"
+ },
+ {
+ "name": "Python: Uvicorn",
+ "type": "python",
+ "request": "launch",
+ "module": "uvicorn",
+ "args": [
+ "app.main:app",
+ "--host",
+ "127.0.0.1",
+ "--port",
+ "8001",
+ "--reload"
+ ],
+ "cwd": "${workspaceFolder}/ai-hub",
+ "envFile": ".env"
+ }
+ ]
+}
\ No newline at end of file
diff --git a/ai-hub/app/core/pipelines/file_selector.py b/ai-hub/app/core/pipelines/file_selector.py
new file mode 100644
index 0000000..0c421a4
--- /dev/null
+++ b/ai-hub/app/core/pipelines/file_selector.py
@@ -0,0 +1,46 @@
+import dspy
+import json
+from typing import List
+from app.db import models
+
+# Assuming SelectFiles and other necessary imports are defined as in the previous example
+
+class SelectFiles(dspy.Signature):
+ """
+ Based on the user's question, communication history, and the code folder's file list, identify the files that are most relevant to answer the question.
+ """
+ question = dspy.InputField(desc="The user's current question.")
+ chat_history = dspy.InputField(desc="The ongoing dialogue between the user and the AI.")
+ code_folder_filename_list = dspy.InputField(desc="A JSON array of objects. Each object represents a file and contains its name, path, size, last modified timestamp, and created timestamp.")
+ answer = dspy.OutputField(format=list, desc="A list of strings containing the names of the most relevant files to examine further.")
+
+class CodeRagFileSelector(dspy.Module):
+ """
+ A single-step module to select relevant files from a list based on a user question.
+ """
+ def __init__(self):
+ super().__init__()
+ self.select_files = dspy.Predict(SelectFiles)
+
+ async def forward(self, question: str, history: List[models.Message], file_list: List[dict]) -> List[str]:
+ # Format history for the signature
+ history_text = self._default_history_formatter(history)
+
+ # Convert the list of dictionaries to a JSON string
+ file_list_json_string = json.dumps(file_list, indent=2)
+
+ # Call the predictor with the necessary inputs
+ prediction = await self.select_files.acall(
+ question=question,
+ chat_history=history_text,
+ code_folder_filename_list=file_list_json_string
+ )
+
+ # The output is expected to be a list of strings
+ return prediction.answer
+
+ 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
+ )
\ No newline at end of file
diff --git a/ai-hub/app/core/pipelines/file_selector_rag.py b/ai-hub/app/core/pipelines/file_selector_rag.py
deleted file mode 100644
index fb28ff6..0000000
--- a/ai-hub/app/core/pipelines/file_selector_rag.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import dspy
-from typing import List
-from app.db import models
-
-# Assuming SelectFiles and other necessary imports are defined as in the previous example
-
-class SelectFiles(dspy.Signature):
- """
- Based on the user's question, communication history, and the code folder's file list, identify the files that are most relevant to answer the question.
- """
- question = dspy.InputField(desc="The user's current question.")
- chat_history = dspy.InputField(desc="The ongoing dialogue between the user and the AI.")
- code_folder_filename_list = dspy.InputField(desc="A list of file names as strings, representing the file structure of the code base.")
- answer = dspy.OutputField(format=list, desc="A list of strings containing the names of the most relevant files to examine further.")
-
-class CodeRagFileSelector(dspy.Module):
- """
- A single-step module to select relevant files from a list based on a user question.
- """
- def __init__(self):
- super().__init__()
- self.select_files = dspy.Predict(SelectFiles)
-
- async def forward(self, question: str, history: List[models.Message], file_list: List[str]) -> List[str]:
- # Format history for the signature
- history_text = self._default_history_formatter(history)
-
- # Call the predictor with the necessary inputs
- prediction = await self.select_files(
- question=question,
- chat_history=history_text,
- code_folder_filename_list="\n".join(file_list)
- )
-
- # The output is expected to be a list of strings
- return prediction.answer
-
- 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
- )
\ No newline at end of file
diff --git a/ai-hub/app/core/services/workspace.py b/ai-hub/app/core/services/workspace.py
index 907d9a8..8ca08b5 100644
--- a/ai-hub/app/core/services/workspace.py
+++ b/ai-hub/app/core/services/workspace.py
@@ -1,8 +1,11 @@
-import asyncio
+import dspy
import json
-from typing import Dict, Any, Callable, Awaitable
+import uuid
+import ast # Import the Abstract Syntax Trees module
+from typing import Dict, Any, Callable, Awaitable, List
from fastapi import WebSocket
-from sqlalchemy.orm import Session
+from app.core.providers.factory import get_llm_provider
+from app.core.pipelines.file_selector import CodeRagFileSelector
# A type hint for our handler functions
MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]]
@@ -15,18 +18,60 @@
def __init__(self):
# The dispatcher map: keys are message types, values are handler functions
self.message_handlers: Dict[str, MessageHandler] = {
- "select_folder": self.handle_select_folder_response,
+ "select_folder_response": self.handle_select_folder_response,
"list_directory_response": self.handle_list_directory_response,
- "file_content_response": self.handle_file_content_response,
+ "file_content_response": self.handle_files_content_response,
"execute_command_response": self.handle_command_output,
# Add more message types here as needed
}
+ # Centralized map of commands that can be sent to the client
+ self.command_map: Dict[str, Dict[str, Any]] = {
+ "list_directory": {"type": "list_directory", "description": "Request a list of files and folders in the current directory."},
+ "get_file_content": {"type": "get_file_content", "description": "Request the content of a specific file."},
+ "execute_command": {"type": "execute_command", "description": "Request to execute a shell command."},
+ # Define more commands here
+ }
+ # Per-websocket session state management
+ self.sessions: Dict[str, Dict[str, Any]] = {}
+
+ def generate_request_id(self) -> str:
+ """Generates a unique request ID."""
+ return str(uuid.uuid4())
+
+ async def send_command(self, websocket: WebSocket, command_name: str, data: Dict[str, Any] = {}):
+ """Helper to send a command to the client with a unique request_id and round number."""
+ if command_name not in self.command_map:
+ raise ValueError(f"Unknown command: {command_name}")
+
+ request_id = self.generate_request_id()
+ session_state = self.sessions.get(websocket.scope["client"], {"round": 0})
+ session_state["round"] += 1
+
+ message_to_send = {
+ "type": self.command_map[command_name]["type"],
+ "request_id": request_id,
+ "round": session_state["round"],
+ **data,
+ }
+
+ await websocket.send_text(json.dumps(message_to_send))
+ print(f"Sent command '{command_name}' to client (request_id: {request_id}, round: {session_state['round']})")
+
+ self.sessions[websocket.scope["client"]] = session_state
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})")
+
handler = self.message_handlers.get(message_type)
if handler:
await handler(websocket, message)
@@ -38,40 +83,80 @@
"""Handles the client's response to a select folder response."""
path = data.get("path")
request_id = data.get("request_id")
+
print(f"Received folder selected (request_id: {request_id}): Path: {path}")
- # After the server received the request that folder selected, we immediately ask for the file lists in the folder.
- await websocket.send_text(json.dumps({
- "type": "list_directory",
- "request_id": request_id
- }))
-
+
+ # 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})
async def handle_list_directory_response(self, websocket: WebSocket, data: Dict[str, Any]):
"""Handles the client's response to a list_directory request."""
- # This is where the AI logic would pick up after getting the file list
files = data.get("files", [])
- folders = data.get("folders", [])
- request_id = data.get("request_id")
- await websocket.send_text(json.dumps({
- "type": "list",
- "content": f"Analyzing the content of file: {files[0]}"
- }))
- print(f"Received directory listing (request_id: {request_id}): Files: {files}, Folders: {folders}")
+ provider_name = data.get("provider_name", "gemini")
+ llm_provider = get_llm_provider(provider_name)
+ cfs = CodeRagFileSelector()
+
+ with dspy.context(lm=llm_provider):
+ raw_answer_text = await cfs(
+ question="Please help to refactor my code",
+ # The history will be retrieved from a database in a real application
+ # For this example, we'll pass an empty history
+ history="",
+ file_list=files
+ )
+
+ try:
+ 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:
+ # Handle cases where the LLM output is not a valid list string.
+ print(f"Error parsing LLM output: {e}")
+ answer_text = [] # Default to an empty list to prevent errors.
+ await websocket.send_text(json.dumps({
+ "type": "thinking_log",
+ "content": f"Warning: AI's file list could not be parsed. Error: {e}"
+ }))
+ return
- async def handle_file_content_response(self, websocket: WebSocket, data: Dict[str, Any]):
- """Handles the content of a file sent by the client."""
- filename = data.get("filename")
- content = data.get("content")
- request_id = data.get("request_id")
-
- print(f"Received content for '{filename}' (request_id: {request_id}). Content length: {len(content)}")
-
await websocket.send_text(json.dumps({
"type": "thinking_log",
- "content": f"Analyzing the content of file: {filename}"
+ "content": f"AI selected files: {answer_text}. Now requesting file content."
}))
+ # After getting the AI's selected files, we send a command to the client to get their content.
+ await self.send_command(websocket, "get_file_content", data={"filenames": answer_text})
+
+ 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."""
+ # The client is expected to send a list of file objects
+ # Each object should have 'filename' and 'content' keys.
+ files_data: List[Dict[str, str]] = data.get("files", [])
+ request_id = data.get("request_id")
+
+ if not files_data:
+ print(f"Warning: No files data received for request_id: {request_id}")
+ return
+
+ print(f"Received content for {len(files_data)} files (request_id: {request_id}).")
+
+ for file_info in files_data:
+ filename = file_info.get("filename")
+ content = file_info.get("content")
+
+ if filename and content:
+ print(f"Processing content for '{filename}'. Content length: {len(content)}")
+ # The AI would analyze this content to determine the next action, e.g.,
+ # generate a plan, perform a refactoring, or ask for more information.
+ await websocket.send_text(json.dumps({
+ "type": "thinking_log",
+ "content": f"Analyzing the content of file: {filename}"
+ }))
+ else:
+ print(f"Warning: Malformed file data in response for request_id: {request_id}")
+
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")
@@ -84,5 +169,4 @@
await websocket.send_text(json.dumps({
"type": "thinking_log",
"content": f"Command '{command}' completed. Analyzing output."
- }))
-
+ }))
\ No newline at end of file
diff --git a/ui/client-app/src/components/InteractionLog.js b/ui/client-app/src/components/InteractionLog.js
index 4057b18..fa2e232 100644
--- a/ui/client-app/src/components/InteractionLog.js
+++ b/ui/client-app/src/components/InteractionLog.js
@@ -14,12 +14,15 @@
? "bg-blue-100 dark:bg-blue-900 text-blue-900 dark:text-blue-100"
: log.type === "user"
? "bg-green-100 dark:bg-green-900 text-green-900 dark:text-green-100"
- : log.type === "system"
+ : log.type === "remote"
? "bg-yellow-100 dark:bg-yellow-900 text-yellow-900 dark:text-yellow-100"
+ : log.type === "local"
+ ? "bg-purple-100 dark:bg-purple-900 text-purple-900 dark:text-purple-100"
: "bg-gray-200 dark:bg-gray-700 text-gray-900 dark:text-gray-100"
}`}
>
+ [{log.round}]
{log.type.charAt(0).toUpperCase() + log.type.slice(1)}:
diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js
index 2c293ab..b878e83 100644
--- a/ui/client-app/src/hooks/useCodeAssistant.js
+++ b/ui/client-app/src/hooks/useCodeAssistant.js
@@ -1,6 +1,7 @@
// src/hooks/useCodeAssistant.js
import { useState, useEffect, useRef, useCallback } from "react";
import { connectToWebSocket } from "../services/websocket";
+import { v4 as uuidv4 } from 'uuid';
const useCodeAssistant = ({ pageContainerRef }) => {
// State variables for the assistant's UI and status
@@ -9,7 +10,7 @@
const [selectedFolder, setSelectedFolder] = useState(null);
const [connectionStatus, setConnectionStatus] = useState("disconnected");
const [isProcessing, setIsProcessing] = useState(false);
- const [isPaused, setIsPaused] = useState(false);
+ const [isPaused, setIsPaused] = useState(false); // Corrected this line
const [errorMessage, setErrorMessage] = useState("");
const [showErrorModal, setShowErrorModal] = useState(false);
const [sessionId, setSessionId] = useState(null);
@@ -28,7 +29,7 @@
const handleThinkingLog = useCallback((message) => {
setThinkingProcess((prev) => [
...prev,
- { type: message.subtype, message: message.content },
+ { type: "remote", message: message.content, round: message.round },
]);
}, []);
@@ -40,50 +41,63 @@
const handleStatusUpdate = useCallback((message) => {
setIsProcessing(message.processing);
- setIsPaused(message.paused);
+ setIsPaused(message.paused); // Corrected this line
setConnectionStatus(message.status);
}, []);
const handleListDirectoryRequest = useCallback(async (message) => {
+ const { request_id, round } = 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 }));
+ ws.current.send(JSON.stringify({ type: "error", content: errorMsg, request_id }));
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 [name, entry] of handle.entries()) {
- const relativePath = path ? `${path}/${name}` : name;
-
+ async function walkDirectory(handle, path = '') {
+ for await (const entry of handle.values()) {
+ const entryPath = `${path}/${entry.name}`;
if (entry.kind === "file") {
- files.push(relativePath); // Store full relative path
+ const file = await entry.getFile();
+ files.push({
+ name: file.name,
+ path: entryPath,
+ size: file.size,
+ lastModified: file.lastModified,
+ });
} else if (entry.kind === "directory") {
- await walkDirectory(entry, relativePath); // Recurse into subdirectory
+ await walkDirectory(entry, entryPath); // Recurse into subdirectory
}
}
}
-
+
await walkDirectory(dirHandle);
-
+
ws.current.send(
JSON.stringify({
type: "list_directory_response",
files,
- request_id: message.request_id,
+ request_id,
+ round,
})
);
-
+
setThinkingProcess((prev) => [
...prev,
{
- type: "system",
- message: `Sent list of file names from folder "${dirHandle.name}" to server. Total files: ${files.length}`,
+ type: "local",
+ message: `Sent list of files (${files.length}) to server.`,
+ round,
},
]);
} catch (error) {
@@ -92,63 +106,77 @@
JSON.stringify({
type: "error",
content: "Failed to access folder contents.",
- request_id: message.request_id,
+ request_id,
+ round,
})
);
}
}, []);
-
- const handleReadFileRequest = useCallback(async (message) => {
+ const handleReadFilesRequest = useCallback(async (message) => {
+ console.log(message);
+ const { filenames, request_id, round } = message;
const dirHandle = dirHandleRef.current;
- const { filename, request_id } = message;
-
if (!dirHandle) {
- ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id }));
+ ws.current.send(JSON.stringify({ type: "error", content: "No folder selected.", request_id, round }));
return;
}
- try {
- const fileHandle = await dirHandle.getFileHandle(filename);
- const file = await fileHandle.getFile();
- const content = await file.text();
+ setThinkingProcess((prev) => [
+ ...prev,
+ { type: "local", message: `Reading content of ${filenames.length} files...`, round },
+ ]);
- ws.current.send(JSON.stringify({
- type: "file_content_response",
- filename,
- content,
- request_id,
- }));
-
- setThinkingProcess((prev) => [
- ...prev,
- { type: "system", message: `Sent content of file "${filename}" to server.` },
- ]);
- } catch (error) {
- console.error(`Failed to read file ${filename}:`, error);
- ws.current.send(JSON.stringify({
- type: "error",
- content: `Could not read file: ${filename}`,
- request_id,
- }));
+ const filesData = [];
+ for (const filename of filenames) {
+ try {
+ const fileHandle = await dirHandle.getFileHandle(filename);
+ const file = await fileHandle.getFile();
+ const content = await file.text();
+ filesData.push({ filename, content });
+ setThinkingProcess((prev) => [
+ ...prev,
+ { type: "local", message: `Read file: ${filename}`, round },
+ ]);
+ } catch (error) {
+ console.error(`Failed to read file ${filename}:`, error);
+ ws.current.send(JSON.stringify({
+ type: "error",
+ content: `Could not read file: ${filename}`,
+ request_id,
+ round,
+ }));
+ }
}
+
+ ws.current.send(JSON.stringify({
+ type: "file_content_response",
+ files: filesData,
+ request_id,
+ round,
+ }));
+
+ setThinkingProcess((prev) => [
+ ...prev,
+ { type: "local", message: `Sent content for ${filesData.length} files to server.`, round },
+ ]);
}, []);
const handleExecuteCommandRequest = useCallback((message) => {
- const { command, request_id } = message;
-
+ const { command, request_id, round } = message;
const output = `Simulated output for command: '${command}'`;
ws.current.send(JSON.stringify({
- type: "command_output",
+ type: "execute_command_response",
command,
output,
request_id,
+ round,
}));
setThinkingProcess((prev) => [
...prev,
- { type: "system", message: `Simulated execution of command: '${command}'` },
+ { type: "system", message: `Simulated execution of command: '${command}'`, round },
]);
}, []);
@@ -170,8 +198,8 @@
case "list_directory":
handleListDirectoryRequest(message);
break;
- case "read_file":
- handleReadFileRequest(message);
+ case "get_file_content":
+ handleReadFilesRequest(message);
break;
case "execute_command":
handleExecuteCommandRequest(message);
@@ -179,7 +207,7 @@
default:
console.log("Unknown message type:", message);
}
- }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFileRequest, handleExecuteCommandRequest]);
+ }, [handleChatMessage, handleThinkingLog, handleError, handleStatusUpdate, handleListDirectoryRequest, handleReadFilesRequest, handleExecuteCommandRequest]);
// --- WebSocket Connection Setup ---
useEffect(() => {
@@ -229,7 +257,6 @@
// Open folder picker and store handle
const handleSelectFolder = useCallback(async (directoryHandle) => {
if (!window.showDirectoryPicker) {
- // Don't use window.alert, use a custom modal
return;
}
@@ -237,7 +264,9 @@
dirHandleRef.current = directoryHandle;
setSelectedFolder(directoryHandle.name);
- ws.current.send(JSON.stringify({ type: "select_folder", path: 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 }));
setThinkingProcess((prev) => [
...prev,
@@ -278,4 +307,4 @@
};
};
-export default useCodeAssistant;
+export default useCodeAssistant;
\ No newline at end of file
diff --git a/ui/run_web.sh b/ui/run_web.sh
index 896415e..901ce81 100644
--- a/ui/run_web.sh
+++ b/ui/run_web.sh
@@ -76,7 +76,7 @@
concurrently \
--prefix "[{name}]" \
--names "aihub,tts-frontend" \
- "LOG_LEVEL=DEBUG uvicorn $APP_MODULE --host $AI_HUB_HOST --log-level debug --port $AI_HUB_PORT $SSL_ARGS" \
+ "LOG_LEVEL=DEBUG uvicorn $APP_MODULE --host $AI_HUB_HOST --log-level debug --port $AI_HUB_PORT $SSL_ARGS --reload" \
"cd $TTS_CLIENT_DIR && $FRONTEND_ENV HOST=0.0.0.0 PORT=8000 npm start"
popd > /dev/null