WIP - communication memory

This commit is contained in:
Fahad
2025-06-10 19:16:51 +04:00
parent bb8a101dbf
commit f5060367a0
23 changed files with 2111 additions and 716 deletions

View File

@@ -88,6 +88,10 @@ class AnalyzeTool(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
},
},
"required": ["files", "question"],
}

View File

@@ -26,7 +26,7 @@ from pydantic import BaseModel, Field
from config import MCP_PROMPT_SIZE_LIMIT
from utils.file_utils import read_file_content, translate_path_for_environment
from .models import ClarificationRequest, ToolOutput
from .models import ClarificationRequest, ContinuationOffer, FollowUpRequest, ToolOutput
class ToolRequest(BaseModel):
@@ -50,6 +50,10 @@ class ToolRequest(BaseModel):
False,
description="Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
)
continuation_id: Optional[str] = Field(
None,
description="Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
)
class BaseTool(ABC):
@@ -375,6 +379,20 @@ If any of these would strengthen your analysis, specify what Claude should searc
# This is delegated to the tool implementation for customization
prompt = await self.prepare_prompt(request)
# Add follow-up instructions for new conversations (not threaded)
continuation_id = getattr(request, "continuation_id", None)
if not continuation_id:
# Import here to avoid circular imports
from server import get_follow_up_instructions
import logging
follow_up_instructions = get_follow_up_instructions(0) # New conversation, turn 0
prompt = f"{prompt}\n\n{follow_up_instructions}"
logging.debug(f"Added follow-up instructions for new {self.name} conversation")
else:
import logging
logging.debug(f"Continuing {self.name} conversation with thread {continuation_id}")
# Extract model configuration from request or use defaults
from config import GEMINI_MODEL
@@ -425,10 +443,10 @@ If any of these would strengthen your analysis, specify what Claude should searc
def _parse_response(self, raw_text: str, request) -> ToolOutput:
"""
Parse the raw response and determine if it's a clarification request.
Parse the raw response and determine if it's a clarification request or follow-up.
Some tools may return JSON indicating they need more information.
This method detects such responses and formats them appropriately.
Some tools may return JSON indicating they need more information or want to
continue the conversation. This method detects such responses and formats them.
Args:
raw_text: The raw text response from the model
@@ -437,6 +455,15 @@ If any of these would strengthen your analysis, specify what Claude should searc
Returns:
ToolOutput: Standardized output object
"""
# Check for follow-up questions in JSON blocks at the end of the response
follow_up_question = self._extract_follow_up_question(raw_text)
import logging
if follow_up_question:
logging.debug(f"Found follow-up question in {self.name} response: {follow_up_question.get('follow_up_question', 'N/A')}")
else:
logging.debug(f"No follow-up question found in {self.name} response")
try:
# Try to parse as JSON to check for clarification requests
potential_json = json.loads(raw_text.strip())
@@ -460,6 +487,20 @@ If any of these would strengthen your analysis, specify what Claude should searc
# Normal text response - format using tool-specific formatting
formatted_content = self.format_response(raw_text, request)
# If we found a follow-up question, prepare the threading response
if follow_up_question:
return self._create_follow_up_response(formatted_content, follow_up_question, request)
# Check if we should offer Claude a continuation opportunity
continuation_offer = self._check_continuation_opportunity(request)
import logging
if continuation_offer:
logging.debug(f"Creating continuation offer for {self.name} with {continuation_offer['remaining_turns']} turns remaining")
return self._create_continuation_offer_response(formatted_content, continuation_offer, request)
else:
logging.debug(f"No continuation offer created for {self.name}")
# Determine content type based on the formatted content
content_type = (
"markdown" if any(marker in formatted_content for marker in ["##", "**", "`", "- ", "1. "]) else "text"
@@ -472,6 +513,230 @@ If any of these would strengthen your analysis, specify what Claude should searc
metadata={"tool_name": self.name},
)
def _extract_follow_up_question(self, text: str) -> Optional[dict]:
"""
Extract follow-up question from JSON blocks in the response.
Looks for JSON blocks containing follow_up_question at the end of responses.
Args:
text: The response text to parse
Returns:
Dict with follow-up data if found, None otherwise
"""
import re
# Look for JSON blocks that contain follow_up_question
# Pattern handles optional leading whitespace and indentation
json_pattern = r'```json\s*\n\s*(\{.*?"follow_up_question".*?\})\s*\n\s*```'
matches = re.findall(json_pattern, text, re.DOTALL)
if not matches:
return None
# Take the last match (most recent follow-up)
try:
# Clean up the JSON string - remove excess whitespace and normalize
json_str = re.sub(r"\n\s+", "\n", matches[-1]).strip()
follow_up_data = json.loads(json_str)
if "follow_up_question" in follow_up_data:
return follow_up_data
except (json.JSONDecodeError, ValueError):
pass
return None
def _create_follow_up_response(self, content: str, follow_up_data: dict, request) -> ToolOutput:
"""
Create a response with follow-up question for conversation threading.
Args:
content: The main response content
follow_up_data: Dict containing follow_up_question and optional suggested_params
request: Original request for context
Returns:
ToolOutput configured for conversation continuation
"""
from utils.conversation_memory import add_turn, create_thread
# Create or get thread ID
continuation_id = getattr(request, "continuation_id", None)
if continuation_id:
# This is a continuation - add this turn to existing thread
request_files = getattr(request, "files", []) or []
success = add_turn(
continuation_id,
"assistant",
content,
follow_up_question=follow_up_data.get("follow_up_question"),
files=request_files,
tool_name=self.name,
)
if not success:
# Thread not found or at limit, return normal response
return ToolOutput(
status="success",
content=content,
content_type="markdown",
metadata={"tool_name": self.name},
)
thread_id = continuation_id
else:
# Create new thread
try:
thread_id = create_thread(
tool_name=self.name, initial_request=request.model_dump() if hasattr(request, "model_dump") else {}
)
# Add the assistant's response with follow-up
request_files = getattr(request, "files", []) or []
add_turn(
thread_id,
"assistant",
content,
follow_up_question=follow_up_data.get("follow_up_question"),
files=request_files,
tool_name=self.name,
)
except Exception as e:
# Threading failed, return normal response
import logging
logging.warning(f"Follow-up threading failed in {self.name}: {str(e)}")
return ToolOutput(
status="success",
content=content,
content_type="markdown",
metadata={"tool_name": self.name, "follow_up_error": str(e)},
)
# Create follow-up request
follow_up_request = FollowUpRequest(
continuation_id=thread_id,
question_to_user=follow_up_data["follow_up_question"],
suggested_tool_params=follow_up_data.get("suggested_params"),
ui_hint=follow_up_data.get("ui_hint"),
)
# Strip the JSON block from the content since it's now in the follow_up_request
clean_content = self._remove_follow_up_json(content)
return ToolOutput(
status="requires_continuation",
content=clean_content,
content_type="markdown",
follow_up_request=follow_up_request,
metadata={"tool_name": self.name, "thread_id": thread_id},
)
def _remove_follow_up_json(self, text: str) -> str:
"""Remove follow-up JSON blocks from the response text"""
import re
# Remove JSON blocks containing follow_up_question
pattern = r'```json\s*\n\s*\{.*?"follow_up_question".*?\}\s*\n\s*```'
return re.sub(pattern, "", text, flags=re.DOTALL).strip()
def _check_continuation_opportunity(self, request) -> Optional[dict]:
"""
Check if we should offer Claude a continuation opportunity.
This is called when Gemini doesn't ask a follow-up question, but we want
to give Claude the chance to continue the conversation if needed.
Args:
request: The original request
Returns:
Dict with continuation data if opportunity should be offered, None otherwise
"""
# Only offer continuation for new conversations (not already threaded)
continuation_id = getattr(request, "continuation_id", None)
if continuation_id:
# This is already a threaded conversation, don't offer continuation
# (either Gemini will ask follow-up or conversation naturally ends)
return None
# Only offer if we haven't reached conversation limits
try:
from utils.conversation_memory import MAX_CONVERSATION_TURNS
# For new conversations, we have MAX_CONVERSATION_TURNS - 1 remaining
# (since this response will be turn 1)
remaining_turns = MAX_CONVERSATION_TURNS - 1
if remaining_turns <= 0:
return None
# Offer continuation opportunity
return {"remaining_turns": remaining_turns, "tool_name": self.name}
except Exception:
# If anything fails, don't offer continuation
return None
def _create_continuation_offer_response(self, content: str, continuation_data: dict, request) -> ToolOutput:
"""
Create a response offering Claude the opportunity to continue conversation.
Args:
content: The main response content
continuation_data: Dict containing remaining_turns and tool_name
request: Original request for context
Returns:
ToolOutput configured with continuation offer
"""
from utils.conversation_memory import create_thread
try:
# Create new thread for potential continuation
thread_id = create_thread(
tool_name=self.name, initial_request=request.model_dump() if hasattr(request, "model_dump") else {}
)
# Add this response as the first turn (assistant turn)
from utils.conversation_memory import add_turn
request_files = getattr(request, "files", []) or []
add_turn(thread_id, "assistant", content, files=request_files, tool_name=self.name)
# Create continuation offer
remaining_turns = continuation_data["remaining_turns"]
continuation_offer = ContinuationOffer(
continuation_id=thread_id,
message_to_user=(
f"If you'd like to continue this analysis or need further details, "
f"you can use the continuation_id '{thread_id}' in your next {self.name} tool call. "
f"You have {remaining_turns} more exchange(s) available in this conversation thread."
),
suggested_tool_params={
"continuation_id": thread_id,
"prompt": "[Your follow-up question or request for additional analysis]",
},
remaining_turns=remaining_turns,
)
return ToolOutput(
status="continuation_available",
content=content,
content_type="markdown",
continuation_offer=continuation_offer,
metadata={"tool_name": self.name, "thread_id": thread_id, "remaining_turns": remaining_turns},
)
except Exception as e:
# If threading fails, return normal response but log the error
import logging
logging.warning(f"Conversation threading failed in {self.name}: {str(e)}")
return ToolOutput(
status="success",
content=content,
content_type="markdown",
metadata={"tool_name": self.name, "threading_error": str(e)},
)
@abstractmethod
async def prepare_prompt(self, request) -> str:
"""

View File

@@ -73,6 +73,10 @@ class ChatTool(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
},
},
"required": ["prompt"],
}

View File

@@ -126,6 +126,10 @@ class CodeReviewTool(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
},
},
"required": ["files", "context"],
}

View File

@@ -91,6 +91,10 @@ class DebugIssueTool(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
},
},
"required": ["error_description"],
}

View File

@@ -7,13 +7,50 @@ from typing import Any, Literal, Optional
from pydantic import BaseModel, Field
class FollowUpRequest(BaseModel):
"""Request for follow-up conversation turn"""
continuation_id: str = Field(..., description="Thread continuation ID for multi-turn conversations")
question_to_user: str = Field(..., description="Follow-up question to ask Claude")
suggested_tool_params: Optional[dict[str, Any]] = Field(
None, description="Suggested parameters for the next tool call"
)
ui_hint: Optional[str] = Field(
None, description="UI hint for Claude (e.g., 'text_input', 'file_select', 'multi_choice')"
)
class ContinuationOffer(BaseModel):
"""Offer for Claude to continue conversation when Gemini doesn't ask follow-up"""
continuation_id: str = Field(..., description="Thread continuation ID for multi-turn conversations")
message_to_user: str = Field(..., description="Message explaining continuation opportunity to Claude")
suggested_tool_params: Optional[dict[str, Any]] = Field(
None, description="Suggested parameters for continued tool usage"
)
remaining_turns: int = Field(..., description="Number of conversation turns remaining")
class ToolOutput(BaseModel):
"""Standardized output format for all tools"""
status: Literal["success", "error", "requires_clarification", "requires_file_prompt"] = "success"
content: str = Field(..., description="The main content/response from the tool")
status: Literal[
"success",
"error",
"requires_clarification",
"requires_file_prompt",
"requires_continuation",
"continuation_available",
] = "success"
content: Optional[str] = Field(None, description="The main content/response from the tool")
content_type: Literal["text", "markdown", "json"] = "text"
metadata: Optional[dict[str, Any]] = Field(default_factory=dict)
follow_up_request: Optional[FollowUpRequest] = Field(
None, description="Optional follow-up request for continued conversation"
)
continuation_offer: Optional[ContinuationOffer] = Field(
None, description="Optional offer for Claude to continue conversation"
)
class ClarificationRequest(BaseModel):

View File

@@ -100,6 +100,12 @@ class Precommit(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
}
# Add continuation_id parameter
if "properties" in schema and "continuation_id" not in schema["properties"]:
schema["properties"]["continuation_id"] = {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
}
return schema
def get_system_prompt(self) -> str:

View File

@@ -87,6 +87,10 @@ class ThinkDeepTool(BaseTool):
"description": "Enable web search for documentation, best practices, and current information. Particularly useful for: brainstorming sessions, architectural design discussions, exploring industry best practices, working with specific frameworks/technologies, researching solutions to complex problems, or when current documentation and community insights would enhance the analysis.",
"default": True,
},
"continuation_id": {
"type": "string",
"description": "Thread continuation ID for multi-turn conversations. Only provide this if continuing a previous conversation thread.",
},
},
"required": ["current_analysis"],
}