diff --git a/.gitignore b/.gitignore
index d7df9e9..3faa417 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,4 @@
**.bin
**.db
ai-hub/data/*
-
+ai-hub/ai_payloads/*
diff --git a/ai-hub/app/core/pipelines/context_compressor.py b/ai-hub/app/core/pipelines/context_compressor.py
new file mode 100644
index 0000000..6563bab
--- /dev/null
+++ b/ai-hub/app/core/pipelines/context_compressor.py
@@ -0,0 +1,96 @@
+import dspy
+import os
+import json
+from datetime import datetime
+
+class SimpleDataCompressor(dspy.Signature):
+ """
+ Condense a long string of concatenated file data into a concise summary that retains the most important information related to the user's question.
+ """
+ question = dspy.InputField(desc="The user's current question.")
+ full_context_string = dspy.InputField(desc="A long string containing file paths and their contents.")
+ summarized_context = dspy.OutputField(
+ desc="A long, detailed, synthesizing the most relevant information from the input data to address the user's question effectively."
+ )
+
+# 1. Define the system prompt string
+SYSTEM_PROMPT = """
+Here is a system prompt designed to address the issues in the input data and guide the model toward a useful, non-truncated response.
+
+### System Prompt for Compression and Contextualization
+
+You are an expert AI assistant tasked with analyzing a complex request and its associated file data. Your primary goal is to provide a **concise, accurate, and actionable response** to the user's question, while handling potentially overwhelming input data.
+
+---
+**Instructions:**
+
+1. **Strictly Filter Irrelevant Files:** First, ignore all non-essential files in the `retrieved_files` array. This includes:
+ * **Compiled files:** Any file with a `.pyc` or similar extension.
+ * **Cache directories:** Any path containing `__pycache__`.
+ * **Documentation or shell scripts:** Files ending in `.md` or `.sh`.
+ * **Initialization files:** `__init__.py` files unless explicitly requested.
+
+2. **Focus on Specificity:** The user's question is "generally polish the `app.core.services.workspace.py` and its test code." This means the most relevant files are `app/core/services/workspace.py` and any file explicitly identified as its test file (e.g., `test_workspace.py`). Prioritize these. All other files are secondary context.
+
+3. **Synthesize, Don't Dump:** Instead of simply printing all the file contents, **synthesize** the most relevant information. Describe the high-level purpose of the `workspace.py` file based on its content, and identify the key functions and classes. For the test code, describe the functionality being tested.
+
+4. **Action-Oriented Response:** Based on your analysis, provide a response that directly answers the user's request to "polish" the code. This might include:
+ * Suggestions for code refactoring or simplification.
+ * Recommendations for improving a function's logic.
+ * Identifying potential bugs or edge cases.
+ * Adding or improving comments and documentation.
+ * Proposing new test cases for better coverage.
+
+5. **Final Output:** Your final response should be a clean, well-formatted markdown text. **Do not include the raw input data or the full list of files in your final output.** Your response must be an answer to the user's request, not a report on the input data itself.
+
+6. **Maintain Professional Tone:** Respond in a clear, concise, and professional tone. Avoid conversational filler or unnecessary explanations about how you processed the data.
+"""
+class StringContextCompressor(dspy.Module):
+ """
+ A pipeline to compress a string of retrieved file data.
+ """
+ def __init__(self):
+ super().__init__()
+ # Use ChainOfThought for multi-step reasoning to produce a good summary.
+ self.compressor = dspy.ChainOfThought(SimpleDataCompressor, system_prompt=SYSTEM_PROMPT)
+
+ async def forward(self, question: str, retrieved_data_string: str) -> str:
+ """
+ Processes a string of data and returns a compressed string.
+ """
+ # Call the ChainOfThought module with the string inputs.
+ input_payload = {
+ "question": question,
+ "full_context_string": retrieved_data_string
+ }
+ prediction = await self.compressor.acall(**input_payload)
+
+ self._log_to_file(input_payload, prediction)
+
+ return prediction.summarized_context
+
+ def _log_to_file(self, history_entry, prediction):
+ """
+ Saves the raw payload sent to the AI and the raw response to a local file.
+ """
+ # Create a log directory if it doesn't exist
+ log_dir = "ai_payloads"
+ os.makedirs(log_dir, exist_ok=True)
+
+ # Generate a unique filename using a timestamp
+ timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ filename = f"{log_dir}/ai_payload_{timestamp}.json"
+
+ log_data = {
+ "timestamp": timestamp,
+ "request_payload": history_entry,
+ "response_payload": {
+ "summarized_context": prediction.summarized_context,
+ }
+ }
+
+ # Save the data to a file
+ with open(filename, "w", encoding="utf-8") as f:
+ json.dump(log_data, f, indent=4)
+
+ print(f"Logged AI payload and response to {filename}")
\ No newline at end of file
diff --git a/ai-hub/app/core/pipelines/question_decider.py b/ai-hub/app/core/pipelines/question_decider.py
index d2c996f..bca7499 100644
--- a/ai-hub/app/core/pipelines/question_decider.py
+++ b/ai-hub/app/core/pipelines/question_decider.py
@@ -1,67 +1,141 @@
import dspy
import json
-from typing import List, Dict, Any
+import os
+from typing import List, Dict, Any, Tuple
+from datetime import datetime
+
class QuestionDecider(dspy.Signature):
"""
- Based on the user's question, chat history, and the content of the retrieved files, decide whether you have enough information to answer the question or if you need to request more files.
- If you have enough information, your decision should be 'answer' and you should provide a detailed answer.
- If you need more information, your decision should be 'files' and you should provide a list of file paths that are needed.
+ Based on the user's question, chat history, and the content of the retrieved files,
+ decide whether you have enough information to answer the question, or if you need to request more files.
+ If the request is to modify code, suggest changes using git diff syntax.
+
+ - If you have enough information, your decision must be 'answer' and you must provide a long, complete, and well-explained answer.
+ - If you need more information, your decision must be 'files' and you must provide a JSON-formatted list of file paths that are needed.
+ - If the request is to modify code, your decision should be 'code_change', and you must provide both a high-level suggestion and code changes in git diff format.
"""
+
question = dspy.InputField(desc="The user's current question.")
chat_history = dspy.InputField(desc="The ongoing dialogue between the user and the AI.")
- retrieved_data = dspy.InputField(desc="A JSON object containing the content of the retrieved files relevant to the question.")
+ retrieved_data = dspy.InputField(desc="A JSON string containing the content of retrieved files relevant to the question.")
+
decision = dspy.OutputField(
- desc="Your decision, which must be either 'answer' or 'files'."
+ desc="Your decision: either 'answer', 'files', or 'code_change'."
)
+
answer = dspy.OutputField(
- desc="If the decision is 'answer', this is your final, complete answer. If the decision is 'files', this is a JSON-formatted list of file paths needed to answer the question."
+ desc=(
+ "If decision is 'answer', provide a detailed, comprehensive, and well-structured response. "
+ "If decision is 'files', return a JSON-formatted list of file paths required to answer the question. "
+ "If decision is 'code_change', provide a high-level suggestion or explanation for the proposed changes. "
+ "For code changes, this summary is optional and can be brief; the git diff is the priority."
+ )
)
+ code_diff = dspy.OutputField(
+ desc=(
+ "If decision is 'code_change', provide the code modifications using a clear git diff syntax. "
+ "This output must be complete and not truncated. "
+ "For 'answer' or 'files' decisions, this field should be empty."
+ )
+ )
+
class CodeRagQuestionDecider(dspy.Module):
"""
- A pipeline to decide whether to answer a question or request more files.
+ A pipeline that uses DSPy to decide whether to answer a user's question,
+ request additional files, or suggest a code change based on the available data.
"""
+
def __init__(self):
super().__init__()
- # Assign the system prompt directly to the dspy.Predict instance
- system_prompt = "You are a helpful AI assistant. Your job is to decide whether to answer the user's question based on the provided context or to request more files. If you have enough information, respond with 'answer' and provide a detailed answer. If not, respond with 'files' and provide a JSON-formatted list of file paths needed."
+
+ # Modified system prompt to prioritize git diff output
+ # Revised system prompt for CodeRagQuestionDecider
+ system_prompt = """
+ You are a highly specialized AI assistant for a software engineering workspace. Your task is to accurately and efficiently handle user requests based on provided file information.
+
+ **Core Directives:**
+ 1. **Analyze the User's Question First:** Read the user's question and identify the specific file(s) they are asking about.
+ 2. **Filter Aggressively:** The provided `retrieved_data` is a JSON string that may contain a large number of files. **You must strictly ignore all irrelevant files.**
+ * **Irrelevant files include, but are not limited to:**
+ * Compiled Python files (`.pyc`).
+ * Documentation files (`.md`).
+ * Shell scripts (`.sh`).
+ * Files in `__pycache__` directories.
+ 3. **Prioritize Relevant Files:**
+ * Find the exact file path mentioned in the question (e.g., `app.core.services.workspace.py`).
+ * Identify any related test files (e.g., `test_app_core_services_workspace.py` or similar).
+ * **Only process the content of these highly relevant files.** The content for other files is not provided and should not be discussed.
+ 4. **Decide Your Action:**
+ * If the user's question is "polish the X.py and its test code" and you have the code for `X.py` in your context, your decision is always **'answer'**. You have enough information.
+ * If you do not have the content for the specifically requested file(s), your decision is **'files'**, and you must provide a list of the required file paths (e.g., `["/path/to/X.py", "/path/to/test_X.py"]`).
+ 5. **Format the Output Strictly:**
+ * If the decision is 'answer', the `answer` field must be a detailed response in markdown.
+ * If the decision is 'files', the `answer` field must be a **valid JSON array of strings**, representing the file paths you need.
+ * The `decision` field must be either 'answer' or 'files'.
+
+ **Your goal is to provide a correct, concise, and focused response based ONLY on the provided relevant file content.** Do not hallucinate or discuss files for which you do not have content.
+ """
self.decider = dspy.Predict(QuestionDecider, system_prompt=system_prompt)
- async def forward(self, question: str, history: List[str], retrieved_data: Dict[str, Any]) -> str | List[str]:
- # Format the chat history for the signature
- history_text = "\n".join(history)
-
- # Convert the retrieved_data to a JSON string
- retrieved_data_json_string = json.dumps(retrieved_data, indent=2)
+ async def forward(
+ self,
+ question: str,
+ history: List[str],
+ retrieved_data: Dict[str, Any]
+ ) -> Tuple[str, str, str]:
+ """
+ Decide whether to answer, request more files, or suggest a code change.
- # Use the decider to determine the next action
- prediction = await self.decider.acall(
- question=question,
- chat_history=history_text,
- retrieved_data=retrieved_data_json_string
- )
-
- # Check the decision and return the appropriate output
- if prediction.decision.lower() == 'answer':
- # The AI has enough information, so return the final answer.
- return prediction.answer
- elif prediction.decision.lower() == 'files':
- # The AI needs more files, so parse and return the list of file paths.
- try:
- # The LLM's output for files should be a JSON array string
- return json.loads(prediction.answer)
- except (json.JSONDecodeError, TypeError) as e:
- # Fallback for malformed JSON output from the LLM
- print(f"Warning: Failed to parse files list from LLM: {e}")
- print(f"Raw LLM output: {prediction.answer}")
- # Return an empty list or a list parsed from lines, depending on how you want to handle bad output.
- # Assuming the model returns a string like 'file1.py,file2.py' or '["file1.py", "file2.py"]'
- if isinstance(prediction.answer, str):
- cleaned_string = prediction.answer.strip("[]'\" \n\t")
- return [item.strip() for item in cleaned_string.split(',') if item.strip()]
- return []
- else:
- # Handle unexpected decisions from the LLM
- print(f"Warning: Unexpected decision from LLM: {prediction.decision}")
- return []
\ No newline at end of file
+ Args:
+ question (str): The user's question.
+ history (List[str]): The chat history.
+ retrieved_data (Dict[str, Any]): The content of files relevant to the question.
+
+ Returns:
+ Tuple[str, str, str]: The model's answer, decision ('answer', 'files', or 'code_change'), and code diff.
+ """
+ history_text = "\n".join(history)
+ retrieved_data_json = json.dumps(retrieved_data, indent=0)
+
+ input_payload = {
+ "question": question,
+ "chat_history": history_text,
+ "retrieved_data": retrieved_data_json
+ }
+
+ prediction = await self.decider.acall(**input_payload)
+
+ self._log_to_file(input_payload, prediction)
+
+ return prediction.answer, prediction.decision.lower(), prediction.code_diff
+
+ def _log_to_file(self, request_payload: Dict[str, Any], prediction: Any) -> None:
+ """
+ Saves the input and output of the AI call to a JSON file.
+
+ Args:
+ request_payload (Dict[str, Any]): The input sent to the AI.
+ prediction (Any): The AI's response.
+ """
+ log_dir = "ai_payloads"
+ os.makedirs(log_dir, exist_ok=True)
+
+ timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ filename = os.path.join(log_dir, f"ai_payload_{timestamp}.json")
+
+ log_data = {
+ "timestamp": timestamp,
+ "request_payload": request_payload,
+ "response_payload": {
+ "decision": prediction.decision,
+ "answer": prediction.answer,
+ "code_diff": prediction.code_diff
+ }
+ }
+
+ with open(filename, "w", encoding="utf-8") as f:
+ json.dump(log_data, f, indent=4)
+
+ print(f"[LOG] AI payload and response saved to {filename}")
\ No newline at end of file
diff --git a/ai-hub/app/core/retrievers/file_retriever.py b/ai-hub/app/core/retrievers/file_retriever.py
index a811c40..ef3a681 100644
--- a/ai-hub/app/core/retrievers/file_retriever.py
+++ b/ai-hub/app/core/retrievers/file_retriever.py
@@ -1,6 +1,6 @@
from typing import Dict, Any, Optional
from app.db import file_retriever_models
-from sqlalchemy.orm import Session,joinedload
+from sqlalchemy.orm import Session, joinedload
import uuid
class FileRetriever:
@@ -46,18 +46,27 @@
"directory_path": request.directory_path,
"session_id": request.session_id,
"created_at": request.created_at.isoformat() if request.created_at else None,
- "retrieved_files": [
- {
- "file_id": file.id,
+ "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,
- "file_name": file.file_name,
"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,
}
- for file in request.retrieved_files
- ]
- }
+ 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
\ 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 e6790a2..33e6786 100644
--- a/ai-hub/app/core/services/workspace.py
+++ b/ai-hub/app/core/services/workspace.py
@@ -1,6 +1,7 @@
import dspy
import json
import uuid
+import re
import logging
from datetime import datetime
import ast # Import the Abstract Syntax Trees module
@@ -14,6 +15,7 @@
from app.core.pipelines.file_selector import CodeRagFileSelector
from app.core.pipelines.dspy_rag import DspyRagPipeline
from app.core.pipelines.question_decider import CodeRagQuestionDecider
+from app.core.pipelines.context_compressor import StringContextCompressor
from app.core.retrievers.file_retriever import FileRetriever
# A type hint for our handler functions
MessageHandler = Callable[[WebSocket, Dict[str, Any]], Awaitable[None]]
@@ -163,7 +165,8 @@
# 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.")
+ # logger.debug(f"File {file_path} is identical or older, skipping.")
+ pass
else:
# Case: This is a newly introduced file.
@@ -195,6 +198,8 @@
# 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] = {}):
if command_name not in self.command_map:
@@ -219,7 +224,7 @@
# 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}")
+ # logger.info(f"Received message: {message}")
handler = self.message_handlers.get(message_type)
if handler:
@@ -304,9 +309,9 @@
await self.send_command(websocket, "get_file_content", data={"filepaths": answer_text, "request_id": request_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."""
- # The client is expected to send a list of file objects
- # Each object should have 'filename' and 'content' keys.
+ """
+ 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")
@@ -316,17 +321,60 @@
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)
- data = self.file_retriever.retrieve_by_request_id(self.db, request_id=request_id)
+
+ # Retrieve the updated context from the database
+ context_data = self.file_retriever.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
+
+ # Use the LLM to make a decision
with dspy.context(lm=get_llm_provider("gemini")):
crqd = CodeRagQuestionDecider()
- answer_text = await crqd(
- question=data.get("question",""),
+ raw_answer_text, decision, code_diff = await crqd(
+ question=context_data.get("question", ""),
history="",
- retrieved_data= data)
- await websocket.send_text(json.dumps({
- "type": "chat_message",
- "content": answer_text
- }))
+ retrieved_data=context_data
+ )
+
+ if decision == "files":
+ await websocket.send_text(json.dumps({
+ "type": "thinking_log",
+ "content": f"AI decided more files are needed: {raw_answer_text}."
+ }))
+ try:
+ # The LLM is instructed to provide a JSON list, so we parse it
+ file_list = json.loads(raw_answer_text)
+ if not isinstance(file_list, list):
+ raise ValueError("Parsed result is not a list.")
+ except (ValueError, json.JSONDecodeError) as e:
+ print(f"Error parsing LLM output: {e}")
+ file_list = []
+ await websocket.send_text(json.dumps({
+ "type": "thinking_log",
+ "content": f"Warning: AI's file list could not be parsed. Error: {e}"
+ }))
+ return
+
+ await self.send_command(websocket, "get_file_content", data={"filepaths": file_list, "request_id": request_id})
+
+ elif decision == "code_change":
+ await websocket.send_text(json.dumps({
+ "type": "chat_message",
+ "content": raw_answer_text,
+ "code_diff": code_diff
+ }))
+
+ else: # decision is "answer"
+ await websocket.send_text(json.dumps({
+ "type": "chat_message",
+ "content": raw_answer_text
+ }))
async def handle_command_output(self, websocket: WebSocket, data: Dict[str, Any]):
"""Handles the output from a command executed by the client."""
diff --git a/ui/client-app/src/components/InteractionLog.js b/ui/client-app/src/components/InteractionLog.js
index fa2e232..3d43c14 100644
--- a/ui/client-app/src/components/InteractionLog.js
+++ b/ui/client-app/src/components/InteractionLog.js
@@ -22,7 +22,6 @@
}`}
>
- [{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 d1fb64f..ade4af9 100644
--- a/ui/client-app/src/hooks/useCodeAssistant.js
+++ b/ui/client-app/src/hooks/useCodeAssistant.js
@@ -19,8 +19,8 @@
const initialized = useRef(false);
const dirHandleRef = useRef(null);
- // --- WebSocket Message Handlers ---
const handleChatMessage = useCallback((message) => {
+ // Update chat history with the formatted content
setChatHistory((prev) => [...prev, { isUser: false, text: message.content }]);
setIsProcessing(false);
}, []);
@@ -120,17 +120,17 @@
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);
- }
+ 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;
+ console.error(`Error navigating to path part '${part}':`, error);
+ // Return null or re-throw based on desired error handling
+ return null;
}
}
return null;
@@ -343,4 +343,4 @@
};
};
-export default useCodeAssistant;
+export default useCodeAssistant;
\ No newline at end of file