diff --git a/app/agentpress/__init__.py b/app/agentpress/__init__.py new file mode 100644 index 0000000..8ab9008 --- /dev/null +++ b/app/agentpress/__init__.py @@ -0,0 +1 @@ +# Utility functions and constants for agent tools \ No newline at end of file diff --git a/app/agentpress/context_manager.py b/app/agentpress/context_manager.py new file mode 100644 index 0000000..11405f4 --- /dev/null +++ b/app/agentpress/context_manager.py @@ -0,0 +1,298 @@ +""" +Context Management for AgentPress Threads. + +This module handles token counting and thread summarization to prevent +reaching the context window limitations of LLM models. +""" + +import json +from typing import List, Dict, Any, Optional + +from litellm import token_counter, completion_cost +from services.supabase import DBConnection +from services.llm import make_llm_api_call +from utils.logger import logger + +# Constants for token management +DEFAULT_TOKEN_THRESHOLD = 120000 # 80k tokens threshold for summarization +SUMMARY_TARGET_TOKENS = 10000 # Target ~10k tokens for the summary message +RESERVE_TOKENS = 5000 # Reserve tokens for new messages + +class ContextManager: + """Manages thread context including token counting and summarization.""" + + def __init__(self, token_threshold: int = DEFAULT_TOKEN_THRESHOLD): + """Initialize the ContextManager. + + Args: + token_threshold: Token count threshold to trigger summarization + """ + self.db = DBConnection() + self.token_threshold = token_threshold + + async def get_thread_token_count(self, thread_id: str) -> int: + """Get the current token count for a thread using LiteLLM. + + Args: + thread_id: ID of the thread to analyze + + Returns: + The total token count for relevant messages in the thread + """ + logger.debug(f"Getting token count for thread {thread_id}") + + try: + # Get messages for the thread + messages = await self.get_messages_for_summarization(thread_id) + + if not messages: + logger.debug(f"No messages found for thread {thread_id}") + return 0 + + # Use litellm's token_counter for accurate model-specific counting + # This is much more accurate than the SQL-based estimation + token_count = token_counter(model="gpt-4", messages=messages) + + logger.info(f"Thread {thread_id} has {token_count} tokens (calculated with litellm)") + return token_count + + except Exception as e: + logger.error(f"Error getting token count: {str(e)}") + return 0 + + async def get_messages_for_summarization(self, thread_id: str) -> List[Dict[str, Any]]: + """Get all LLM messages from the thread that need to be summarized. + + This gets messages after the most recent summary or all messages if + no summary exists. Unlike get_llm_messages, this includes ALL messages + since the last summary, even if we're generating a new summary. + + Args: + thread_id: ID of the thread to get messages from + + Returns: + List of message objects to summarize + """ + logger.debug(f"Getting messages for summarization for thread {thread_id}") + client = await self.db.client + + try: + # Find the most recent summary message + summary_result = await client.table('messages').select('created_at') \ + .eq('thread_id', thread_id) \ + .eq('type', 'summary') \ + .eq('is_llm_message', True) \ + .order('created_at', desc=True) \ + .limit(1) \ + .execute() + + # Get messages after the most recent summary or all messages if no summary + if summary_result.data and len(summary_result.data) > 0: + last_summary_time = summary_result.data[0]['created_at'] + logger.debug(f"Found last summary at {last_summary_time}") + + # Get all messages after the summary, but NOT including the summary itself + messages_result = await client.table('messages').select('*') \ + .eq('thread_id', thread_id) \ + .eq('is_llm_message', True) \ + .gt('created_at', last_summary_time) \ + .order('created_at') \ + .execute() + else: + logger.debug("No previous summary found, getting all messages") + # Get all messages + messages_result = await client.table('messages').select('*') \ + .eq('thread_id', thread_id) \ + .eq('is_llm_message', True) \ + .order('created_at') \ + .execute() + + # Parse the message content if needed + messages = [] + for msg in messages_result.data: + # Skip existing summary messages - we don't want to summarize summaries + if msg.get('type') == 'summary': + logger.debug(f"Skipping summary message from {msg.get('created_at')}") + continue + + # Parse content if it's a string + content = msg['content'] + if isinstance(content, str): + try: + content = json.loads(content) + except json.JSONDecodeError: + pass # Keep as string if not valid JSON + + # Ensure we have the proper format for the LLM + if 'role' not in content and 'type' in msg: + # Convert message type to role if needed + role = msg['type'] + if role == 'assistant' or role == 'user' or role == 'system' or role == 'tool': + content = {'role': role, 'content': content} + + messages.append(content) + + logger.info(f"Got {len(messages)} messages to summarize for thread {thread_id}") + return messages + + except Exception as e: + logger.error(f"Error getting messages for summarization: {str(e)}", exc_info=True) + return [] + + async def create_summary( + self, + thread_id: str, + messages: List[Dict[str, Any]], + model: str = "gpt-4o-mini" + ) -> Optional[Dict[str, Any]]: + """Generate a summary of conversation messages. + + Args: + thread_id: ID of the thread to summarize + messages: Messages to summarize + model: LLM model to use for summarization + + Returns: + Summary message object or None if summarization failed + """ + if not messages: + logger.warning("No messages to summarize") + return None + + logger.info(f"Creating summary for thread {thread_id} with {len(messages)} messages") + + # Create system message with summarization instructions + system_message = { + "role": "system", + "content": f"""You are a specialized summarization assistant. Your task is to create a concise but comprehensive summary of the conversation history. + +The summary should: +1. Preserve all key information including decisions, conclusions, and important context +2. Include any tools that were used and their results +3. Maintain chronological order of events +4. Be presented as a narrated list of key points with section headers +5. Include only factual information from the conversation (no new information) +6. Be concise but detailed enough that the conversation can continue with this summary as context + +VERY IMPORTANT: This summary will replace older parts of the conversation in the LLM's context window, so ensure it contains ALL key information and LATEST STATE OF THE CONVERSATION - SO WE WILL KNOW HOW TO PICK UP WHERE WE LEFT OFF. + + +THE CONVERSATION HISTORY TO SUMMARIZE IS AS FOLLOWS: +=============================================================== +==================== CONVERSATION HISTORY ==================== +{messages} +==================== END OF CONVERSATION HISTORY ==================== +=============================================================== +""" + } + + try: + # Call LLM to generate summary + response = await make_llm_api_call( + model_name=model, + messages=[system_message, {"role": "user", "content": "PLEASE PROVIDE THE SUMMARY NOW."}], + temperature=0, + max_tokens=SUMMARY_TARGET_TOKENS, + stream=False + ) + + if response and hasattr(response, 'choices') and response.choices: + summary_content = response.choices[0].message.content + + # Track token usage + try: + token_count = token_counter(model=model, messages=[{"role": "user", "content": summary_content}]) + cost = completion_cost(model=model, prompt="", completion=summary_content) + logger.info(f"Summary generated with {token_count} tokens at cost ${cost:.6f}") + except Exception as e: + logger.error(f"Error calculating token usage: {str(e)}") + + # Format the summary message with clear beginning and end markers + formatted_summary = f""" +======== CONVERSATION HISTORY SUMMARY ======== + +{summary_content} + +======== END OF SUMMARY ======== + +The above is a summary of the conversation history. The conversation continues below. +""" + + # Format the summary message + summary_message = { + "role": "user", + "content": formatted_summary + } + + return summary_message + else: + logger.error("Failed to generate summary: Invalid response") + return None + + except Exception as e: + logger.error(f"Error creating summary: {str(e)}", exc_info=True) + return None + + async def check_and_summarize_if_needed( + self, + thread_id: str, + add_message_callback, + model: str = "gpt-4o-mini", + force: bool = False + ) -> bool: + """Check if thread needs summarization and summarize if so. + + Args: + thread_id: ID of the thread to check + add_message_callback: Callback to add the summary message to the thread + model: LLM model to use for summarization + force: Whether to force summarization regardless of token count + + Returns: + True if summarization was performed, False otherwise + """ + try: + # Get token count using LiteLLM (accurate model-specific counting) + token_count = await self.get_thread_token_count(thread_id) + + # If token count is below threshold and not forcing, no summarization needed + if token_count < self.token_threshold and not force: + logger.debug(f"Thread {thread_id} has {token_count} tokens, below threshold {self.token_threshold}") + return False + + # Log reason for summarization + if force: + logger.info(f"Forced summarization of thread {thread_id} with {token_count} tokens") + else: + logger.info(f"Thread {thread_id} exceeds token threshold ({token_count} >= {self.token_threshold}), summarizing...") + + # Get messages to summarize + messages = await self.get_messages_for_summarization(thread_id) + + # If there are too few messages, don't summarize + if len(messages) < 3: + logger.info(f"Thread {thread_id} has too few messages ({len(messages)}) to summarize") + return False + + # Create summary + summary = await self.create_summary(thread_id, messages, model) + + if summary: + # Add summary message to thread + await add_message_callback( + thread_id=thread_id, + type="summary", + content=summary, + is_llm_message=True, + metadata={"token_count": token_count} + ) + + logger.info(f"Successfully added summary to thread {thread_id}") + return True + else: + logger.error(f"Failed to create summary for thread {thread_id}") + return False + + except Exception as e: + logger.error(f"Error in check_and_summarize_if_needed: {str(e)}", exc_info=True) + return False \ No newline at end of file diff --git a/app/agentpress/response_processor.py b/app/agentpress/response_processor.py new file mode 100644 index 0000000..68cd6bd --- /dev/null +++ b/app/agentpress/response_processor.py @@ -0,0 +1,1890 @@ +""" +Response processing module for AgentPress. + +This module handles the processing of LLM responses, including: +- Streaming and non-streaming response handling +- XML and native tool call detection and parsing +- Tool execution orchestration +- Message formatting and persistence +""" + +import json +import re +import uuid +import asyncio +from datetime import datetime, timezone +from typing import List, Dict, Any, Optional, AsyncGenerator, Tuple, Union, Callable, Literal +from dataclasses import dataclass +from utils.logger import logger +from agentpress.tool import ToolResult +from agentpress.tool_registry import ToolRegistry +from agentpress.xml_tool_parser import XMLToolParser +from langfuse.client import StatefulTraceClient +from services.langfuse import langfuse +from agentpress.utils.json_helpers import ( + ensure_dict, ensure_list, safe_json_parse, + to_json_string, format_for_yield +) +from litellm import token_counter + +# Type alias for XML result adding strategy +XmlAddingStrategy = Literal["user_message", "assistant_message", "inline_edit"] + +# Type alias for tool execution strategy +ToolExecutionStrategy = Literal["sequential", "parallel"] + +@dataclass +class ToolExecutionContext: + """Context for a tool execution including call details, result, and display info.""" + tool_call: Dict[str, Any] + tool_index: int + result: Optional[ToolResult] = None + function_name: Optional[str] = None + xml_tag_name: Optional[str] = None + error: Optional[Exception] = None + assistant_message_id: Optional[str] = None + parsing_details: Optional[Dict[str, Any]] = None + +@dataclass +class ProcessorConfig: + """ + Configuration for response processing and tool execution. + + This class controls how the LLM's responses are processed, including how tool calls + are detected, executed, and their results handled. + + Attributes: + xml_tool_calling: Enable XML-based tool call detection (...) + native_tool_calling: Enable OpenAI-style function calling format + execute_tools: Whether to automatically execute detected tool calls + execute_on_stream: For streaming, execute tools as they appear vs. at the end + tool_execution_strategy: How to execute multiple tools ("sequential" or "parallel") + xml_adding_strategy: How to add XML tool results to the conversation + max_xml_tool_calls: Maximum number of XML tool calls to process (0 = no limit) + """ + + xml_tool_calling: bool = True + native_tool_calling: bool = False + + execute_tools: bool = True + execute_on_stream: bool = False + tool_execution_strategy: ToolExecutionStrategy = "sequential" + xml_adding_strategy: XmlAddingStrategy = "assistant_message" + max_xml_tool_calls: int = 0 # 0 means no limit + + def __post_init__(self): + """Validate configuration after initialization.""" + if self.xml_tool_calling is False and self.native_tool_calling is False and self.execute_tools: + raise ValueError("At least one tool calling format (XML or native) must be enabled if execute_tools is True") + + if self.xml_adding_strategy not in ["user_message", "assistant_message", "inline_edit"]: + raise ValueError("xml_adding_strategy must be 'user_message', 'assistant_message', or 'inline_edit'") + + if self.max_xml_tool_calls < 0: + raise ValueError("max_xml_tool_calls must be a non-negative integer (0 = no limit)") + +class ResponseProcessor: + """Processes LLM responses, extracting and executing tool calls.""" + + def __init__(self, tool_registry: ToolRegistry, add_message_callback: Callable, trace: Optional[StatefulTraceClient] = None, is_agent_builder: bool = False, target_agent_id: Optional[str] = None, agent_config: Optional[dict] = None): + """Initialize the ResponseProcessor. + + Args: + tool_registry: Registry of available tools + add_message_callback: Callback function to add messages to the thread. + MUST return the full saved message object (dict) or None. + agent_config: Optional agent configuration with version information + """ + self.tool_registry = tool_registry + self.add_message = add_message_callback + self.trace = trace + if not self.trace: + self.trace = langfuse.trace(name="anonymous:response_processor") + # Initialize the XML parser with backwards compatibility + self.xml_parser = XMLToolParser(strict_mode=False) + self.is_agent_builder = is_agent_builder + self.target_agent_id = target_agent_id + self.agent_config = agent_config + + async def _yield_message(self, message_obj: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Helper to yield a message with proper formatting. + + Ensures that content and metadata are JSON strings for client compatibility. + """ + if message_obj: + return format_for_yield(message_obj) + + async def _add_message_with_agent_info( + self, + thread_id: str, + type: str, + content: Union[Dict[str, Any], List[Any], str], + is_llm_message: bool = False, + metadata: Optional[Dict[str, Any]] = None + ): + """Helper to add a message with agent version information if available.""" + agent_id = None + agent_version_id = None + + if self.agent_config: + agent_id = self.agent_config.get('agent_id') + agent_version_id = self.agent_config.get('current_version_id') + + return await self.add_message( + thread_id=thread_id, + type=type, + content=content, + is_llm_message=is_llm_message, + metadata=metadata, + agent_id=agent_id, + agent_version_id=agent_version_id + ) + + async def process_streaming_response( + self, + llm_response: AsyncGenerator, + thread_id: str, + prompt_messages: List[Dict[str, Any]], + llm_model: str, + config: ProcessorConfig = ProcessorConfig(), + ) -> AsyncGenerator[Dict[str, Any], None]: + """Process a streaming LLM response, handling tool calls and execution. + + Args: + llm_response: Streaming response from the LLM + thread_id: ID of the conversation thread + prompt_messages: List of messages sent to the LLM (the prompt) + llm_model: The name of the LLM model used + config: Configuration for parsing and execution + + Yields: + Complete message objects matching the DB schema, except for content chunks. + """ + accumulated_content = "" + tool_calls_buffer = {} + current_xml_content = "" + xml_chunks_buffer = [] + pending_tool_executions = [] + yielded_tool_indices = set() # Stores indices of tools whose *status* has been yielded + tool_index = 0 + xml_tool_call_count = 0 + finish_reason = None + last_assistant_message_object = None # Store the final saved assistant message object + tool_result_message_objects = {} # tool_index -> full saved message object + has_printed_thinking_prefix = False # Flag for printing thinking prefix only once + agent_should_terminate = False # Flag to track if a terminating tool has been executed + complete_native_tool_calls = [] # Initialize early for use in assistant_response_end + + # Collect metadata for reconstructing LiteLLM response object + streaming_metadata = { + "model": llm_model, + "created": None, + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + }, + "response_ms": None, + "first_chunk_time": None, + "last_chunk_time": None + } + + logger.info(f"Streaming Config: XML={config.xml_tool_calling}, Native={config.native_tool_calling}, " + f"Execute on stream={config.execute_on_stream}, Strategy={config.tool_execution_strategy}") + + thread_run_id = str(uuid.uuid4()) + + try: + # --- Save and Yield Start Events --- + start_content = {"status_type": "thread_run_start", "thread_run_id": thread_run_id} + start_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=start_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if start_msg_obj: yield format_for_yield(start_msg_obj) + + assist_start_content = {"status_type": "assistant_response_start"} + assist_start_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=assist_start_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if assist_start_msg_obj: yield format_for_yield(assist_start_msg_obj) + # --- End Start Events --- + + __sequence = 0 + + async for chunk in llm_response: + # Extract streaming metadata from chunks + current_time = datetime.now(timezone.utc).timestamp() + if streaming_metadata["first_chunk_time"] is None: + streaming_metadata["first_chunk_time"] = current_time + streaming_metadata["last_chunk_time"] = current_time + + # Extract metadata from chunk attributes + if hasattr(chunk, 'created') and chunk.created: + streaming_metadata["created"] = chunk.created + if hasattr(chunk, 'model') and chunk.model: + streaming_metadata["model"] = chunk.model + if hasattr(chunk, 'usage') and chunk.usage: + # Update usage information if available (including zero values) + if hasattr(chunk.usage, 'prompt_tokens') and chunk.usage.prompt_tokens is not None: + streaming_metadata["usage"]["prompt_tokens"] = chunk.usage.prompt_tokens + if hasattr(chunk.usage, 'completion_tokens') and chunk.usage.completion_tokens is not None: + streaming_metadata["usage"]["completion_tokens"] = chunk.usage.completion_tokens + if hasattr(chunk.usage, 'total_tokens') and chunk.usage.total_tokens is not None: + streaming_metadata["usage"]["total_tokens"] = chunk.usage.total_tokens + + if hasattr(chunk, 'choices') and chunk.choices and hasattr(chunk.choices[0], 'finish_reason') and chunk.choices[0].finish_reason: + finish_reason = chunk.choices[0].finish_reason + logger.debug(f"Detected finish_reason: {finish_reason}") + + if hasattr(chunk, 'choices') and chunk.choices: + delta = chunk.choices[0].delta if hasattr(chunk.choices[0], 'delta') else None + + # Check for and log Anthropic thinking content + if delta and hasattr(delta, 'reasoning_content') and delta.reasoning_content: + if not has_printed_thinking_prefix: + # print("[THINKING]: ", end='', flush=True) + has_printed_thinking_prefix = True + # print(delta.reasoning_content, end='', flush=True) + # Append reasoning to main content to be saved in the final message + accumulated_content += delta.reasoning_content + + # Process content chunk + if delta and hasattr(delta, 'content') and delta.content: + chunk_content = delta.content + # print(chunk_content, end='', flush=True) + accumulated_content += chunk_content + current_xml_content += chunk_content + + if not (config.max_xml_tool_calls > 0 and xml_tool_call_count >= config.max_xml_tool_calls): + # Yield ONLY content chunk (don't save) + now_chunk = datetime.now(timezone.utc).isoformat() + yield { + "sequence": __sequence, + "message_id": None, "thread_id": thread_id, "type": "assistant", + "is_llm_message": True, + "content": to_json_string({"role": "assistant", "content": chunk_content}), + "metadata": to_json_string({"stream_status": "chunk", "thread_run_id": thread_run_id}), + "created_at": now_chunk, "updated_at": now_chunk + } + __sequence += 1 + else: + logger.info("XML tool call limit reached - not yielding more content chunks") + self.trace.event(name="xml_tool_call_limit_reached", level="DEFAULT", status_message=(f"XML tool call limit reached - not yielding more content chunks")) + + # --- Process XML Tool Calls (if enabled and limit not reached) --- + if config.xml_tool_calling and not (config.max_xml_tool_calls > 0 and xml_tool_call_count >= config.max_xml_tool_calls): + xml_chunks = self._extract_xml_chunks(current_xml_content) + for xml_chunk in xml_chunks: + current_xml_content = current_xml_content.replace(xml_chunk, "", 1) + xml_chunks_buffer.append(xml_chunk) + result = self._parse_xml_tool_call(xml_chunk) + if result: + tool_call, parsing_details = result + xml_tool_call_count += 1 + current_assistant_id = last_assistant_message_object['message_id'] if last_assistant_message_object else None + context = self._create_tool_context( + tool_call, tool_index, current_assistant_id, parsing_details + ) + + if config.execute_tools and config.execute_on_stream: + # Save and Yield tool_started status + started_msg_obj = await self._yield_and_save_tool_started(context, thread_id, thread_run_id) + if started_msg_obj: yield format_for_yield(started_msg_obj) + yielded_tool_indices.add(tool_index) # Mark status as yielded + + execution_task = asyncio.create_task(self._execute_tool(tool_call)) + pending_tool_executions.append({ + "task": execution_task, "tool_call": tool_call, + "tool_index": tool_index, "context": context + }) + tool_index += 1 + + if config.max_xml_tool_calls > 0 and xml_tool_call_count >= config.max_xml_tool_calls: + logger.debug(f"Reached XML tool call limit ({config.max_xml_tool_calls})") + finish_reason = "xml_tool_limit_reached" + break # Stop processing more XML chunks in this delta + + # --- Process Native Tool Call Chunks --- + if config.native_tool_calling and delta and hasattr(delta, 'tool_calls') and delta.tool_calls: + for tool_call_chunk in delta.tool_calls: + # Yield Native Tool Call Chunk (transient status, not saved) + # ... (safe extraction logic for tool_call_data_chunk) ... + tool_call_data_chunk = {} # Placeholder for extracted data + if hasattr(tool_call_chunk, 'model_dump'): tool_call_data_chunk = tool_call_chunk.model_dump() + else: # Manual extraction... + if hasattr(tool_call_chunk, 'id'): tool_call_data_chunk['id'] = tool_call_chunk.id + if hasattr(tool_call_chunk, 'index'): tool_call_data_chunk['index'] = tool_call_chunk.index + if hasattr(tool_call_chunk, 'type'): tool_call_data_chunk['type'] = tool_call_chunk.type + if hasattr(tool_call_chunk, 'function'): + tool_call_data_chunk['function'] = {} + if hasattr(tool_call_chunk.function, 'name'): tool_call_data_chunk['function']['name'] = tool_call_chunk.function.name + if hasattr(tool_call_chunk.function, 'arguments'): tool_call_data_chunk['function']['arguments'] = tool_call_chunk.function.arguments if isinstance(tool_call_chunk.function.arguments, str) else to_json_string(tool_call_chunk.function.arguments) + + + now_tool_chunk = datetime.now(timezone.utc).isoformat() + yield { + "message_id": None, "thread_id": thread_id, "type": "status", "is_llm_message": True, + "content": to_json_string({"role": "assistant", "status_type": "tool_call_chunk", "tool_call_chunk": tool_call_data_chunk}), + "metadata": to_json_string({"thread_run_id": thread_run_id}), + "created_at": now_tool_chunk, "updated_at": now_tool_chunk + } + + # --- Buffer and Execute Complete Native Tool Calls --- + if not hasattr(tool_call_chunk, 'function'): continue + idx = tool_call_chunk.index if hasattr(tool_call_chunk, 'index') else 0 + # ... (buffer update logic remains same) ... + # ... (check complete logic remains same) ... + has_complete_tool_call = False # Placeholder + if (tool_calls_buffer.get(idx) and + tool_calls_buffer[idx]['id'] and + tool_calls_buffer[idx]['function']['name'] and + tool_calls_buffer[idx]['function']['arguments']): + try: + safe_json_parse(tool_calls_buffer[idx]['function']['arguments']) + has_complete_tool_call = True + except json.JSONDecodeError: pass + + + if has_complete_tool_call and config.execute_tools and config.execute_on_stream: + current_tool = tool_calls_buffer[idx] + tool_call_data = { + "function_name": current_tool['function']['name'], + "arguments": safe_json_parse(current_tool['function']['arguments']), + "id": current_tool['id'] + } + current_assistant_id = last_assistant_message_object['message_id'] if last_assistant_message_object else None + context = self._create_tool_context( + tool_call_data, tool_index, current_assistant_id + ) + + # Save and Yield tool_started status + started_msg_obj = await self._yield_and_save_tool_started(context, thread_id, thread_run_id) + if started_msg_obj: yield format_for_yield(started_msg_obj) + yielded_tool_indices.add(tool_index) # Mark status as yielded + + execution_task = asyncio.create_task(self._execute_tool(tool_call_data)) + pending_tool_executions.append({ + "task": execution_task, "tool_call": tool_call_data, + "tool_index": tool_index, "context": context + }) + tool_index += 1 + + if finish_reason == "xml_tool_limit_reached": + logger.info("Stopping stream processing after loop due to XML tool call limit") + self.trace.event(name="stopping_stream_processing_after_loop_due_to_xml_tool_call_limit", level="DEFAULT", status_message=(f"Stopping stream processing after loop due to XML tool call limit")) + break + + # print() # Add a final newline after the streaming loop finishes + + # --- After Streaming Loop --- + + if ( + streaming_metadata["usage"]["total_tokens"] == 0 + ): + logger.info("🔥 No usage data from provider, counting with litellm.token_counter") + + try: + # prompt side + prompt_tokens = token_counter( + model=llm_model, + messages=prompt_messages # chat or plain; token_counter handles both + ) + + # completion side + completion_tokens = token_counter( + model=llm_model, + text=accumulated_content or "" # empty string safe + ) + + streaming_metadata["usage"]["prompt_tokens"] = prompt_tokens + streaming_metadata["usage"]["completion_tokens"] = completion_tokens + streaming_metadata["usage"]["total_tokens"] = prompt_tokens + completion_tokens + + logger.info( + f"🔥 Estimated tokens – prompt: {prompt_tokens}, " + f"completion: {completion_tokens}, total: {prompt_tokens + completion_tokens}" + ) + self.trace.event(name="usage_calculated_with_litellm_token_counter", level="DEFAULT", status_message=(f"Usage calculated with litellm.token_counter")) + except Exception as e: + logger.warning(f"Failed to calculate usage: {str(e)}") + self.trace.event(name="failed_to_calculate_usage", level="WARNING", status_message=(f"Failed to calculate usage: {str(e)}")) + + + # Wait for pending tool executions from streaming phase + tool_results_buffer = [] # Stores (tool_call, result, tool_index, context) + if pending_tool_executions: + logger.info(f"Waiting for {len(pending_tool_executions)} pending streamed tool executions") + self.trace.event(name="waiting_for_pending_streamed_tool_executions", level="DEFAULT", status_message=(f"Waiting for {len(pending_tool_executions)} pending streamed tool executions")) + # ... (asyncio.wait logic) ... + pending_tasks = [execution["task"] for execution in pending_tool_executions] + done, _ = await asyncio.wait(pending_tasks) + + for execution in pending_tool_executions: + tool_idx = execution.get("tool_index", -1) + context = execution["context"] + tool_name = context.function_name + + # Check if status was already yielded during stream run + if tool_idx in yielded_tool_indices: + logger.debug(f"Status for tool index {tool_idx} already yielded.") + # Still need to process the result for the buffer + try: + if execution["task"].done(): + result = execution["task"].result() + context.result = result + tool_results_buffer.append((execution["tool_call"], result, tool_idx, context)) + + if tool_name in ['ask', 'complete']: + logger.info(f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.") + self.trace.event(name="terminating_tool_completed_during_streaming", level="DEFAULT", status_message=(f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")) + agent_should_terminate = True + + else: # Should not happen with asyncio.wait + logger.warning(f"Task for tool index {tool_idx} not done after wait.") + self.trace.event(name="task_for_tool_index_not_done_after_wait", level="WARNING", status_message=(f"Task for tool index {tool_idx} not done after wait.")) + except Exception as e: + logger.error(f"Error getting result for pending tool execution {tool_idx}: {str(e)}") + self.trace.event(name="error_getting_result_for_pending_tool_execution", level="ERROR", status_message=(f"Error getting result for pending tool execution {tool_idx}: {str(e)}")) + context.error = e + # Save and Yield tool error status message (even if started was yielded) + error_msg_obj = await self._yield_and_save_tool_error(context, thread_id, thread_run_id) + if error_msg_obj: yield format_for_yield(error_msg_obj) + continue # Skip further status yielding for this tool index + + # If status wasn't yielded before (shouldn't happen with current logic), yield it now + try: + if execution["task"].done(): + result = execution["task"].result() + context.result = result + tool_results_buffer.append((execution["tool_call"], result, tool_idx, context)) + + # Check if this is a terminating tool + if tool_name in ['ask', 'complete']: + logger.info(f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.") + self.trace.event(name="terminating_tool_completed_during_streaming", level="DEFAULT", status_message=(f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")) + agent_should_terminate = True + + # Save and Yield tool completed/failed status + completed_msg_obj = await self._yield_and_save_tool_completed( + context, None, thread_id, thread_run_id + ) + if completed_msg_obj: yield format_for_yield(completed_msg_obj) + yielded_tool_indices.add(tool_idx) + except Exception as e: + logger.error(f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}") + self.trace.event(name="error_getting_result_yielding_status_for_pending_tool_execution", level="ERROR", status_message=(f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}")) + context.error = e + # Save and Yield tool error status + error_msg_obj = await self._yield_and_save_tool_error(context, thread_id, thread_run_id) + if error_msg_obj: yield format_for_yield(error_msg_obj) + yielded_tool_indices.add(tool_idx) + + + # Save and yield finish status if limit was reached + if finish_reason == "xml_tool_limit_reached": + finish_content = {"status_type": "finish", "finish_reason": "xml_tool_limit_reached"} + finish_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=finish_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if finish_msg_obj: yield format_for_yield(finish_msg_obj) + logger.info(f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls") + self.trace.event(name="stream_finished_with_reason_xml_tool_limit_reached_after_xml_tool_calls", level="DEFAULT", status_message=(f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls")) + + # --- SAVE and YIELD Final Assistant Message --- + if accumulated_content: + # ... (Truncate accumulated_content logic) ... + if config.max_xml_tool_calls > 0 and xml_tool_call_count >= config.max_xml_tool_calls and xml_chunks_buffer: + last_xml_chunk = xml_chunks_buffer[-1] + last_chunk_end_pos = accumulated_content.find(last_xml_chunk) + len(last_xml_chunk) + if last_chunk_end_pos > 0: + accumulated_content = accumulated_content[:last_chunk_end_pos] + + # ... (Extract complete_native_tool_calls logic) ... + # Update complete_native_tool_calls from buffer (initialized earlier) + if config.native_tool_calling: + for idx, tc_buf in tool_calls_buffer.items(): + if tc_buf['id'] and tc_buf['function']['name'] and tc_buf['function']['arguments']: + try: + args = safe_json_parse(tc_buf['function']['arguments']) + complete_native_tool_calls.append({ + "id": tc_buf['id'], "type": "function", + "function": {"name": tc_buf['function']['name'],"arguments": args} + }) + except json.JSONDecodeError: continue + + message_data = { # Dict to be saved in 'content' + "role": "assistant", "content": accumulated_content, + "tool_calls": complete_native_tool_calls or None + } + + last_assistant_message_object = await self._add_message_with_agent_info( + thread_id=thread_id, type="assistant", content=message_data, + is_llm_message=True, metadata={"thread_run_id": thread_run_id} + ) + + if last_assistant_message_object: + # Yield the complete saved object, adding stream_status metadata just for yield + yield_metadata = ensure_dict(last_assistant_message_object.get('metadata'), {}) + yield_metadata['stream_status'] = 'complete' + # Format the message for yielding + yield_message = last_assistant_message_object.copy() + yield_message['metadata'] = yield_metadata + yield format_for_yield(yield_message) + else: + logger.error(f"Failed to save final assistant message for thread {thread_id}") + self.trace.event(name="failed_to_save_final_assistant_message_for_thread", level="ERROR", status_message=(f"Failed to save final assistant message for thread {thread_id}")) + # Save and yield an error status + err_content = {"role": "system", "status_type": "error", "message": "Failed to save final assistant message"} + err_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=err_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if err_msg_obj: yield format_for_yield(err_msg_obj) + + # --- Process All Tool Results Now --- + if config.execute_tools: + final_tool_calls_to_process = [] + # ... (Gather final_tool_calls_to_process from native and XML buffers) ... + # Gather native tool calls from buffer + if config.native_tool_calling and complete_native_tool_calls: + for tc in complete_native_tool_calls: + final_tool_calls_to_process.append({ + "function_name": tc["function"]["name"], + "arguments": tc["function"]["arguments"], # Already parsed object + "id": tc["id"] + }) + # Gather XML tool calls from buffer (up to limit) + parsed_xml_data = [] + if config.xml_tool_calling: + # Reparse remaining content just in case (should be empty if processed correctly) + xml_chunks = self._extract_xml_chunks(current_xml_content) + xml_chunks_buffer.extend(xml_chunks) + # Process only chunks not already handled in the stream loop + remaining_limit = config.max_xml_tool_calls - xml_tool_call_count if config.max_xml_tool_calls > 0 else len(xml_chunks_buffer) + xml_chunks_to_process = xml_chunks_buffer[:remaining_limit] # Ensure limit is respected + + for chunk in xml_chunks_to_process: + parsed_result = self._parse_xml_tool_call(chunk) + if parsed_result: + tool_call, parsing_details = parsed_result + # Avoid adding if already processed during streaming + if not any(exec['tool_call'] == tool_call for exec in pending_tool_executions): + final_tool_calls_to_process.append(tool_call) + parsed_xml_data.append({'tool_call': tool_call, 'parsing_details': parsing_details}) + + + all_tool_data_map = {} # tool_index -> {'tool_call': ..., 'parsing_details': ...} + # Add native tool data + native_tool_index = 0 + if config.native_tool_calling and complete_native_tool_calls: + for tc in complete_native_tool_calls: + # Find the corresponding entry in final_tool_calls_to_process if needed + # For now, assume order matches if only native used + exec_tool_call = { + "function_name": tc["function"]["name"], + "arguments": tc["function"]["arguments"], + "id": tc["id"] + } + all_tool_data_map[native_tool_index] = {"tool_call": exec_tool_call, "parsing_details": None} + native_tool_index += 1 + + # Add XML tool data + xml_tool_index_start = native_tool_index + for idx, item in enumerate(parsed_xml_data): + all_tool_data_map[xml_tool_index_start + idx] = item + + + tool_results_map = {} # tool_index -> (tool_call, result, context) + + # Populate from buffer if executed on stream + if config.execute_on_stream and tool_results_buffer: + logger.info(f"Processing {len(tool_results_buffer)} buffered tool results") + self.trace.event(name="processing_buffered_tool_results", level="DEFAULT", status_message=(f"Processing {len(tool_results_buffer)} buffered tool results")) + for tool_call, result, tool_idx, context in tool_results_buffer: + if last_assistant_message_object: context.assistant_message_id = last_assistant_message_object['message_id'] + tool_results_map[tool_idx] = (tool_call, result, context) + + # Or execute now if not streamed + elif final_tool_calls_to_process and not config.execute_on_stream: + logger.info(f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream") + self.trace.event(name="executing_tools_after_stream", level="DEFAULT", status_message=(f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream")) + results_list = await self._execute_tools(final_tool_calls_to_process, config.tool_execution_strategy) + current_tool_idx = 0 + for tc, res in results_list: + # Map back using all_tool_data_map which has correct indices + if current_tool_idx in all_tool_data_map: + tool_data = all_tool_data_map[current_tool_idx] + context = self._create_tool_context( + tc, current_tool_idx, + last_assistant_message_object['message_id'] if last_assistant_message_object else None, + tool_data.get('parsing_details') + ) + context.result = res + tool_results_map[current_tool_idx] = (tc, res, context) + else: + logger.warning(f"Could not map result for tool index {current_tool_idx}") + self.trace.event(name="could_not_map_result_for_tool_index", level="WARNING", status_message=(f"Could not map result for tool index {current_tool_idx}")) + current_tool_idx += 1 + + # Save and Yield each result message + if tool_results_map: + logger.info(f"Saving and yielding {len(tool_results_map)} final tool result messages") + self.trace.event(name="saving_and_yielding_final_tool_result_messages", level="DEFAULT", status_message=(f"Saving and yielding {len(tool_results_map)} final tool result messages")) + for tool_idx in sorted(tool_results_map.keys()): + tool_call, result, context = tool_results_map[tool_idx] + context.result = result + if not context.assistant_message_id and last_assistant_message_object: + context.assistant_message_id = last_assistant_message_object['message_id'] + + # Yield start status ONLY IF executing non-streamed (already yielded if streamed) + if not config.execute_on_stream and tool_idx not in yielded_tool_indices: + started_msg_obj = await self._yield_and_save_tool_started(context, thread_id, thread_run_id) + if started_msg_obj: yield format_for_yield(started_msg_obj) + yielded_tool_indices.add(tool_idx) # Mark status yielded + + # Save the tool result message to DB + saved_tool_result_object = await self._add_tool_result( # Returns full object or None + thread_id, tool_call, result, config.xml_adding_strategy, + context.assistant_message_id, context.parsing_details + ) + + # Yield completed/failed status (linked to saved result ID if available) + completed_msg_obj = await self._yield_and_save_tool_completed( + context, + saved_tool_result_object['message_id'] if saved_tool_result_object else None, + thread_id, thread_run_id + ) + if completed_msg_obj: yield format_for_yield(completed_msg_obj) + # Don't add to yielded_tool_indices here, completion status is separate yield + + # Yield the saved tool result object + if saved_tool_result_object: + tool_result_message_objects[tool_idx] = saved_tool_result_object + yield format_for_yield(saved_tool_result_object) + else: + logger.error(f"Failed to save tool result for index {tool_idx}, not yielding result message.") + self.trace.event(name="failed_to_save_tool_result_for_index", level="ERROR", status_message=(f"Failed to save tool result for index {tool_idx}, not yielding result message.")) + # Optionally yield error status for saving failure? + + # --- Final Finish Status --- + if finish_reason and finish_reason != "xml_tool_limit_reached": + finish_content = {"status_type": "finish", "finish_reason": finish_reason} + finish_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=finish_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if finish_msg_obj: yield format_for_yield(finish_msg_obj) + + # Check if agent should terminate after processing pending tools + if agent_should_terminate: + logger.info("Agent termination requested after executing ask/complete tool. Stopping further processing.") + self.trace.event(name="agent_termination_requested", level="DEFAULT", status_message="Agent termination requested after executing ask/complete tool. Stopping further processing.") + + # Set finish reason to indicate termination + finish_reason = "agent_terminated" + + # Save and yield termination status + finish_content = {"status_type": "finish", "finish_reason": "agent_terminated"} + finish_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=finish_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if finish_msg_obj: yield format_for_yield(finish_msg_obj) + + # Save assistant_response_end BEFORE terminating + if last_assistant_message_object: + try: + # Calculate response time if we have timing data + if streaming_metadata["first_chunk_time"] and streaming_metadata["last_chunk_time"]: + streaming_metadata["response_ms"] = (streaming_metadata["last_chunk_time"] - streaming_metadata["first_chunk_time"]) * 1000 + + # Create a LiteLLM-like response object for streaming (before termination) + # Check if we have any actual usage data + has_usage_data = ( + streaming_metadata["usage"]["prompt_tokens"] > 0 or + streaming_metadata["usage"]["completion_tokens"] > 0 or + streaming_metadata["usage"]["total_tokens"] > 0 + ) + + assistant_end_content = { + "choices": [ + { + "finish_reason": finish_reason or "stop", + "index": 0, + "message": { + "role": "assistant", + "content": accumulated_content, + "tool_calls": complete_native_tool_calls or None + } + } + ], + "created": streaming_metadata.get("created"), + "model": streaming_metadata.get("model", llm_model), + "usage": streaming_metadata["usage"], # Always include usage like LiteLLM does + "streaming": True, # Add flag to indicate this was reconstructed from streaming + } + + # Only include response_ms if we have timing data + if streaming_metadata.get("response_ms"): + assistant_end_content["response_ms"] = streaming_metadata["response_ms"] + + await self.add_message( + thread_id=thread_id, + type="assistant_response_end", + content=assistant_end_content, + is_llm_message=False, + metadata={"thread_run_id": thread_run_id} + ) + logger.info("Assistant response end saved for stream (before termination)") + except Exception as e: + logger.error(f"Error saving assistant response end for stream (before termination): {str(e)}") + self.trace.event(name="error_saving_assistant_response_end_for_stream_before_termination", level="ERROR", status_message=(f"Error saving assistant response end for stream (before termination): {str(e)}")) + + # Skip all remaining processing and go to finally block + return + + # --- Save and Yield assistant_response_end --- + if last_assistant_message_object: # Only save if assistant message was saved + try: + # Calculate response time if we have timing data + if streaming_metadata["first_chunk_time"] and streaming_metadata["last_chunk_time"]: + streaming_metadata["response_ms"] = (streaming_metadata["last_chunk_time"] - streaming_metadata["first_chunk_time"]) * 1000 + + # Create a LiteLLM-like response object for streaming + # Check if we have any actual usage data + has_usage_data = ( + streaming_metadata["usage"]["prompt_tokens"] > 0 or + streaming_metadata["usage"]["completion_tokens"] > 0 or + streaming_metadata["usage"]["total_tokens"] > 0 + ) + + assistant_end_content = { + "choices": [ + { + "finish_reason": finish_reason or "stop", + "index": 0, + "message": { + "role": "assistant", + "content": accumulated_content, + "tool_calls": complete_native_tool_calls or None + } + } + ], + "created": streaming_metadata.get("created"), + "model": streaming_metadata.get("model", llm_model), + "usage": streaming_metadata["usage"], # Always include usage like LiteLLM does + "streaming": True, # Add flag to indicate this was reconstructed from streaming + } + + # Only include response_ms if we have timing data + if streaming_metadata.get("response_ms"): + assistant_end_content["response_ms"] = streaming_metadata["response_ms"] + + await self.add_message( + thread_id=thread_id, + type="assistant_response_end", + content=assistant_end_content, + is_llm_message=False, + metadata={"thread_run_id": thread_run_id} + ) + logger.info("Assistant response end saved for stream") + except Exception as e: + logger.error(f"Error saving assistant response end for stream: {str(e)}") + self.trace.event(name="error_saving_assistant_response_end_for_stream", level="ERROR", status_message=(f"Error saving assistant response end for stream: {str(e)}")) + + except Exception as e: + logger.error(f"Error processing stream: {str(e)}", exc_info=True) + self.trace.event(name="error_processing_stream", level="ERROR", status_message=(f"Error processing stream: {str(e)}")) + # Save and yield error status message + err_content = {"role": "system", "status_type": "error", "message": str(e)} + err_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=err_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None} + ) + if err_msg_obj: yield format_for_yield(err_msg_obj) # Yield the saved error message + + # Re-raise the same exception (not a new one) to ensure proper error propagation + logger.critical(f"Re-raising error to stop further processing: {str(e)}") + self.trace.event(name="re_raising_error_to_stop_further_processing", level="ERROR", status_message=(f"Re-raising error to stop further processing: {str(e)}")) + raise # Use bare 'raise' to preserve the original exception with its traceback + + finally: + # Save and Yield the final thread_run_end status + try: + end_content = {"status_type": "thread_run_end"} + end_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=end_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None} + ) + if end_msg_obj: yield format_for_yield(end_msg_obj) + except Exception as final_e: + logger.error(f"Error in finally block: {str(final_e)}", exc_info=True) + self.trace.event(name="error_in_finally_block", level="ERROR", status_message=(f"Error in finally block: {str(final_e)}")) + + async def process_non_streaming_response( + self, + llm_response: Any, + thread_id: str, + prompt_messages: List[Dict[str, Any]], + llm_model: str, + config: ProcessorConfig = ProcessorConfig(), + ) -> AsyncGenerator[Dict[str, Any], None]: + """Process a non-streaming LLM response, handling tool calls and execution. + + Args: + llm_response: Response from the LLM + thread_id: ID of the conversation thread + prompt_messages: List of messages sent to the LLM (the prompt) + llm_model: The name of the LLM model used + config: Configuration for parsing and execution + + Yields: + Complete message objects matching the DB schema. + """ + content = "" + thread_run_id = str(uuid.uuid4()) + all_tool_data = [] # Stores {'tool_call': ..., 'parsing_details': ...} + tool_index = 0 + assistant_message_object = None + tool_result_message_objects = {} + finish_reason = None + native_tool_calls_for_message = [] + + try: + # Save and Yield thread_run_start status message + start_content = {"status_type": "thread_run_start", "thread_run_id": thread_run_id} + start_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=start_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if start_msg_obj: yield format_for_yield(start_msg_obj) + + # Extract finish_reason, content, tool calls + if hasattr(llm_response, 'choices') and llm_response.choices: + if hasattr(llm_response.choices[0], 'finish_reason'): + finish_reason = llm_response.choices[0].finish_reason + logger.info(f"Non-streaming finish_reason: {finish_reason}") + self.trace.event(name="non_streaming_finish_reason", level="DEFAULT", status_message=(f"Non-streaming finish_reason: {finish_reason}")) + response_message = llm_response.choices[0].message if hasattr(llm_response.choices[0], 'message') else None + if response_message: + if hasattr(response_message, 'content') and response_message.content: + content = response_message.content + if config.xml_tool_calling: + parsed_xml_data = self._parse_xml_tool_calls(content) + if config.max_xml_tool_calls > 0 and len(parsed_xml_data) > config.max_xml_tool_calls: + # Truncate content and tool data if limit exceeded + # ... (Truncation logic similar to streaming) ... + if parsed_xml_data: + xml_chunks = self._extract_xml_chunks(content)[:config.max_xml_tool_calls] + if xml_chunks: + last_chunk = xml_chunks[-1] + last_chunk_pos = content.find(last_chunk) + if last_chunk_pos >= 0: content = content[:last_chunk_pos + len(last_chunk)] + parsed_xml_data = parsed_xml_data[:config.max_xml_tool_calls] + finish_reason = "xml_tool_limit_reached" + all_tool_data.extend(parsed_xml_data) + + if config.native_tool_calling and hasattr(response_message, 'tool_calls') and response_message.tool_calls: + for tool_call in response_message.tool_calls: + if hasattr(tool_call, 'function'): + exec_tool_call = { + "function_name": tool_call.function.name, + "arguments": safe_json_parse(tool_call.function.arguments) if isinstance(tool_call.function.arguments, str) else tool_call.function.arguments, + "id": tool_call.id if hasattr(tool_call, 'id') else str(uuid.uuid4()) + } + all_tool_data.append({"tool_call": exec_tool_call, "parsing_details": None}) + native_tool_calls_for_message.append({ + "id": exec_tool_call["id"], "type": "function", + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments if isinstance(tool_call.function.arguments, str) else to_json_string(tool_call.function.arguments) + } + }) + + + # --- SAVE and YIELD Final Assistant Message --- + message_data = {"role": "assistant", "content": content, "tool_calls": native_tool_calls_for_message or None} + assistant_message_object = await self._add_message_with_agent_info( + thread_id=thread_id, type="assistant", content=message_data, + is_llm_message=True, metadata={"thread_run_id": thread_run_id} + ) + if assistant_message_object: + yield assistant_message_object + else: + logger.error(f"Failed to save non-streaming assistant message for thread {thread_id}") + self.trace.event(name="failed_to_save_non_streaming_assistant_message_for_thread", level="ERROR", status_message=(f"Failed to save non-streaming assistant message for thread {thread_id}")) + err_content = {"role": "system", "status_type": "error", "message": "Failed to save assistant message"} + err_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=err_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if err_msg_obj: yield format_for_yield(err_msg_obj) + + # --- Execute Tools and Yield Results --- + tool_calls_to_execute = [item['tool_call'] for item in all_tool_data] + if config.execute_tools and tool_calls_to_execute: + logger.info(f"Executing {len(tool_calls_to_execute)} tools with strategy: {config.tool_execution_strategy}") + self.trace.event(name="executing_tools_with_strategy", level="DEFAULT", status_message=(f"Executing {len(tool_calls_to_execute)} tools with strategy: {config.tool_execution_strategy}")) + tool_results = await self._execute_tools(tool_calls_to_execute, config.tool_execution_strategy) + + for i, (returned_tool_call, result) in enumerate(tool_results): + original_data = all_tool_data[i] + tool_call_from_data = original_data['tool_call'] + parsing_details = original_data['parsing_details'] + current_assistant_id = assistant_message_object['message_id'] if assistant_message_object else None + + context = self._create_tool_context( + tool_call_from_data, tool_index, current_assistant_id, parsing_details + ) + context.result = result + + # Save and Yield start status + started_msg_obj = await self._yield_and_save_tool_started(context, thread_id, thread_run_id) + if started_msg_obj: yield format_for_yield(started_msg_obj) + + # Save tool result + saved_tool_result_object = await self._add_tool_result( + thread_id, tool_call_from_data, result, config.xml_adding_strategy, + current_assistant_id, parsing_details + ) + + # Save and Yield completed/failed status + completed_msg_obj = await self._yield_and_save_tool_completed( + context, + saved_tool_result_object['message_id'] if saved_tool_result_object else None, + thread_id, thread_run_id + ) + if completed_msg_obj: yield format_for_yield(completed_msg_obj) + + # Yield the saved tool result object + if saved_tool_result_object: + tool_result_message_objects[tool_index] = saved_tool_result_object + yield format_for_yield(saved_tool_result_object) + else: + logger.error(f"Failed to save tool result for index {tool_index}") + self.trace.event(name="failed_to_save_tool_result_for_index", level="ERROR", status_message=(f"Failed to save tool result for index {tool_index}")) + + tool_index += 1 + + # --- Save and Yield Final Status --- + if finish_reason: + finish_content = {"status_type": "finish", "finish_reason": finish_reason} + finish_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=finish_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id} + ) + if finish_msg_obj: yield format_for_yield(finish_msg_obj) + + # --- Save and Yield assistant_response_end --- + if assistant_message_object: # Only save if assistant message was saved + try: + # Save the full LiteLLM response object directly in content + await self.add_message( + thread_id=thread_id, + type="assistant_response_end", + content=llm_response, + is_llm_message=False, + metadata={"thread_run_id": thread_run_id} + ) + logger.info("Assistant response end saved for non-stream") + except Exception as e: + logger.error(f"Error saving assistant response end for non-stream: {str(e)}") + self.trace.event(name="error_saving_assistant_response_end_for_non_stream", level="ERROR", status_message=(f"Error saving assistant response end for non-stream: {str(e)}")) + + except Exception as e: + logger.error(f"Error processing non-streaming response: {str(e)}", exc_info=True) + self.trace.event(name="error_processing_non_streaming_response", level="ERROR", status_message=(f"Error processing non-streaming response: {str(e)}")) + # Save and yield error status + err_content = {"role": "system", "status_type": "error", "message": str(e)} + err_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=err_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None} + ) + if err_msg_obj: yield format_for_yield(err_msg_obj) + + # Re-raise the same exception (not a new one) to ensure proper error propagation + logger.critical(f"Re-raising error to stop further processing: {str(e)}") + self.trace.event(name="re_raising_error_to_stop_further_processing", level="CRITICAL", status_message=(f"Re-raising error to stop further processing: {str(e)}")) + raise # Use bare 'raise' to preserve the original exception with its traceback + + finally: + # Save and Yield the final thread_run_end status + end_content = {"status_type": "thread_run_end"} + end_msg_obj = await self.add_message( + thread_id=thread_id, type="status", content=end_content, + is_llm_message=False, metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None} + ) + if end_msg_obj: yield format_for_yield(end_msg_obj) + + # XML parsing methods + def _extract_tag_content(self, xml_chunk: str, tag_name: str) -> Tuple[Optional[str], Optional[str]]: + """Extract content between opening and closing tags, handling nested tags.""" + start_tag = f'<{tag_name}' + end_tag = f'' + + try: + # Find start tag position + start_pos = xml_chunk.find(start_tag) + if start_pos == -1: + return None, xml_chunk + + # Find end of opening tag + tag_end = xml_chunk.find('>', start_pos) + if tag_end == -1: + return None, xml_chunk + + # Find matching closing tag + content_start = tag_end + 1 + nesting_level = 1 + pos = content_start + + while nesting_level > 0 and pos < len(xml_chunk): + next_start = xml_chunk.find(start_tag, pos) + next_end = xml_chunk.find(end_tag, pos) + + if next_end == -1: + return None, xml_chunk + + if next_start != -1 and next_start < next_end: + nesting_level += 1 + pos = next_start + len(start_tag) + else: + nesting_level -= 1 + if nesting_level == 0: + content = xml_chunk[content_start:next_end] + remaining = xml_chunk[next_end + len(end_tag):] + return content, remaining + else: + pos = next_end + len(end_tag) + + return None, xml_chunk + + except Exception as e: + logger.error(f"Error extracting tag content: {e}") + self.trace.event(name="error_extracting_tag_content", level="ERROR", status_message=(f"Error extracting tag content: {e}")) + return None, xml_chunk + + def _extract_attribute(self, opening_tag: str, attr_name: str) -> Optional[str]: + """Extract attribute value from opening tag.""" + try: + # Handle both single and double quotes with raw strings + patterns = [ + fr'{attr_name}="([^"]*)"', # Double quotes + fr"{attr_name}='([^']*)'", # Single quotes + fr'{attr_name}=([^\s/>;]+)' # No quotes - fixed escape sequence + ] + + for pattern in patterns: + match = re.search(pattern, opening_tag) + if match: + value = match.group(1) + # Unescape common XML entities + value = value.replace('"', '"').replace(''', "'") + value = value.replace('<', '<').replace('>', '>') + value = value.replace('&', '&') + return value + + return None + + except Exception as e: + logger.error(f"Error extracting attribute: {e}") + self.trace.event(name="error_extracting_attribute", level="ERROR", status_message=(f"Error extracting attribute: {e}")) + return None + + def _extract_xml_chunks(self, content: str) -> List[str]: + """Extract complete XML chunks using start and end pattern matching.""" + chunks = [] + pos = 0 + + try: + # First, look for new format blocks + start_pattern = '' + end_pattern = '' + + while pos < len(content): + # Find the next function_calls block + start_pos = content.find(start_pattern, pos) + if start_pos == -1: + break + + # Find the matching end tag + end_pos = content.find(end_pattern, start_pos) + if end_pos == -1: + break + + # Extract the complete block including tags + chunk_end = end_pos + len(end_pattern) + chunk = content[start_pos:chunk_end] + chunks.append(chunk) + + # Move position past this chunk + pos = chunk_end + + # If no new format found, fall back to old format for backwards compatibility + if not chunks: + pos = 0 + while pos < len(content): + # Find the next tool tag + next_tag_start = -1 + current_tag = None + + # Find the earliest occurrence of any registered tag + for tag_name in self.tool_registry.xml_tools.keys(): + start_pattern = f'<{tag_name}' + tag_pos = content.find(start_pattern, pos) + + if tag_pos != -1 and (next_tag_start == -1 or tag_pos < next_tag_start): + next_tag_start = tag_pos + current_tag = tag_name + + if next_tag_start == -1 or not current_tag: + break + + # Find the matching end tag + end_pattern = f'' + tag_stack = [] + chunk_start = next_tag_start + current_pos = next_tag_start + + while current_pos < len(content): + # Look for next start or end tag of the same type + next_start = content.find(f'<{current_tag}', current_pos + 1) + next_end = content.find(end_pattern, current_pos) + + if next_end == -1: # No closing tag found + break + + if next_start != -1 and next_start < next_end: + # Found nested start tag + tag_stack.append(next_start) + current_pos = next_start + 1 + else: + # Found end tag + if not tag_stack: # This is our matching end tag + chunk_end = next_end + len(end_pattern) + chunk = content[chunk_start:chunk_end] + chunks.append(chunk) + pos = chunk_end + break + else: + # Pop nested tag + tag_stack.pop() + current_pos = next_end + 1 + + if current_pos >= len(content): # Reached end without finding closing tag + break + + pos = max(pos + 1, current_pos) + + except Exception as e: + logger.error(f"Error extracting XML chunks: {e}") + logger.error(f"Content was: {content}") + self.trace.event(name="error_extracting_xml_chunks", level="ERROR", status_message=(f"Error extracting XML chunks: {e}"), metadata={"content": content}) + + return chunks + + def _parse_xml_tool_call(self, xml_chunk: str) -> Optional[Tuple[Dict[str, Any], Dict[str, Any]]]: + """Parse XML chunk into tool call format and return parsing details. + + Returns: + Tuple of (tool_call, parsing_details) or None if parsing fails. + - tool_call: Dict with 'function_name', 'xml_tag_name', 'arguments' + - parsing_details: Dict with 'attributes', 'elements', 'text_content', 'root_content' + """ + try: + # Check if this is the new format (contains ) + if '' in xml_chunk and ']+)', xml_chunk) + if not tag_match: + logger.error(f"No tag found in XML chunk: {xml_chunk}") + self.trace.event(name="no_tag_found_in_xml_chunk", level="ERROR", status_message=(f"No tag found in XML chunk: {xml_chunk}")) + return None + + # This is the XML tag as it appears in the text (e.g., "create-file") + xml_tag_name = tag_match.group(1) + logger.info(f"Found XML tag: {xml_tag_name}") + self.trace.event(name="found_xml_tag", level="DEFAULT", status_message=(f"Found XML tag: {xml_tag_name}")) + + # Get tool info and schema from registry + tool_info = self.tool_registry.get_xml_tool(xml_tag_name) + if not tool_info or not tool_info['schema'].xml_schema: + logger.error(f"No tool or schema found for tag: {xml_tag_name}") + self.trace.event(name="no_tool_or_schema_found_for_tag", level="ERROR", status_message=(f"No tool or schema found for tag: {xml_tag_name}")) + return None + + # This is the actual function name to call (e.g., "create_file") + function_name = tool_info['method'] + + schema = tool_info['schema'].xml_schema + params = {} + remaining_chunk = xml_chunk + + # --- Store detailed parsing info --- + parsing_details = { + "attributes": {}, + "elements": {}, + "text_content": None, + "root_content": None, + "raw_chunk": xml_chunk # Store the original chunk for reference + } + # --- + + # Process each mapping + for mapping in schema.mappings: + try: + if mapping.node_type == "attribute": + # Extract attribute from opening tag + opening_tag = remaining_chunk.split('>', 1)[0] + value = self._extract_attribute(opening_tag, mapping.param_name) + if value is not None: + params[mapping.param_name] = value + parsing_details["attributes"][mapping.param_name] = value # Store raw attribute + # logger.info(f"Found attribute {mapping.param_name}: {value}") + + elif mapping.node_type == "element": + # Extract element content + content, remaining_chunk = self._extract_tag_content(remaining_chunk, mapping.path) + if content is not None: + params[mapping.param_name] = content.strip() + parsing_details["elements"][mapping.param_name] = content.strip() # Store raw element content + # logger.info(f"Found element {mapping.param_name}: {content.strip()}") + + elif mapping.node_type == "text": + # Extract text content + content, _ = self._extract_tag_content(remaining_chunk, xml_tag_name) + if content is not None: + params[mapping.param_name] = content.strip() + parsing_details["text_content"] = content.strip() # Store raw text content + # logger.info(f"Found text content for {mapping.param_name}: {content.strip()}") + + elif mapping.node_type == "content": + # Extract root content + content, _ = self._extract_tag_content(remaining_chunk, xml_tag_name) + if content is not None: + params[mapping.param_name] = content.strip() + parsing_details["root_content"] = content.strip() # Store raw root content + # logger.info(f"Found root content for {mapping.param_name}") + + except Exception as e: + logger.error(f"Error processing mapping {mapping}: {e}") + self.trace.event(name="error_processing_mapping", level="ERROR", status_message=(f"Error processing mapping {mapping}: {e}")) + continue + + # Create tool call with clear separation between function_name and xml_tag_name + tool_call = { + "function_name": function_name, # The actual method to call (e.g., create_file) + "xml_tag_name": xml_tag_name, # The original XML tag (e.g., create-file) + "arguments": params # The extracted parameters + } + + # logger.debug(f"Parsed old format tool call: {tool_call["function_name"]}") + return tool_call, parsing_details # Return both dicts + + except Exception as e: + logger.error(f"Error parsing XML chunk: {e}") + logger.error(f"XML chunk was: {xml_chunk}") + self.trace.event(name="error_parsing_xml_chunk", level="ERROR", status_message=(f"Error parsing XML chunk: {e}"), metadata={"xml_chunk": xml_chunk}) + return None + + def _parse_xml_tool_calls(self, content: str) -> List[Dict[str, Any]]: + """Parse XML tool calls from content string. + + Returns: + List of dictionaries, each containing {'tool_call': ..., 'parsing_details': ...} + """ + parsed_data = [] + + try: + xml_chunks = self._extract_xml_chunks(content) + + for xml_chunk in xml_chunks: + result = self._parse_xml_tool_call(xml_chunk) + if result: + tool_call, parsing_details = result + parsed_data.append({ + "tool_call": tool_call, + "parsing_details": parsing_details + }) + + except Exception as e: + logger.error(f"Error parsing XML tool calls: {e}", exc_info=True) + self.trace.event(name="error_parsing_xml_tool_calls", level="ERROR", status_message=(f"Error parsing XML tool calls: {e}"), metadata={"content": content}) + + return parsed_data + + # Tool execution methods + async def _execute_tool(self, tool_call: Dict[str, Any]) -> ToolResult: + """Execute a single tool call and return the result.""" + span = self.trace.span(name=f"execute_tool.{tool_call['function_name']}", input=tool_call["arguments"]) + try: + function_name = tool_call["function_name"] + arguments = tool_call["arguments"] + + logger.info(f"Executing tool: {function_name} with arguments: {arguments}") + self.trace.event(name="executing_tool", level="DEFAULT", status_message=(f"Executing tool: {function_name} with arguments: {arguments}")) + + if isinstance(arguments, str): + try: + arguments = safe_json_parse(arguments) + except json.JSONDecodeError: + arguments = {"text": arguments} + + # Get available functions from tool registry + available_functions = self.tool_registry.get_available_functions() + + # Look up the function by name + tool_fn = available_functions.get(function_name) + if not tool_fn: + logger.error(f"Tool function '{function_name}' not found in registry") + span.end(status_message="tool_not_found", level="ERROR") + return ToolResult(success=False, output=f"Tool function '{function_name}' not found") + + logger.debug(f"Found tool function for '{function_name}', executing...") + result = await tool_fn(**arguments) + logger.info(f"Tool execution complete: {function_name} -> {result}") + span.end(status_message="tool_executed", output=result) + return result + except Exception as e: + logger.error(f"Error executing tool {tool_call['function_name']}: {str(e)}", exc_info=True) + span.end(status_message="tool_execution_error", output=f"Error executing tool: {str(e)}", level="ERROR") + return ToolResult(success=False, output=f"Error executing tool: {str(e)}") + + async def _execute_tools( + self, + tool_calls: List[Dict[str, Any]], + execution_strategy: ToolExecutionStrategy = "sequential" + ) -> List[Tuple[Dict[str, Any], ToolResult]]: + """Execute tool calls with the specified strategy. + + This is the main entry point for tool execution. It dispatches to the appropriate + execution method based on the provided strategy. + + Args: + tool_calls: List of tool calls to execute + execution_strategy: Strategy for executing tools: + - "sequential": Execute tools one after another, waiting for each to complete + - "parallel": Execute all tools simultaneously for better performance + + Returns: + List of tuples containing the original tool call and its result + """ + logger.info(f"Executing {len(tool_calls)} tools with strategy: {execution_strategy}") + self.trace.event(name="executing_tools_with_strategy", level="DEFAULT", status_message=(f"Executing {len(tool_calls)} tools with strategy: {execution_strategy}")) + + if execution_strategy == "sequential": + return await self._execute_tools_sequentially(tool_calls) + elif execution_strategy == "parallel": + return await self._execute_tools_in_parallel(tool_calls) + else: + logger.warning(f"Unknown execution strategy: {execution_strategy}, falling back to sequential") + return await self._execute_tools_sequentially(tool_calls) + + async def _execute_tools_sequentially(self, tool_calls: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], ToolResult]]: + """Execute tool calls sequentially and return results. + + This method executes tool calls one after another, waiting for each tool to complete + before starting the next one. This is useful when tools have dependencies on each other. + + Args: + tool_calls: List of tool calls to execute + + Returns: + List of tuples containing the original tool call and its result + """ + if not tool_calls: + return [] + + try: + tool_names = [t.get('function_name', 'unknown') for t in tool_calls] + logger.info(f"Executing {len(tool_calls)} tools sequentially: {tool_names}") + self.trace.event(name="executing_tools_sequentially", level="DEFAULT", status_message=(f"Executing {len(tool_calls)} tools sequentially: {tool_names}")) + + results = [] + for index, tool_call in enumerate(tool_calls): + tool_name = tool_call.get('function_name', 'unknown') + logger.debug(f"Executing tool {index+1}/{len(tool_calls)}: {tool_name}") + + try: + result = await self._execute_tool(tool_call) + results.append((tool_call, result)) + logger.debug(f"Completed tool {tool_name} with success={result.success}") + + # Check if this is a terminating tool (ask or complete) + if tool_name in ['ask', 'complete']: + logger.info(f"Terminating tool '{tool_name}' executed. Stopping further tool execution.") + self.trace.event(name="terminating_tool_executed", level="DEFAULT", status_message=(f"Terminating tool '{tool_name}' executed. Stopping further tool execution.")) + break # Stop executing remaining tools + + except Exception as e: + logger.error(f"Error executing tool {tool_name}: {str(e)}") + self.trace.event(name="error_executing_tool", level="ERROR", status_message=(f"Error executing tool {tool_name}: {str(e)}")) + error_result = ToolResult(success=False, output=f"Error executing tool: {str(e)}") + results.append((tool_call, error_result)) + + logger.info(f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)") + self.trace.event(name="sequential_execution_completed", level="DEFAULT", status_message=(f"Sequential execution completed for {len(results)} tools (out of {len(tool_calls)} total)")) + return results + + except Exception as e: + logger.error(f"Error in sequential tool execution: {str(e)}", exc_info=True) + # Return partial results plus error results for remaining tools + completed_tool_names = [r[0].get('function_name', 'unknown') for r in results] if 'results' in locals() else [] + remaining_tools = [t for t in tool_calls if t.get('function_name', 'unknown') not in completed_tool_names] + + # Add error results for remaining tools + error_results = [(tool, ToolResult(success=False, output=f"Execution error: {str(e)}")) + for tool in remaining_tools] + + return (results if 'results' in locals() else []) + error_results + + async def _execute_tools_in_parallel(self, tool_calls: List[Dict[str, Any]]) -> List[Tuple[Dict[str, Any], ToolResult]]: + """Execute tool calls in parallel and return results. + + This method executes all tool calls simultaneously using asyncio.gather, which + can significantly improve performance when executing multiple independent tools. + + Args: + tool_calls: List of tool calls to execute + + Returns: + List of tuples containing the original tool call and its result + """ + if not tool_calls: + return [] + + try: + tool_names = [t.get('function_name', 'unknown') for t in tool_calls] + logger.info(f"Executing {len(tool_calls)} tools in parallel: {tool_names}") + self.trace.event(name="executing_tools_in_parallel", level="DEFAULT", status_message=(f"Executing {len(tool_calls)} tools in parallel: {tool_names}")) + + # Create tasks for all tool calls + tasks = [self._execute_tool(tool_call) for tool_call in tool_calls] + + # Execute all tasks concurrently with error handling + results = await asyncio.gather(*tasks, return_exceptions=True) + + # Process results and handle any exceptions + processed_results = [] + for i, (tool_call, result) in enumerate(zip(tool_calls, results)): + if isinstance(result, Exception): + logger.error(f"Error executing tool {tool_call.get('function_name', 'unknown')}: {str(result)}") + self.trace.event(name="error_executing_tool", level="ERROR", status_message=(f"Error executing tool {tool_call.get('function_name', 'unknown')}: {str(result)}")) + # Create error result + error_result = ToolResult(success=False, output=f"Error executing tool: {str(result)}") + processed_results.append((tool_call, error_result)) + else: + processed_results.append((tool_call, result)) + + logger.info(f"Parallel execution completed for {len(tool_calls)} tools") + self.trace.event(name="parallel_execution_completed", level="DEFAULT", status_message=(f"Parallel execution completed for {len(tool_calls)} tools")) + return processed_results + + except Exception as e: + logger.error(f"Error in parallel tool execution: {str(e)}", exc_info=True) + self.trace.event(name="error_in_parallel_tool_execution", level="ERROR", status_message=(f"Error in parallel tool execution: {str(e)}")) + # Return error results for all tools if the gather itself fails + return [(tool_call, ToolResult(success=False, output=f"Execution error: {str(e)}")) + for tool_call in tool_calls] + + async def _add_tool_result( + self, + thread_id: str, + tool_call: Dict[str, Any], + result: ToolResult, + strategy: Union[XmlAddingStrategy, str] = "assistant_message", + assistant_message_id: Optional[str] = None, + parsing_details: Optional[Dict[str, Any]] = None + ) -> Optional[Dict[str, Any]]: # Return the full message object + """Add a tool result to the conversation thread based on the specified format. + + This method formats tool results and adds them to the conversation history, + making them visible to the LLM in subsequent interactions. Results can be + added either as native tool messages (OpenAI format) or as XML-wrapped content + with a specified role (user or assistant). + + Args: + thread_id: ID of the conversation thread + tool_call: The original tool call that produced this result + result: The result from the tool execution + strategy: How to add XML tool results to the conversation + ("user_message", "assistant_message", or "inline_edit") + assistant_message_id: ID of the assistant message that generated this tool call + parsing_details: Detailed parsing info for XML calls (attributes, elements, etc.) + """ + try: + message_obj = None # Initialize message_obj + + # Create metadata with assistant_message_id if provided + metadata = {} + if assistant_message_id: + metadata["assistant_message_id"] = assistant_message_id + logger.info(f"Linking tool result to assistant message: {assistant_message_id}") + self.trace.event(name="linking_tool_result_to_assistant_message", level="DEFAULT", status_message=(f"Linking tool result to assistant message: {assistant_message_id}")) + + # --- Add parsing details to metadata if available --- + if parsing_details: + metadata["parsing_details"] = parsing_details + logger.info("Adding parsing_details to tool result metadata") + self.trace.event(name="adding_parsing_details_to_tool_result_metadata", level="DEFAULT", status_message=(f"Adding parsing_details to tool result metadata"), metadata={"parsing_details": parsing_details}) + # --- + + # Check if this is a native function call (has id field) + if "id" in tool_call: + # Format as a proper tool message according to OpenAI spec + function_name = tool_call.get("function_name", "") + + # Format the tool result content - tool role needs string content + if isinstance(result, str): + content = result + elif hasattr(result, 'output'): + # If it's a ToolResult object + if isinstance(result.output, dict) or isinstance(result.output, list): + # If output is already a dict or list, convert to JSON string + content = json.dumps(result.output) + else: + # Otherwise just use the string representation + content = str(result.output) + else: + # Fallback to string representation of the whole result + content = str(result) + + logger.info(f"Formatted tool result content: {content[:100]}...") + self.trace.event(name="formatted_tool_result_content", level="DEFAULT", status_message=(f"Formatted tool result content: {content[:100]}...")) + + # Create the tool response message with proper format + tool_message = { + "role": "tool", + "tool_call_id": tool_call["id"], + "name": function_name, + "content": content + } + + logger.info(f"Adding native tool result for tool_call_id={tool_call['id']} with role=tool") + self.trace.event(name="adding_native_tool_result_for_tool_call_id", level="DEFAULT", status_message=(f"Adding native tool result for tool_call_id={tool_call['id']} with role=tool")) + + # Add as a tool message to the conversation history + # This makes the result visible to the LLM in the next turn + message_obj = await self.add_message( + thread_id=thread_id, + type="tool", # Special type for tool responses + content=tool_message, + is_llm_message=True, + metadata=metadata + ) + return message_obj # Return the full message object + + # Check if this is an MCP tool (function_name starts with "call_mcp_tool") + function_name = tool_call.get("function_name", "") + + # Check if this is an MCP tool - either the old call_mcp_tool or a dynamically registered MCP tool + is_mcp_tool = False + if function_name == "call_mcp_tool": + is_mcp_tool = True + else: + # Check if the result indicates it's an MCP tool by looking for MCP metadata + if hasattr(result, 'output') and isinstance(result.output, str): + # Check for MCP metadata pattern in the output + if "MCP Tool Result from" in result.output and "Tool Metadata:" in result.output: + is_mcp_tool = True + # Also check for MCP metadata in JSON format + elif "mcp_metadata" in result.output: + is_mcp_tool = True + + if is_mcp_tool: + # Special handling for MCP tools - make content prominent and LLM-friendly + result_role = "user" if strategy == "user_message" else "assistant" + + # Extract the actual content from the ToolResult + if hasattr(result, 'output'): + mcp_content = str(result.output) + else: + mcp_content = str(result) + + # Create a simple, LLM-friendly message format that puts content first + simple_message = { + "role": result_role, + "content": mcp_content # Direct content, no complex nesting + } + + logger.info(f"Adding MCP tool result with simplified format for LLM visibility") + self.trace.event(name="adding_mcp_tool_result_simplified", level="DEFAULT", status_message="Adding MCP tool result with simplified format for LLM visibility") + + message_obj = await self.add_message( + thread_id=thread_id, + type="tool", + content=simple_message, + is_llm_message=True, + metadata=metadata + ) + return message_obj + + # For XML and other non-native tools, use the new structured format + # Determine message role based on strategy + result_role = "user" if strategy == "user_message" else "assistant" + + # Create the new structured tool result format + structured_result = self._create_structured_tool_result(tool_call, result, parsing_details) + + # Add the message with the appropriate role to the conversation history + # This allows the LLM to see the tool result in subsequent interactions + result_message = { + "role": result_role, + "content": json.dumps(structured_result) + } + message_obj = await self.add_message( + thread_id=thread_id, + type="tool", + content=result_message, + is_llm_message=True, + metadata=metadata + ) + return message_obj # Return the full message object + except Exception as e: + logger.error(f"Error adding tool result: {str(e)}", exc_info=True) + self.trace.event(name="error_adding_tool_result", level="ERROR", status_message=(f"Error adding tool result: {str(e)}"), metadata={"tool_call": tool_call, "result": result, "strategy": strategy, "assistant_message_id": assistant_message_id, "parsing_details": parsing_details}) + # Fallback to a simple message + try: + fallback_message = { + "role": "user", + "content": str(result) + } + message_obj = await self.add_message( + thread_id=thread_id, + type="tool", + content=fallback_message, + is_llm_message=True, + metadata={"assistant_message_id": assistant_message_id} if assistant_message_id else {} + ) + return message_obj # Return the full message object + except Exception as e2: + logger.error(f"Failed even with fallback message: {str(e2)}", exc_info=True) + self.trace.event(name="failed_even_with_fallback_message", level="ERROR", status_message=(f"Failed even with fallback message: {str(e2)}"), metadata={"tool_call": tool_call, "result": result, "strategy": strategy, "assistant_message_id": assistant_message_id, "parsing_details": parsing_details}) + return None # Return None on error + + def _create_structured_tool_result(self, tool_call: Dict[str, Any], result: ToolResult, parsing_details: Optional[Dict[str, Any]] = None): + """Create a structured tool result format that's tool-agnostic and provides rich information. + + Args: + tool_call: The original tool call that was executed + result: The result from the tool execution + parsing_details: Optional parsing details for XML calls + + Returns: + Structured dictionary containing tool execution information + """ + # Extract tool information + function_name = tool_call.get("function_name", "unknown") + xml_tag_name = tool_call.get("xml_tag_name") + arguments = tool_call.get("arguments", {}) + tool_call_id = tool_call.get("id") + logger.info(f"Creating structured tool result for tool_call: {tool_call}") + + # Process the output - if it's a JSON string, parse it back to an object + output = result.output if hasattr(result, 'output') else str(result) + if isinstance(output, str): + try: + # Try to parse as JSON to provide structured data to frontend + parsed_output = safe_json_parse(output) + # If parsing succeeded and we got a dict/list, use the parsed version + if isinstance(parsed_output, (dict, list)): + output = parsed_output + # Otherwise keep the original string + except Exception: + # If parsing fails, keep the original string + pass + + # Create the structured result + structured_result_v1 = { + "tool_execution": { + "function_name": function_name, + "xml_tag_name": xml_tag_name, + "tool_call_id": tool_call_id, + "arguments": arguments, + "result": { + "success": result.success if hasattr(result, 'success') else True, + "output": output, # Now properly structured for frontend + "error": getattr(result, 'error', None) if hasattr(result, 'error') else None + }, + # "execution_details": { + # "timestamp": datetime.now(timezone.utc).isoformat(), + # "parsing_details": parsing_details + # } + } + } + + # STRUCTURED_OUTPUT_TOOLS = { + # "str_replace", + # "get_data_provider_endpoints", + # } + + # summary_output = result.output if hasattr(result, 'output') else str(result) + + # if xml_tag_name: + # status = "completed successfully" if structured_result_v1["tool_execution"]["result"]["success"] else "failed" + # summary = f"Tool '{xml_tag_name}' {status}. Output: {summary_output}" + # else: + # status = "completed successfully" if structured_result_v1["tool_execution"]["result"]["success"] else "failed" + # summary = f"Function '{function_name}' {status}. Output: {summary_output}" + + # if self.is_agent_builder: + # return summary + # if function_name in STRUCTURED_OUTPUT_TOOLS: + # return structured_result_v1 + # else: + # return summary + + summary_output = result.output if hasattr(result, 'output') else str(result) + success_status = structured_result_v1["tool_execution"]["result"]["success"] + + # # Create a more comprehensive summary for the LLM + # if xml_tag_name: + # status = "completed successfully" if structured_result_v1["tool_execution"]["result"]["success"] else "failed" + # summary = f"Tool '{xml_tag_name}' {status}. Output: {summary_output}" + # else: + # status = "completed successfully" if structured_result_v1["tool_execution"]["result"]["success"] else "failed" + # summary = f"Function '{function_name}' {status}. Output: {summary_output}" + + # if self.is_agent_builder: + # return summary + # elif function_name == "get_data_provider_endpoints": + # logger.info(f"Returning sumnary for data provider call: {summary}") + # return summary + + return structured_result_v1 + + def _format_xml_tool_result(self, tool_call: Dict[str, Any], result: ToolResult) -> str: + """Format a tool result wrapped in a tag. + + DEPRECATED: This method is kept for backwards compatibility. + New implementations should use _create_structured_tool_result instead. + + Args: + tool_call: The tool call that was executed + result: The result of the tool execution + + Returns: + String containing the formatted result wrapped in tag + """ + # Always use xml_tag_name if it exists + if "xml_tag_name" in tool_call: + xml_tag_name = tool_call["xml_tag_name"] + return f" <{xml_tag_name}> {str(result)} " + + # Non-XML tool, just return the function result + function_name = tool_call["function_name"] + return f"Result for {function_name}: {str(result)}" + + def _create_tool_context(self, tool_call: Dict[str, Any], tool_index: int, assistant_message_id: Optional[str] = None, parsing_details: Optional[Dict[str, Any]] = None) -> ToolExecutionContext: + """Create a tool execution context with display name and parsing details populated.""" + context = ToolExecutionContext( + tool_call=tool_call, + tool_index=tool_index, + assistant_message_id=assistant_message_id, + parsing_details=parsing_details + ) + + # Set function_name and xml_tag_name fields + if "xml_tag_name" in tool_call: + context.xml_tag_name = tool_call["xml_tag_name"] + context.function_name = tool_call.get("function_name", tool_call["xml_tag_name"]) + else: + # For non-XML tools, use function name directly + context.function_name = tool_call.get("function_name", "unknown") + context.xml_tag_name = None + + return context + + async def _yield_and_save_tool_started(self, context: ToolExecutionContext, thread_id: str, thread_run_id: str) -> Optional[Dict[str, Any]]: + """Formats, saves, and returns a tool started status message.""" + tool_name = context.xml_tag_name or context.function_name + content = { + "role": "assistant", "status_type": "tool_started", + "function_name": context.function_name, "xml_tag_name": context.xml_tag_name, + "message": f"Starting execution of {tool_name}", "tool_index": context.tool_index, + "tool_call_id": context.tool_call.get("id") # Include tool_call ID if native + } + metadata = {"thread_run_id": thread_run_id} + saved_message_obj = await self.add_message( + thread_id=thread_id, type="status", content=content, is_llm_message=False, metadata=metadata + ) + return saved_message_obj # Return the full object (or None if saving failed) + + async def _yield_and_save_tool_completed(self, context: ToolExecutionContext, tool_message_id: Optional[str], thread_id: str, thread_run_id: str) -> Optional[Dict[str, Any]]: + """Formats, saves, and returns a tool completed/failed status message.""" + if not context.result: + # Delegate to error saving if result is missing (e.g., execution failed) + return await self._yield_and_save_tool_error(context, thread_id, thread_run_id) + + tool_name = context.xml_tag_name or context.function_name + status_type = "tool_completed" if context.result.success else "tool_failed" + message_text = f"Tool {tool_name} {'completed successfully' if context.result.success else 'failed'}" + + content = { + "role": "assistant", "status_type": status_type, + "function_name": context.function_name, "xml_tag_name": context.xml_tag_name, + "message": message_text, "tool_index": context.tool_index, + "tool_call_id": context.tool_call.get("id") + } + metadata = {"thread_run_id": thread_run_id} + # Add the *actual* tool result message ID to the metadata if available and successful + if context.result.success and tool_message_id: + metadata["linked_tool_result_message_id"] = tool_message_id + + # <<< ADDED: Signal if this is a terminating tool >>> + if context.function_name in ['ask', 'complete']: + metadata["agent_should_terminate"] = True + logger.info(f"Marking tool status for '{context.function_name}' with termination signal.") + self.trace.event(name="marking_tool_status_for_termination", level="DEFAULT", status_message=(f"Marking tool status for '{context.function_name}' with termination signal.")) + # <<< END ADDED >>> + + saved_message_obj = await self.add_message( + thread_id=thread_id, type="status", content=content, is_llm_message=False, metadata=metadata + ) + return saved_message_obj + + async def _yield_and_save_tool_error(self, context: ToolExecutionContext, thread_id: str, thread_run_id: str) -> Optional[Dict[str, Any]]: + """Formats, saves, and returns a tool error status message.""" + error_msg = str(context.error) if context.error else "Unknown error during tool execution" + tool_name = context.xml_tag_name or context.function_name + content = { + "role": "assistant", "status_type": "tool_error", + "function_name": context.function_name, "xml_tag_name": context.xml_tag_name, + "message": f"Error executing tool {tool_name}: {error_msg}", + "tool_index": context.tool_index, + "tool_call_id": context.tool_call.get("id") + } + metadata = {"thread_run_id": thread_run_id} + # Save the status message with is_llm_message=False + saved_message_obj = await self.add_message( + thread_id=thread_id, type="status", content=content, is_llm_message=False, metadata=metadata + ) + return saved_message_obj diff --git a/app/agentpress/thread_manager.py b/app/agentpress/thread_manager.py new file mode 100644 index 0000000..2eb963a --- /dev/null +++ b/app/agentpress/thread_manager.py @@ -0,0 +1,787 @@ +""" +Conversation thread management system for AgentPress. + +This module provides comprehensive conversation management, including: +- Thread creation and persistence +- Message handling with support for text and images +- Tool registration and execution +- LLM interaction with streaming support +- Error handling and cleanup +- Context summarization to manage token limits +""" + +import json +from typing import List, Dict, Any, Optional, Type, Union, AsyncGenerator, Literal +from services.llm import make_llm_api_call +from agentpress.tool import Tool +from agentpress.tool_registry import ToolRegistry +from agentpress.context_manager import ContextManager +from agentpress.response_processor import ( + ResponseProcessor, + ProcessorConfig +) +from services.supabase import DBConnection +from utils.logger import logger +from langfuse.client import StatefulGenerationClient, StatefulTraceClient +from services.langfuse import langfuse +import datetime +from litellm import token_counter + +# Type alias for tool choice +ToolChoice = Literal["auto", "required", "none"] + +class ThreadManager: + """Manages conversation threads with LLM models and tool execution. + + Provides comprehensive conversation management, handling message threading, + tool registration, and LLM interactions with support for both standard and + XML-based tool execution patterns. + """ + + def __init__(self, trace: Optional[StatefulTraceClient] = None, is_agent_builder: bool = False, target_agent_id: Optional[str] = None, agent_config: Optional[dict] = None): + """Initialize ThreadManager. + + Args: + trace: Optional trace client for logging + is_agent_builder: Whether this is an agent builder session + target_agent_id: ID of the agent being built (if in agent builder mode) + agent_config: Optional agent configuration with version information + """ + self.db = DBConnection() + self.tool_registry = ToolRegistry() + self.trace = trace + self.is_agent_builder = is_agent_builder + self.target_agent_id = target_agent_id + self.agent_config = agent_config + if not self.trace: + self.trace = langfuse.trace(name="anonymous:thread_manager") + self.response_processor = ResponseProcessor( + tool_registry=self.tool_registry, + add_message_callback=self.add_message, + trace=self.trace, + is_agent_builder=self.is_agent_builder, + target_agent_id=self.target_agent_id, + agent_config=self.agent_config + ) + self.context_manager = ContextManager() + + def _is_tool_result_message(self, msg: Dict[str, Any]) -> bool: + if not ("content" in msg and msg['content']): + return False + content = msg['content'] + if isinstance(content, str) and "ToolResult" in content: return True + if isinstance(content, dict) and "tool_execution" in content: return True + if isinstance(content, dict) and "interactive_elements" in content: return True + if isinstance(content, str): + try: + parsed_content = json.loads(content) + if isinstance(parsed_content, dict) and "tool_execution" in parsed_content: return True + if isinstance(parsed_content, dict) and "interactive_elements" in content: return True + except (json.JSONDecodeError, TypeError): + pass + return False + + def _compress_message(self, msg_content: Union[str, dict], message_id: Optional[str] = None, max_length: int = 3000) -> Union[str, dict]: + """Compress the message content.""" + # print("max_length", max_length) + if isinstance(msg_content, str): + if len(msg_content) > max_length: + return msg_content[:max_length] + "... (truncated)" + f"\n\nmessage_id \"{message_id}\"\nUse expand-message tool to see contents" + else: + return msg_content + elif isinstance(msg_content, dict): + if len(json.dumps(msg_content)) > max_length: + return json.dumps(msg_content)[:max_length] + "... (truncated)" + f"\n\nmessage_id \"{message_id}\"\nUse expand-message tool to see contents" + else: + return msg_content + + def _safe_truncate(self, msg_content: Union[str, dict], max_length: int = 100000) -> Union[str, dict]: + """Truncate the message content safely by removing the middle portion.""" + max_length = min(max_length, 100000) + if isinstance(msg_content, str): + if len(msg_content) > max_length: + # Calculate how much to keep from start and end + keep_length = max_length - 150 # Reserve space for truncation message + start_length = keep_length // 2 + end_length = keep_length - start_length + + start_part = msg_content[:start_length] + end_part = msg_content[-end_length:] if end_length > 0 else "" + + return start_part + f"\n\n... (middle truncated) ...\n\n" + end_part + f"\n\nThis message is too long, repeat relevant information in your response to remember it" + else: + return msg_content + elif isinstance(msg_content, dict): + json_str = json.dumps(msg_content) + if len(json_str) > max_length: + # Calculate how much to keep from start and end + keep_length = max_length - 150 # Reserve space for truncation message + start_length = keep_length // 2 + end_length = keep_length - start_length + + start_part = json_str[:start_length] + end_part = json_str[-end_length:] if end_length > 0 else "" + + return start_part + f"\n\n... (middle truncated) ...\n\n" + end_part + f"\n\nThis message is too long, repeat relevant information in your response to remember it" + else: + return msg_content + + def _compress_tool_result_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]: + """Compress the tool result messages except the most recent one.""" + uncompressed_total_token_count = token_counter(model=llm_model, messages=messages) + + if uncompressed_total_token_count > (max_tokens or (100 * 1000)): + _i = 0 # Count the number of ToolResult messages + for msg in reversed(messages): # Start from the end and work backwards + if self._is_tool_result_message(msg): # Only compress ToolResult messages + _i += 1 # Count the number of ToolResult messages + msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message + if msg_token_count > token_threshold: # If the message is too long + if _i > 1: # If this is not the most recent ToolResult message + message_id = msg.get('message_id') # Get the message_id + if message_id: + msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3) + else: + logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}") + else: + msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2)) + return messages + + def _compress_user_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]: + """Compress the user messages except the most recent one.""" + uncompressed_total_token_count = token_counter(model=llm_model, messages=messages) + + if uncompressed_total_token_count > (max_tokens or (100 * 1000)): + _i = 0 # Count the number of User messages + for msg in reversed(messages): # Start from the end and work backwards + if msg.get('role') == 'user': # Only compress User messages + _i += 1 # Count the number of User messages + msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message + if msg_token_count > token_threshold: # If the message is too long + if _i > 1: # If this is not the most recent User message + message_id = msg.get('message_id') # Get the message_id + if message_id: + msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3) + else: + logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}") + else: + msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2)) + return messages + + def _compress_assistant_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]: + """Compress the assistant messages except the most recent one.""" + uncompressed_total_token_count = token_counter(model=llm_model, messages=messages) + if uncompressed_total_token_count > (max_tokens or (100 * 1000)): + _i = 0 # Count the number of Assistant messages + for msg in reversed(messages): # Start from the end and work backwards + if msg.get('role') == 'assistant': # Only compress Assistant messages + _i += 1 # Count the number of Assistant messages + msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message + if msg_token_count > token_threshold: # If the message is too long + if _i > 1: # If this is not the most recent Assistant message + message_id = msg.get('message_id') # Get the message_id + if message_id: + msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3) + else: + logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}") + else: + msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2)) + + return messages + + + def _remove_meta_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Remove meta messages from the messages.""" + result: List[Dict[str, Any]] = [] + for msg in messages: + msg_content = msg.get('content') + # Try to parse msg_content as JSON if it's a string + if isinstance(msg_content, str): + try: msg_content = json.loads(msg_content) + except json.JSONDecodeError: pass + if isinstance(msg_content, dict): + # Create a copy to avoid modifying the original + msg_content_copy = msg_content.copy() + if "tool_execution" in msg_content_copy: + tool_execution = msg_content_copy["tool_execution"].copy() + if "arguments" in tool_execution: + del tool_execution["arguments"] + msg_content_copy["tool_execution"] = tool_execution + # Create a new message dict with the modified content + new_msg = msg.copy() + new_msg["content"] = json.dumps(msg_content_copy) + result.append(new_msg) + else: + result.append(msg) + return result + + def _compress_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int] = 41000, token_threshold: Optional[int] = 4096, max_iterations: int = 5) -> List[Dict[str, Any]]: + """Compress the messages. + token_threshold: must be a power of 2 + """ + + if 'sonnet' in llm_model.lower(): + max_tokens = 200 * 1000 - 64000 - 28000 + elif 'gpt' in llm_model.lower(): + max_tokens = 128 * 1000 - 28000 + elif 'gemini' in llm_model.lower(): + max_tokens = 1000 * 1000 - 300000 + elif 'deepseek' in llm_model.lower(): + max_tokens = 128 * 1000 - 28000 + else: + max_tokens = 41 * 1000 - 10000 + + result = messages + result = self._remove_meta_messages(result) + + uncompressed_total_token_count = token_counter(model=llm_model, messages=result) + + result = self._compress_tool_result_messages(result, llm_model, max_tokens, token_threshold) + result = self._compress_user_messages(result, llm_model, max_tokens, token_threshold) + result = self._compress_assistant_messages(result, llm_model, max_tokens, token_threshold) + + compressed_token_count = token_counter(model=llm_model, messages=result) + + logger.info(f"_compress_messages: {uncompressed_total_token_count} -> {compressed_token_count}") # Log the token compression for debugging later + + if max_iterations <= 0: + logger.warning(f"_compress_messages: Max iterations reached, omitting messages") + result = self._compress_messages_by_omitting_messages(messages, llm_model, max_tokens) + return result + + if (compressed_token_count > max_tokens): + logger.warning(f"Further token compression is needed: {compressed_token_count} > {max_tokens}") + result = self._compress_messages(messages, llm_model, max_tokens, int(token_threshold / 2), max_iterations - 1) + + return self._middle_out_messages(result) + + def _compress_messages_by_omitting_messages( + self, + messages: List[Dict[str, Any]], + llm_model: str, + max_tokens: Optional[int] = 41000, + removal_batch_size: int = 10, + min_messages_to_keep: int = 10 + ) -> List[Dict[str, Any]]: + """Compress the messages by omitting messages from the middle. + + Args: + messages: List of messages to compress + llm_model: Model name for token counting + max_tokens: Maximum allowed tokens + removal_batch_size: Number of messages to remove per iteration + min_messages_to_keep: Minimum number of messages to preserve + """ + if not messages: + return messages + + result = messages + result = self._remove_meta_messages(result) + + # Early exit if no compression needed + initial_token_count = token_counter(model=llm_model, messages=result) + max_allowed_tokens = max_tokens or (100 * 1000) + + if initial_token_count <= max_allowed_tokens: + return result + + # Separate system message (assumed to be first) from conversation messages + system_message = messages[0] if messages and messages[0].get('role') == 'system' else None + conversation_messages = result[1:] if system_message else result + + safety_limit = 500 + current_token_count = initial_token_count + + while current_token_count > max_allowed_tokens and safety_limit > 0: + safety_limit -= 1 + + if len(conversation_messages) <= min_messages_to_keep: + logger.warning(f"Cannot compress further: only {len(conversation_messages)} messages remain (min: {min_messages_to_keep})") + break + + # Calculate removal strategy based on current message count + if len(conversation_messages) > (removal_batch_size * 2): + # Remove from middle, keeping recent and early context + middle_start = len(conversation_messages) // 2 - (removal_batch_size // 2) + middle_end = middle_start + removal_batch_size + conversation_messages = conversation_messages[:middle_start] + conversation_messages[middle_end:] + else: + # Remove from earlier messages, preserving recent context + messages_to_remove = min(removal_batch_size, len(conversation_messages) // 2) + if messages_to_remove > 0: + conversation_messages = conversation_messages[messages_to_remove:] + else: + # Can't remove any more messages + break + + # Recalculate token count + messages_to_count = ([system_message] + conversation_messages) if system_message else conversation_messages + current_token_count = token_counter(model=llm_model, messages=messages_to_count) + + # Prepare final result + final_messages = ([system_message] + conversation_messages) if system_message else conversation_messages + final_token_count = token_counter(model=llm_model, messages=final_messages) + + logger.info(f"_compress_messages_by_omitting_messages: {initial_token_count} -> {final_token_count} tokens ({len(messages)} -> {len(final_messages)} messages)") + + return final_messages + + def _middle_out_messages(self, messages: List[Dict[str, Any]], max_messages: int = 320) -> List[Dict[str, Any]]: + """Remove messages from the middle of the list, keeping max_messages total.""" + if len(messages) <= max_messages: + return messages + + # Keep half from the beginning and half from the end + keep_start = max_messages // 2 + keep_end = max_messages - keep_start + + return messages[:keep_start] + messages[-keep_end:] + + + def add_tool(self, tool_class: Type[Tool], function_names: Optional[List[str]] = None, **kwargs): + """Add a tool to the ThreadManager.""" + self.tool_registry.register_tool(tool_class, function_names, **kwargs) + + async def add_message( + self, + thread_id: str, + type: str, + content: Union[Dict[str, Any], List[Any], str], + is_llm_message: bool = False, + metadata: Optional[Dict[str, Any]] = None, + agent_id: Optional[str] = None, + agent_version_id: Optional[str] = None + ): + """Add a message to the thread in the database. + + Args: + thread_id: The ID of the thread to add the message to. + type: The type of the message (e.g., 'text', 'image_url', 'tool_call', 'tool', 'user', 'assistant'). + content: The content of the message. Can be a dictionary, list, or string. + It will be stored as JSONB in the database. + is_llm_message: Flag indicating if the message originated from the LLM. + Defaults to False (user message). + metadata: Optional dictionary for additional message metadata. + Defaults to None, stored as an empty JSONB object if None. + agent_id: Optional ID of the agent associated with this message. + agent_version_id: Optional ID of the specific agent version used. + """ + logger.debug(f"Adding message of type '{type}' to thread {thread_id} (agent: {agent_id}, version: {agent_version_id})") + client = await self.db.client + + # Prepare data for insertion + data_to_insert = { + 'thread_id': thread_id, + 'type': type, + 'content': content, + 'is_llm_message': is_llm_message, + 'metadata': metadata or {}, + } + + # Add agent information if provided + if agent_id: + data_to_insert['agent_id'] = agent_id + if agent_version_id: + data_to_insert['agent_version_id'] = agent_version_id + + try: + # Add returning='representation' to get the inserted row data including the id + result = await client.table('messages').insert(data_to_insert, returning='representation').execute() + logger.info(f"Successfully added message to thread {thread_id}") + + if result.data and len(result.data) > 0 and isinstance(result.data[0], dict) and 'message_id' in result.data[0]: + return result.data[0] + else: + logger.error(f"Insert operation failed or did not return expected data structure for thread {thread_id}. Result data: {result.data}") + return None + except Exception as e: + logger.error(f"Failed to add message to thread {thread_id}: {str(e)}", exc_info=True) + raise + + async def get_llm_messages(self, thread_id: str) -> List[Dict[str, Any]]: + """Get all messages for a thread. + + This method uses the SQL function which handles context truncation + by considering summary messages. + + Args: + thread_id: The ID of the thread to get messages for. + + Returns: + List of message objects. + """ + logger.debug(f"Getting messages for thread {thread_id}") + client = await self.db.client + + try: + # result = await client.rpc('get_llm_formatted_messages', {'p_thread_id': thread_id}).execute() + + # Fetch messages in batches of 1000 to avoid overloading the database + all_messages = [] + batch_size = 1000 + offset = 0 + + while True: + result = await client.table('messages').select('message_id, content').eq('thread_id', thread_id).eq('is_llm_message', True).order('created_at').range(offset, offset + batch_size - 1).execute() + + if not result.data or len(result.data) == 0: + break + + all_messages.extend(result.data) + + # If we got fewer than batch_size records, we've reached the end + if len(result.data) < batch_size: + break + + offset += batch_size + + # Use all_messages instead of result.data in the rest of the method + result_data = all_messages + + # Parse the returned data which might be stringified JSON + if not result_data: + return [] + + # Return properly parsed JSON objects + messages = [] + for item in result_data: + if isinstance(item['content'], str): + try: + parsed_item = json.loads(item['content']) + parsed_item['message_id'] = item['message_id'] + messages.append(parsed_item) + except json.JSONDecodeError: + logger.error(f"Failed to parse message: {item['content']}") + else: + content = item['content'] + content['message_id'] = item['message_id'] + messages.append(content) + + return messages + + except Exception as e: + logger.error(f"Failed to get messages for thread {thread_id}: {str(e)}", exc_info=True) + return [] + + async def run_thread( + self, + thread_id: str, + system_prompt: Dict[str, Any], + stream: bool = True, + temporary_message: Optional[Dict[str, Any]] = None, + llm_model: str = "gpt-4o", + llm_temperature: float = 0, + llm_max_tokens: Optional[int] = None, + processor_config: Optional[ProcessorConfig] = None, + tool_choice: ToolChoice = "auto", + native_max_auto_continues: int = 25, + max_xml_tool_calls: int = 0, + include_xml_examples: bool = False, + enable_thinking: Optional[bool] = False, + reasoning_effort: Optional[str] = 'low', + enable_context_manager: bool = True, + generation: Optional[StatefulGenerationClient] = None, + ) -> Union[Dict[str, Any], AsyncGenerator]: + """Run a conversation thread with LLM integration and tool execution. + + Args: + thread_id: The ID of the thread to run + system_prompt: System message to set the assistant's behavior + stream: Use streaming API for the LLM response + temporary_message: Optional temporary user message for this run only + llm_model: The name of the LLM model to use + llm_temperature: Temperature parameter for response randomness (0-1) + llm_max_tokens: Maximum tokens in the LLM response + processor_config: Configuration for the response processor + tool_choice: Tool choice preference ("auto", "required", "none") + native_max_auto_continues: Maximum number of automatic continuations when + finish_reason="tool_calls" (0 disables auto-continue) + max_xml_tool_calls: Maximum number of XML tool calls to allow (0 = no limit) + include_xml_examples: Whether to include XML tool examples in the system prompt + enable_thinking: Whether to enable thinking before making a decision + reasoning_effort: The effort level for reasoning + enable_context_manager: Whether to enable automatic context summarization. + + Returns: + An async generator yielding response chunks or error dict + """ + + logger.info(f"Starting thread execution for thread {thread_id}") + logger.info(f"Using model: {llm_model}") + # Log parameters + logger.info(f"Parameters: model={llm_model}, temperature={llm_temperature}, max_tokens={llm_max_tokens}") + logger.info(f"Auto-continue: max={native_max_auto_continues}, XML tool limit={max_xml_tool_calls}") + + # Log model info + logger.info(f"🤖 Thread {thread_id}: Using model {llm_model}") + + # Apply max_xml_tool_calls if specified and not already set in config + if max_xml_tool_calls > 0 and not processor_config.max_xml_tool_calls: + processor_config.max_xml_tool_calls = max_xml_tool_calls + + # Create a working copy of the system prompt to potentially modify + working_system_prompt = system_prompt.copy() + + # Add XML examples to system prompt if requested, do this only ONCE before the loop + if include_xml_examples and processor_config.xml_tool_calling: + xml_examples = self.tool_registry.get_xml_examples() + if xml_examples: + examples_content = """ +--- XML TOOL CALLING --- + +In this environment you have access to a set of tools you can use to answer the user's question. The tools are specified in XML format. +Format your tool calls using the specified XML tags. Place parameters marked as 'attribute' within the opening tag (e.g., ``). Place parameters marked as 'content' between the opening and closing tags. Place parameters marked as 'element' within their own child tags (e.g., `value`). Refer to the examples provided below for the exact structure of each tool. +String and scalar parameters should be specified as attributes, while content goes between tags. +Note that spaces for string values are not stripped. The output is parsed with regular expressions. + +Here are the XML tools available with examples: +""" + for tag_name, example in xml_examples.items(): + examples_content += f"<{tag_name}> Example: {example}\\n" + + # # Save examples content to a file + # try: + # with open('xml_examples.txt', 'w') as f: + # f.write(examples_content) + # logger.debug("Saved XML examples to xml_examples.txt") + # except Exception as e: + # logger.error(f"Failed to save XML examples to file: {e}") + + system_content = working_system_prompt.get('content') + + if isinstance(system_content, str): + working_system_prompt['content'] += examples_content + logger.debug("Appended XML examples to string system prompt content.") + elif isinstance(system_content, list): + appended = False + for item in working_system_prompt['content']: # Modify the copy + if isinstance(item, dict) and item.get('type') == 'text' and 'text' in item: + item['text'] += examples_content + logger.debug("Appended XML examples to the first text block in list system prompt content.") + appended = True + break + if not appended: + logger.warning("System prompt content is a list but no text block found to append XML examples.") + else: + logger.warning(f"System prompt content is of unexpected type ({type(system_content)}), cannot add XML examples.") + # Control whether we need to auto-continue due to tool_calls finish reason + auto_continue = True + auto_continue_count = 0 + + # Define inner function to handle a single run + async def _run_once(temp_msg=None): + try: + # Ensure processor_config is available in this scope + nonlocal processor_config + # Note: processor_config is now guaranteed to exist due to check above + + # 1. Get messages from thread for LLM call + messages = await self.get_llm_messages(thread_id) + + # 2. Check token count before proceeding + token_count = 0 + try: + # Use the potentially modified working_system_prompt for token counting + token_count = token_counter(model=llm_model, messages=[working_system_prompt] + messages) + token_threshold = self.context_manager.token_threshold + logger.info(f"Thread {thread_id} token count: {token_count}/{token_threshold} ({(token_count/token_threshold)*100:.1f}%)") + + # if token_count >= token_threshold and enable_context_manager: + # logger.info(f"Thread token count ({token_count}) exceeds threshold ({token_threshold}), summarizing...") + # summarized = await self.context_manager.check_and_summarize_if_needed( + # thread_id=thread_id, + # add_message_callback=self.add_message, + # model=llm_model, + # force=True + # ) + # if summarized: + # logger.info("Summarization complete, fetching updated messages with summary") + # messages = await self.get_llm_messages(thread_id) + # # Recount tokens after summarization, using the modified prompt + # new_token_count = token_counter(model=llm_model, messages=[working_system_prompt] + messages) + # logger.info(f"After summarization: token count reduced from {token_count} to {new_token_count}") + # else: + # logger.warning("Summarization failed or wasn't needed - proceeding with original messages") + # elif not enable_context_manager: + # logger.info("Automatic summarization disabled. Skipping token count check and summarization.") + + except Exception as e: + logger.error(f"Error counting tokens or summarizing: {str(e)}") + + # 3. Prepare messages for LLM call + add temporary message if it exists + # Use the working_system_prompt which may contain the XML examples + prepared_messages = [working_system_prompt] + + # Find the last user message index + last_user_index = -1 + for i, msg in enumerate(messages): + if msg.get('role') == 'user': + last_user_index = i + + # Insert temporary message before the last user message if it exists + if temp_msg and last_user_index >= 0: + prepared_messages.extend(messages[:last_user_index]) + prepared_messages.append(temp_msg) + prepared_messages.extend(messages[last_user_index:]) + logger.debug("Added temporary message before the last user message") + else: + # If no user message or no temporary message, just add all messages + prepared_messages.extend(messages) + if temp_msg: + prepared_messages.append(temp_msg) + logger.debug("Added temporary message to the end of prepared messages") + + # 4. Prepare tools for LLM call + openapi_tool_schemas = None + if processor_config.native_tool_calling: + openapi_tool_schemas = self.tool_registry.get_openapi_schemas() + logger.debug(f"Retrieved {len(openapi_tool_schemas) if openapi_tool_schemas else 0} OpenAPI tool schemas") + + prepared_messages = self._compress_messages(prepared_messages, llm_model) + + # 5. Make LLM API call + logger.debug("Making LLM API call") + try: + if generation: + generation.update( + input=prepared_messages, + start_time=datetime.datetime.now(datetime.timezone.utc), + model=llm_model, + model_parameters={ + "max_tokens": llm_max_tokens, + "temperature": llm_temperature, + "enable_thinking": enable_thinking, + "reasoning_effort": reasoning_effort, + "tool_choice": tool_choice, + "tools": openapi_tool_schemas, + } + ) + llm_response = await make_llm_api_call( + prepared_messages, # Pass the potentially modified messages + llm_model, + temperature=llm_temperature, + max_tokens=llm_max_tokens, + tools=openapi_tool_schemas, + tool_choice=tool_choice if processor_config.native_tool_calling else None, + stream=stream, + enable_thinking=enable_thinking, + reasoning_effort=reasoning_effort + ) + logger.debug("Successfully received raw LLM API response stream/object") + + except Exception as e: + logger.error(f"Failed to make LLM API call: {str(e)}", exc_info=True) + raise + + # 6. Process LLM response using the ResponseProcessor + if stream: + logger.debug("Processing streaming response") + response_generator = self.response_processor.process_streaming_response( + llm_response=llm_response, + thread_id=thread_id, + config=processor_config, + prompt_messages=prepared_messages, + llm_model=llm_model, + ) + + return response_generator + else: + logger.debug("Processing non-streaming response") + # Pass through the response generator without try/except to let errors propagate up + response_generator = self.response_processor.process_non_streaming_response( + llm_response=llm_response, + thread_id=thread_id, + config=processor_config, + prompt_messages=prepared_messages, + llm_model=llm_model, + ) + return response_generator # Return the generator + + except Exception as e: + logger.error(f"Error in run_thread: {str(e)}", exc_info=True) + # Return the error as a dict to be handled by the caller + return { + "type": "status", + "status": "error", + "message": str(e) + } + + # Define a wrapper generator that handles auto-continue logic + async def auto_continue_wrapper(): + nonlocal auto_continue, auto_continue_count + + while auto_continue and (native_max_auto_continues == 0 or auto_continue_count < native_max_auto_continues): + # Reset auto_continue for this iteration + auto_continue = False + + # Run the thread once, passing the potentially modified system prompt + # Pass temp_msg only on the first iteration + try: + response_gen = await _run_once(temporary_message if auto_continue_count == 0 else None) + + # Handle error responses + if isinstance(response_gen, dict) and "status" in response_gen and response_gen["status"] == "error": + logger.error(f"Error in auto_continue_wrapper: {response_gen.get('message', 'Unknown error')}") + yield response_gen + return # Exit the generator on error + + # Process each chunk + try: + async for chunk in response_gen: + # Check if this is a finish reason chunk with tool_calls or xml_tool_limit_reached + if chunk.get('type') == 'finish': + if chunk.get('finish_reason') == 'tool_calls': + # Only auto-continue if enabled (max > 0) + if native_max_auto_continues > 0: + logger.info(f"Detected finish_reason='tool_calls', auto-continuing ({auto_continue_count + 1}/{native_max_auto_continues})") + auto_continue = True + auto_continue_count += 1 + # Don't yield the finish chunk to avoid confusing the client + continue + elif chunk.get('finish_reason') == 'xml_tool_limit_reached': + # Don't auto-continue if XML tool limit was reached + logger.info(f"Detected finish_reason='xml_tool_limit_reached', stopping auto-continue") + auto_continue = False + # Still yield the chunk to inform the client + + # Otherwise just yield the chunk normally + yield chunk + + # If not auto-continuing, we're done + if not auto_continue: + break + except Exception as e: + # If there's an exception, log it, yield an error status, and stop execution + logger.error(f"Error in auto_continue_wrapper generator: {str(e)}", exc_info=True) + yield { + "type": "status", + "status": "error", + "message": f"Error in thread processing: {str(e)}" + } + return # Exit the generator on any error + except Exception as outer_e: + # Catch exceptions from _run_once itself + logger.error(f"Error executing thread: {str(outer_e)}", exc_info=True) + yield { + "type": "status", + "status": "error", + "message": f"Error executing thread: {str(outer_e)}" + } + return # Exit immediately on exception from _run_once + + # If we've reached the max auto-continues, log a warning + if auto_continue and auto_continue_count >= native_max_auto_continues: + logger.warning(f"Reached maximum auto-continue limit ({native_max_auto_continues}), stopping.") + yield { + "type": "content", + "content": f"\n[Agent reached maximum auto-continue limit of {native_max_auto_continues}]" + } + + # If auto-continue is disabled (max=0), just run once + if native_max_auto_continues == 0: + logger.info("Auto-continue is disabled (native_max_auto_continues=0)") + # Pass the potentially modified system prompt and temp message + return await _run_once(temporary_message) + + # Otherwise return the auto-continue wrapper generator + return auto_continue_wrapper() diff --git a/app/agentpress/tool.py b/app/agentpress/tool.py new file mode 100644 index 0000000..de7a504 --- /dev/null +++ b/app/agentpress/tool.py @@ -0,0 +1,240 @@ +""" +Core tool system providing the foundation for creating and managing tools. + +This module defines the base classes and decorators for creating tools in AgentPress: +- Tool base class for implementing tool functionality +- Schema decorators for OpenAPI and XML tool definitions +- Result containers for standardized tool outputs +""" + +from typing import Dict, Any, Union, Optional, List +from dataclasses import dataclass, field +from abc import ABC +import json +import inspect +from enum import Enum +from utils.logger import logger + +class SchemaType(Enum): + """Enumeration of supported schema types for tool definitions.""" + OPENAPI = "openapi" + XML = "xml" + CUSTOM = "custom" + +@dataclass +class XMLNodeMapping: + """Maps an XML node to a function parameter. + + Attributes: + param_name (str): Name of the function parameter + node_type (str): Type of node ("element", "attribute", or "content") + path (str): XPath-like path to the node ("." means root element) + required (bool): Whether the parameter is required (defaults to True) + """ + param_name: str + node_type: str = "element" + path: str = "." + required: bool = True + +@dataclass +class XMLTagSchema: + """Schema definition for XML tool tags. + + Attributes: + tag_name (str): Root tag name for the tool + mappings (List[XMLNodeMapping]): Parameter mappings for the tag + example (str, optional): Example showing tag usage + + Methods: + add_mapping: Add a new parameter mapping to the schema + """ + tag_name: str + mappings: List[XMLNodeMapping] = field(default_factory=list) + example: Optional[str] = None + + def add_mapping(self, param_name: str, node_type: str = "element", path: str = ".", required: bool = True) -> None: + """Add a new node mapping to the schema. + + Args: + param_name: Name of the function parameter + node_type: Type of node ("element", "attribute", or "content") + path: XPath-like path to the node + required: Whether the parameter is required + """ + self.mappings.append(XMLNodeMapping( + param_name=param_name, + node_type=node_type, + path=path, + required=required + )) + logger.debug(f"Added XML mapping for parameter '{param_name}' with type '{node_type}' at path '{path}', required={required}") + +@dataclass +class ToolSchema: + """Container for tool schemas with type information. + + Attributes: + schema_type (SchemaType): Type of schema (OpenAPI, XML, or Custom) + schema (Dict[str, Any]): The actual schema definition + xml_schema (XMLTagSchema, optional): XML-specific schema if applicable + """ + schema_type: SchemaType + schema: Dict[str, Any] + xml_schema: Optional[XMLTagSchema] = None + +@dataclass +class ToolResult: + """Container for tool execution results. + + Attributes: + success (bool): Whether the tool execution succeeded + output (str): Output message or error description + """ + success: bool + output: str + +class Tool(ABC): + """Abstract base class for all tools. + + Provides the foundation for implementing tools with schema registration + and result handling capabilities. + + Attributes: + _schemas (Dict[str, List[ToolSchema]]): Registered schemas for tool methods + + Methods: + get_schemas: Get all registered tool schemas + success_response: Create a successful result + fail_response: Create a failed result + """ + + def __init__(self): + """Initialize tool with empty schema registry.""" + self._schemas: Dict[str, List[ToolSchema]] = {} + logger.debug(f"Initializing tool class: {self.__class__.__name__}") + self._register_schemas() + + def _register_schemas(self): + """Register schemas from all decorated methods.""" + for name, method in inspect.getmembers(self, predicate=inspect.ismethod): + if hasattr(method, 'tool_schemas'): + self._schemas[name] = method.tool_schemas + logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}") + + def get_schemas(self) -> Dict[str, List[ToolSchema]]: + """Get all registered tool schemas. + + Returns: + Dict mapping method names to their schema definitions + """ + return self._schemas + + def success_response(self, data: Union[Dict[str, Any], str]) -> ToolResult: + """Create a successful tool result. + + Args: + data: Result data (dictionary or string) + + Returns: + ToolResult with success=True and formatted output + """ + if isinstance(data, str): + text = data + else: + text = json.dumps(data, indent=2) + logger.debug(f"Created success response for {self.__class__.__name__}") + return ToolResult(success=True, output=text) + + def fail_response(self, msg: str) -> ToolResult: + """Create a failed tool result. + + Args: + msg: Error message describing the failure + + Returns: + ToolResult with success=False and error message + """ + logger.debug(f"Tool {self.__class__.__name__} returned failed result: {msg}") + return ToolResult(success=False, output=msg) + +def _add_schema(func, schema: ToolSchema): + """Helper to add schema to a function.""" + if not hasattr(func, 'tool_schemas'): + func.tool_schemas = [] + func.tool_schemas.append(schema) + logger.debug(f"Added {schema.schema_type.value} schema to function {func.__name__}") + return func + +def openapi_schema(schema: Dict[str, Any]): + """Decorator for OpenAPI schema tools.""" + def decorator(func): + logger.debug(f"Applying OpenAPI schema to function {func.__name__}") + return _add_schema(func, ToolSchema( + schema_type=SchemaType.OPENAPI, + schema=schema + )) + return decorator + +def xml_schema( + tag_name: str, + mappings: List[Dict[str, Any]] = None, + example: str = None +): + """ + Decorator for XML schema tools with improved node mapping. + + Args: + tag_name: Name of the root XML tag + mappings: List of mapping definitions, each containing: + - param_name: Name of the function parameter + - node_type: "element", "attribute", or "content" + - path: Path to the node (default "." for root) + - required: Whether the parameter is required (default True) + example: Optional example showing how to use the XML tag + + Example: + @xml_schema( + tag_name="str-replace", + mappings=[ + {"param_name": "file_path", "node_type": "attribute", "path": "."}, + {"param_name": "old_str", "node_type": "element", "path": "old_str"}, + {"param_name": "new_str", "node_type": "element", "path": "new_str"} + ], + example=''' + + text to replace + replacement text + + ''' + ) + """ + def decorator(func): + logger.debug(f"Applying XML schema with tag '{tag_name}' to function {func.__name__}") + xml_schema = XMLTagSchema(tag_name=tag_name, example=example) + + # Add mappings + if mappings: + for mapping in mappings: + xml_schema.add_mapping( + param_name=mapping["param_name"], + node_type=mapping.get("node_type", "element"), + path=mapping.get("path", "."), + required=mapping.get("required", True) + ) + + return _add_schema(func, ToolSchema( + schema_type=SchemaType.XML, + schema={}, # OpenAPI schema could be added here if needed + xml_schema=xml_schema + )) + return decorator + +def custom_schema(schema: Dict[str, Any]): + """Decorator for custom schema tools.""" + def decorator(func): + logger.debug(f"Applying custom schema to function {func.__name__}") + return _add_schema(func, ToolSchema( + schema_type=SchemaType.CUSTOM, + schema=schema + )) + return decorator diff --git a/app/agentpress/tool_registry.py b/app/agentpress/tool_registry.py new file mode 100644 index 0000000..b50438a --- /dev/null +++ b/app/agentpress/tool_registry.py @@ -0,0 +1,152 @@ +from typing import Dict, Type, Any, List, Optional, Callable +from agentpress.tool import Tool, SchemaType +from utils.logger import logger + + +class ToolRegistry: + """Registry for managing and accessing tools. + + Maintains a collection of tool instances and their schemas, allowing for + selective registration of tool functions and easy access to tool capabilities. + + Attributes: + tools (Dict[str, Dict[str, Any]]): OpenAPI-style tools and schemas + xml_tools (Dict[str, Dict[str, Any]]): XML-style tools and schemas + + Methods: + register_tool: Register a tool with optional function filtering + get_tool: Get a specific tool by name + get_xml_tool: Get a tool by XML tag name + get_openapi_schemas: Get OpenAPI schemas for function calling + get_xml_examples: Get examples of XML tool usage + """ + + def __init__(self): + """Initialize a new ToolRegistry instance.""" + self.tools = {} + self.xml_tools = {} + logger.debug("Initialized new ToolRegistry instance") + + def register_tool(self, tool_class: Type[Tool], function_names: Optional[List[str]] = None, **kwargs): + """Register a tool with optional function filtering. + + Args: + tool_class: The tool class to register + function_names: Optional list of specific functions to register + **kwargs: Additional arguments passed to tool initialization + + Notes: + - If function_names is None, all functions are registered + - Handles both OpenAPI and XML schema registration + """ + logger.debug(f"Registering tool class: {tool_class.__name__}") + tool_instance = tool_class(**kwargs) + schemas = tool_instance.get_schemas() + + logger.debug(f"Available schemas for {tool_class.__name__}: {list(schemas.keys())}") + + registered_openapi = 0 + registered_xml = 0 + + for func_name, schema_list in schemas.items(): + if function_names is None or func_name in function_names: + for schema in schema_list: + if schema.schema_type == SchemaType.OPENAPI: + self.tools[func_name] = { + "instance": tool_instance, + "schema": schema + } + registered_openapi += 1 + logger.debug(f"Registered OpenAPI function {func_name} from {tool_class.__name__}") + + if schema.schema_type == SchemaType.XML and schema.xml_schema: + self.xml_tools[schema.xml_schema.tag_name] = { + "instance": tool_instance, + "method": func_name, + "schema": schema + } + registered_xml += 1 + logger.debug(f"Registered XML tag {schema.xml_schema.tag_name} -> {func_name} from {tool_class.__name__}") + + logger.debug(f"Tool registration complete for {tool_class.__name__}: {registered_openapi} OpenAPI functions, {registered_xml} XML tags") + + def get_available_functions(self) -> Dict[str, Callable]: + """Get all available tool functions. + + Returns: + Dict mapping function names to their implementations + """ + available_functions = {} + + # Get OpenAPI tool functions + for tool_name, tool_info in self.tools.items(): + tool_instance = tool_info['instance'] + function_name = tool_name + function = getattr(tool_instance, function_name) + available_functions[function_name] = function + + # Get XML tool functions + for tag_name, tool_info in self.xml_tools.items(): + tool_instance = tool_info['instance'] + method_name = tool_info['method'] + function = getattr(tool_instance, method_name) + available_functions[method_name] = function + + logger.debug(f"Retrieved {len(available_functions)} available functions") + return available_functions + + def get_tool(self, tool_name: str) -> Dict[str, Any]: + """Get a specific tool by name. + + Args: + tool_name: Name of the tool function + + Returns: + Dict containing tool instance and schema, or empty dict if not found + """ + tool = self.tools.get(tool_name, {}) + if not tool: + logger.warning(f"Tool not found: {tool_name}") + return tool + + def get_xml_tool(self, tag_name: str) -> Dict[str, Any]: + """Get tool info by XML tag name. + + Args: + tag_name: XML tag name for the tool + + Returns: + Dict containing tool instance, method name, and schema + """ + tool = self.xml_tools.get(tag_name, {}) + if not tool: + logger.warning(f"XML tool not found for tag: {tag_name}") + return tool + + def get_openapi_schemas(self) -> List[Dict[str, Any]]: + """Get OpenAPI schemas for function calling. + + Returns: + List of OpenAPI-compatible schema definitions + """ + schemas = [ + tool_info['schema'].schema + for tool_info in self.tools.values() + if tool_info['schema'].schema_type == SchemaType.OPENAPI + ] + logger.debug(f"Retrieved {len(schemas)} OpenAPI schemas") + return schemas + + def get_xml_examples(self) -> Dict[str, str]: + """Get all XML tag examples. + + Returns: + Dict mapping tag names to their example usage + """ + examples = {} + for tool_info in self.xml_tools.values(): + schema = tool_info['schema'] + if schema.xml_schema and schema.xml_schema.example: + examples[schema.xml_schema.tag_name] = schema.xml_schema.example + logger.debug(f"Retrieved {len(examples)} XML examples") + return examples diff --git a/app/agentpress/utils/__init__.py b/app/agentpress/utils/__init__.py new file mode 100644 index 0000000..e0dabab --- /dev/null +++ b/app/agentpress/utils/__init__.py @@ -0,0 +1 @@ +# Utils module for AgentPress \ No newline at end of file diff --git a/app/agentpress/utils/json_helpers.py b/app/agentpress/utils/json_helpers.py new file mode 100644 index 0000000..3bda909 --- /dev/null +++ b/app/agentpress/utils/json_helpers.py @@ -0,0 +1,174 @@ +""" +JSON helper utilities for handling both legacy (string) and new (dict/list) formats. + +These utilities help with the transition from storing JSON as strings to storing +them as proper JSONB objects in the database. +""" + +import json +from typing import Any, Union, Dict, List + + +def ensure_dict(value: Union[str, Dict[str, Any], None], default: Dict[str, Any] = None) -> Dict[str, Any]: + """ + Ensure a value is a dictionary. + + Handles: + - None -> returns default or {} + - Dict -> returns as-is + - JSON string -> parses and returns dict + - Other -> returns default or {} + + Args: + value: The value to ensure is a dict + default: Default value if conversion fails + + Returns: + A dictionary + """ + if default is None: + default = {} + + if value is None: + return default + + if isinstance(value, dict): + return value + + if isinstance(value, str): + try: + parsed = json.loads(value) + if isinstance(parsed, dict): + return parsed + return default + except (json.JSONDecodeError, TypeError): + return default + + return default + + +def ensure_list(value: Union[str, List[Any], None], default: List[Any] = None) -> List[Any]: + """ + Ensure a value is a list. + + Handles: + - None -> returns default or [] + - List -> returns as-is + - JSON string -> parses and returns list + - Other -> returns default or [] + + Args: + value: The value to ensure is a list + default: Default value if conversion fails + + Returns: + A list + """ + if default is None: + default = [] + + if value is None: + return default + + if isinstance(value, list): + return value + + if isinstance(value, str): + try: + parsed = json.loads(value) + if isinstance(parsed, list): + return parsed + return default + except (json.JSONDecodeError, TypeError): + return default + + return default + + +def safe_json_parse(value: Union[str, Dict, List, Any], default: Any = None) -> Any: + """ + Safely parse a value that might be JSON string or already parsed. + + This handles the transition period where some data might be stored as + JSON strings (old format) and some as proper objects (new format). + + Args: + value: The value to parse + default: Default value if parsing fails + + Returns: + Parsed value or default + """ + if value is None: + return default + + # If it's already a dict or list, return as-is + if isinstance(value, (dict, list)): + return value + + # If it's a string, try to parse it + if isinstance(value, str): + try: + return json.loads(value) + except (json.JSONDecodeError, TypeError): + # If it's not valid JSON, return the string itself + return value + + # For any other type, return as-is + return value + + +def to_json_string(value: Any) -> str: + """ + Convert a value to a JSON string if needed. + + This is used for backwards compatibility when yielding data that + expects JSON strings. + + Args: + value: The value to convert + + Returns: + JSON string representation + """ + if isinstance(value, str): + # If it's already a string, check if it's valid JSON + try: + json.loads(value) + return value # It's already a JSON string + except (json.JSONDecodeError, TypeError): + # It's a plain string, encode it as JSON + return json.dumps(value) + + # For all other types, convert to JSON + return json.dumps(value) + + +def format_for_yield(message_object: Dict[str, Any]) -> Dict[str, Any]: + """ + Format a message object for yielding, ensuring content and metadata are JSON strings. + + This maintains backward compatibility with clients expecting JSON strings + while the database now stores proper objects. + + Args: + message_object: The message object from the database + + Returns: + Message object with content and metadata as JSON strings + """ + if not message_object: + return message_object + + # Create a copy to avoid modifying the original + formatted = message_object.copy() + + # Ensure content is a JSON string + if 'content' in formatted and not isinstance(formatted['content'], str): + formatted['content'] = json.dumps(formatted['content']) + + # Ensure metadata is a JSON string + if 'metadata' in formatted and not isinstance(formatted['metadata'], str): + formatted['metadata'] = json.dumps(formatted['metadata']) + + return formatted \ No newline at end of file diff --git a/app/agentpress/xml_tool_parser.py b/app/agentpress/xml_tool_parser.py new file mode 100644 index 0000000..e882a66 --- /dev/null +++ b/app/agentpress/xml_tool_parser.py @@ -0,0 +1,300 @@ +""" +XML Tool Call Parser Module + +This module provides a reliable XML tool call parsing system that supports +the Cursor-style format with structured function_calls blocks. +""" + +import re +import xml.etree.ElementTree as ET +from typing import List, Dict, Any, Optional, Tuple +from dataclasses import dataclass +import json +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class XMLToolCall: + """Represents a parsed XML tool call.""" + function_name: str + parameters: Dict[str, Any] + raw_xml: str + parsing_details: Dict[str, Any] + + +class XMLToolParser: + """ + Parser for XML tool calls using the Cursor-style format: + + + + param_value + ... + + + """ + + # Regex patterns for extracting XML blocks + FUNCTION_CALLS_PATTERN = re.compile( + r'(.*?)', + re.DOTALL | re.IGNORECASE + ) + + INVOKE_PATTERN = re.compile( + r'(.*?)', + re.DOTALL | re.IGNORECASE + ) + + PARAMETER_PATTERN = re.compile( + r'(.*?)', + re.DOTALL | re.IGNORECASE + ) + + def __init__(self, strict_mode: bool = False): + """ + Initialize the XML tool parser. + + Args: + strict_mode: If True, only accept the exact format. If False, + also try to parse legacy formats for backwards compatibility. + """ + self.strict_mode = strict_mode + + def parse_content(self, content: str) -> List[XMLToolCall]: + """ + Parse XML tool calls from content. + + Args: + content: The text content potentially containing XML tool calls + + Returns: + List of parsed XMLToolCall objects + """ + tool_calls = [] + + # First, try to find function_calls blocks + function_calls_matches = self.FUNCTION_CALLS_PATTERN.findall(content) + + for fc_content in function_calls_matches: + # Find all invoke blocks within this function_calls block + invoke_matches = self.INVOKE_PATTERN.findall(fc_content) + + for function_name, invoke_content in invoke_matches: + try: + tool_call = self._parse_invoke_block( + function_name, + invoke_content, + fc_content + ) + if tool_call: + tool_calls.append(tool_call) + except Exception as e: + logger.error(f"Error parsing invoke block for {function_name}: {e}") + + # If not in strict mode and no tool calls found, try legacy format + if not self.strict_mode and not tool_calls: + tool_calls.extend(self._parse_legacy_format(content)) + + return tool_calls + + def _parse_invoke_block( + self, + function_name: str, + invoke_content: str, + full_block: str + ) -> Optional[XMLToolCall]: + """Parse a single invoke block into an XMLToolCall.""" + parameters = {} + parsing_details = { + "format": "v2", + "function_name": function_name, + "raw_parameters": {} + } + + # Extract all parameters + param_matches = self.PARAMETER_PATTERN.findall(invoke_content) + + for param_name, param_value in param_matches: + # Clean up the parameter value + param_value = param_value.strip() + + # Try to parse as JSON if it looks like JSON + parsed_value = self._parse_parameter_value(param_value) + + parameters[param_name] = parsed_value + parsing_details["raw_parameters"][param_name] = param_value + + # Extract the raw XML for this specific invoke + invoke_pattern = re.compile( + rf'.*?', + re.DOTALL | re.IGNORECASE + ) + raw_xml_match = invoke_pattern.search(full_block) + raw_xml = raw_xml_match.group(0) if raw_xml_match else f"..." + + return XMLToolCall( + function_name=function_name, + parameters=parameters, + raw_xml=raw_xml, + parsing_details=parsing_details + ) + + def _parse_parameter_value(self, value: str) -> Any: + """ + Parse a parameter value, attempting to convert to appropriate type. + + Args: + value: The string value to parse + + Returns: + Parsed value (could be dict, list, bool, int, float, or str) + """ + value = value.strip() + + # Try to parse as JSON first + if value.startswith(('{', '[')): + try: + return json.loads(value) + except json.JSONDecodeError: + pass + + # Try to parse as boolean + if value.lower() in ('true', 'false'): + return value.lower() == 'true' + + # Try to parse as number + try: + if '.' in value: + return float(value) + else: + return int(value) + except ValueError: + pass + + # Return as string + return value + + def _parse_legacy_format(self, content: str) -> List[XMLToolCall]: + """ + Parse legacy XML tool formats for backwards compatibility. + This handles formats like ... or + ... + """ + tool_calls = [] + + # Pattern for finding XML-like tags + tag_pattern = re.compile(r'<([a-zA-Z][\w\-]*)((?:\s+[\w\-]+=["\'][^"\']*["\'])*)\s*>(.*?)', re.DOTALL) + + for match in tag_pattern.finditer(content): + tag_name = match.group(1) + attributes_str = match.group(2) + inner_content = match.group(3) + + # Skip our own format tags + if tag_name in ('function_calls', 'invoke', 'parameter'): + continue + + parameters = {} + parsing_details = { + "format": "legacy", + "tag_name": tag_name, + "attributes": {}, + "inner_content": inner_content.strip() + } + + # Parse attributes + if attributes_str: + attr_pattern = re.compile(r'([\w\-]+)=["\']([^"\']*)["\']') + for attr_match in attr_pattern.finditer(attributes_str): + attr_name = attr_match.group(1) + attr_value = attr_match.group(2) + parameters[attr_name] = self._parse_parameter_value(attr_value) + parsing_details["attributes"][attr_name] = attr_value + + # If there's inner content and no attributes, use it as a 'content' parameter + if inner_content.strip() and not parameters: + parameters['content'] = inner_content.strip() + + # Convert tag name to function name (e.g., create-file -> create_file) + function_name = tag_name.replace('-', '_') + + tool_calls.append(XMLToolCall( + function_name=function_name, + parameters=parameters, + raw_xml=match.group(0), + parsing_details=parsing_details + )) + + return tool_calls + + def format_tool_call(self, function_name: str, parameters: Dict[str, Any]) -> str: + """ + Format a tool call in the Cursor-style XML format. + + Args: + function_name: Name of the function to call + parameters: Dictionary of parameters + + Returns: + Formatted XML string + """ + lines = ['', ''.format(function_name)] + + for param_name, param_value in parameters.items(): + # Convert value to string representation + if isinstance(param_value, (dict, list)): + value_str = json.dumps(param_value) + elif isinstance(param_value, bool): + value_str = str(param_value).lower() + else: + value_str = str(param_value) + + lines.append('{}'.format( + param_name, value_str + )) + + lines.extend(['', '']) + return '\n'.join(lines) + + def validate_tool_call(self, tool_call: XMLToolCall, expected_params: Optional[Dict[str, type]] = None) -> Tuple[bool, Optional[str]]: + """ + Validate a tool call against expected parameters. + + Args: + tool_call: The XMLToolCall to validate + expected_params: Optional dict of parameter names to expected types + + Returns: + Tuple of (is_valid, error_message) + """ + if not tool_call.function_name: + return False, "Function name is required" + + if expected_params: + for param_name, expected_type in expected_params.items(): + if param_name not in tool_call.parameters: + return False, f"Missing required parameter: {param_name}" + + param_value = tool_call.parameters[param_name] + if not isinstance(param_value, expected_type): + return False, f"Parameter {param_name} should be of type {expected_type.__name__}" + + return True, None + + +# Convenience function for quick parsing +def parse_xml_tool_calls(content: str, strict_mode: bool = False) -> List[XMLToolCall]: + """ + Parse XML tool calls from content. + + Args: + content: The text content potentially containing XML tool calls + strict_mode: If True, only accept the Cursor-style format + + Returns: + List of parsed XMLToolCall objects + """ + parser = XMLToolParser(strict_mode=strict_mode) + return parser.parse_content(content) \ No newline at end of file diff --git a/app/daytona/tool_base.py b/app/daytona/tool_base.py index 90446e4..70cb02d 100644 --- a/app/daytona/tool_base.py +++ b/app/daytona/tool_base.py @@ -4,7 +4,7 @@ from typing import Optional from agentpress.thread_manager import ThreadManager from app.tool.base import BaseTool, ToolResult from daytona_sdk import Sandbox -from sandbox.sandbox import get_or_start_sandbox +from daytona.sandbox import get_or_start_sandbox from utils.logger import logger from utils.files_utils import clean_path diff --git a/app/services/billing.py b/app/services/billing.py new file mode 100644 index 0000000..67fa9bb --- /dev/null +++ b/app/services/billing.py @@ -0,0 +1,969 @@ +""" +Stripe Billing API implementation for Suna on top of Basejump. ONLY HAS SUPPOT FOR USER ACCOUNTS – no team accounts. As we are using the user_id as account_id as is the case with personal accounts. In personal accounts, the account_id equals the user_id. In team accounts, the account_id is unique. + +stripe listen --forward-to localhost:8000/api/billing/webhook +""" + +from fastapi import APIRouter, HTTPException, Depends, Request +from typing import Optional, Dict, Tuple +import stripe +from datetime import datetime, timezone +from utils.logger import logger +from utils.config import config, EnvMode +from services.supabase import DBConnection +from utils.auth_utils import get_current_user_id_from_jwt +from pydantic import BaseModel +from utils.constants import MODEL_ACCESS_TIERS, MODEL_NAME_ALIASES +import os + +# Initialize Stripe +stripe.api_key = config.STRIPE_SECRET_KEY + +# Initialize router +router = APIRouter(prefix="/billing", tags=["billing"]) + + +SUBSCRIPTION_TIERS = { + config.STRIPE_FREE_TIER_ID: {'name': 'free', 'minutes': 60}, + config.STRIPE_TIER_2_20_ID: {'name': 'tier_2_20', 'minutes': 120}, # 2 hours + config.STRIPE_TIER_6_50_ID: {'name': 'tier_6_50', 'minutes': 360}, # 6 hours + config.STRIPE_TIER_12_100_ID: {'name': 'tier_12_100', 'minutes': 720}, # 12 hours + config.STRIPE_TIER_25_200_ID: {'name': 'tier_25_200', 'minutes': 1500}, # 25 hours + config.STRIPE_TIER_50_400_ID: {'name': 'tier_50_400', 'minutes': 3000}, # 50 hours + config.STRIPE_TIER_125_800_ID: {'name': 'tier_125_800', 'minutes': 7500}, # 125 hours + config.STRIPE_TIER_200_1000_ID: {'name': 'tier_200_1000', 'minutes': 12000}, # 200 hours +} + +# Pydantic models for request/response validation +class CreateCheckoutSessionRequest(BaseModel): + price_id: str + success_url: str + cancel_url: str + tolt_referral: Optional[str] = None + +class CreatePortalSessionRequest(BaseModel): + return_url: str + +class SubscriptionStatus(BaseModel): + status: str # e.g., 'active', 'trialing', 'past_due', 'scheduled_downgrade', 'no_subscription' + plan_name: Optional[str] = None + price_id: Optional[str] = None # Added price ID + current_period_end: Optional[datetime] = None + cancel_at_period_end: bool = False + trial_end: Optional[datetime] = None + minutes_limit: Optional[int] = None + current_usage: Optional[float] = None + # Fields for scheduled changes + has_schedule: bool = False + scheduled_plan_name: Optional[str] = None + scheduled_price_id: Optional[str] = None # Added scheduled price ID + scheduled_change_date: Optional[datetime] = None + +# Helper functions +async def get_stripe_customer_id(client, user_id: str) -> Optional[str]: + """Get the Stripe customer ID for a user.""" + result = await client.schema('basejump').from_('billing_customers') \ + .select('id') \ + .eq('account_id', user_id) \ + .execute() + + if result.data and len(result.data) > 0: + return result.data[0]['id'] + return None + +async def create_stripe_customer(client, user_id: str, email: str) -> str: + """Create a new Stripe customer for a user.""" + # Create customer in Stripe + customer = stripe.Customer.create( + email=email, + metadata={"user_id": user_id} + ) + + # Store customer ID in Supabase + await client.schema('basejump').from_('billing_customers').insert({ + 'id': customer.id, + 'account_id': user_id, + 'email': email, + 'provider': 'stripe' + }).execute() + + return customer.id + +async def get_user_subscription(user_id: str) -> Optional[Dict]: + """Get the current subscription for a user from Stripe.""" + try: + # Get customer ID + db = DBConnection() + client = await db.client + customer_id = await get_stripe_customer_id(client, user_id) + + if not customer_id: + return None + + # Get all active subscriptions for the customer + subscriptions = stripe.Subscription.list( + customer=customer_id, + status='active' + ) + # print("Found subscriptions:", subscriptions) + + # Check if we have any subscriptions + if not subscriptions or not subscriptions.get('data'): + return None + + # Filter subscriptions to only include our product's subscriptions + our_subscriptions = [] + for sub in subscriptions['data']: + # Get the first subscription item + if sub.get('items') and sub['items'].get('data') and len(sub['items']['data']) > 0: + item = sub['items']['data'][0] + if item.get('price') and item['price'].get('id') in [ + config.STRIPE_FREE_TIER_ID, + config.STRIPE_TIER_2_20_ID, + config.STRIPE_TIER_6_50_ID, + config.STRIPE_TIER_12_100_ID, + config.STRIPE_TIER_25_200_ID, + config.STRIPE_TIER_50_400_ID, + config.STRIPE_TIER_125_800_ID, + config.STRIPE_TIER_200_1000_ID + ]: + our_subscriptions.append(sub) + + if not our_subscriptions: + return None + + # If there are multiple active subscriptions, we need to handle this + if len(our_subscriptions) > 1: + logger.warning(f"User {user_id} has multiple active subscriptions: {[sub['id'] for sub in our_subscriptions]}") + + # Get the most recent subscription + most_recent = max(our_subscriptions, key=lambda x: x['created']) + + # Cancel all other subscriptions + for sub in our_subscriptions: + if sub['id'] != most_recent['id']: + try: + stripe.Subscription.modify( + sub['id'], + cancel_at_period_end=True + ) + logger.info(f"Cancelled subscription {sub['id']} for user {user_id}") + except Exception as e: + logger.error(f"Error cancelling subscription {sub['id']}: {str(e)}") + + return most_recent + + return our_subscriptions[0] + + except Exception as e: + logger.error(f"Error getting subscription from Stripe: {str(e)}") + return None + +async def calculate_monthly_usage(client, user_id: str) -> float: + """Calculate total agent run minutes for the current month for a user.""" + # Get start of current month in UTC + now = datetime.now(timezone.utc) + start_of_month = datetime(now.year, now.month, 1, tzinfo=timezone.utc) + + # First get all threads for this user + threads_result = await client.table('threads') \ + .select('thread_id') \ + .eq('account_id', user_id) \ + .execute() + + if not threads_result.data: + return 0.0 + + thread_ids = [t['thread_id'] for t in threads_result.data] + + # Then get all agent runs for these threads in current month + runs_result = await client.table('agent_runs') \ + .select('started_at, completed_at') \ + .in_('thread_id', thread_ids) \ + .gte('started_at', start_of_month.isoformat()) \ + .execute() + + if not runs_result.data: + return 0.0 + + # Calculate total minutes + total_seconds = 0 + now_ts = now.timestamp() + + for run in runs_result.data: + start_time = datetime.fromisoformat(run['started_at'].replace('Z', '+00:00')).timestamp() + if run['completed_at']: + end_time = datetime.fromisoformat(run['completed_at'].replace('Z', '+00:00')).timestamp() + if start_time < end_time - 7200: + continue + else: + # if the start time is more than an hour ago, don't consider that time in total. else use the current time + if start_time < now_ts - 3600: + continue + else: + end_time = now_ts + + total_seconds += (end_time - start_time) + + return total_seconds / 60 # Convert to minutes + +async def get_allowed_models_for_user(client, user_id: str): + """ + Get the list of models allowed for a user based on their subscription tier. + + Returns: + List of model names allowed for the user's subscription tier. + """ + + subscription = await get_user_subscription(user_id) + tier_name = 'free' + + if subscription: + price_id = None + if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0: + price_id = subscription['items']['data'][0]['price']['id'] + else: + price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID) + + # Get tier info for this price_id + tier_info = SUBSCRIPTION_TIERS.get(price_id) + if tier_info: + tier_name = tier_info['name'] + + # Return allowed models for this tier + return MODEL_ACCESS_TIERS.get(tier_name, MODEL_ACCESS_TIERS['free']) # Default to free tier if unknown + + +async def can_use_model(client, user_id: str, model_name: str): + if config.ENV_MODE == EnvMode.LOCAL: + logger.info("Running in local development mode - billing checks are disabled") + return True, "Local development mode - billing disabled", { + "price_id": "local_dev", + "plan_name": "Local Development", + "minutes_limit": "no limit" + } + + allowed_models = await get_allowed_models_for_user(client, user_id) + resolved_model = MODEL_NAME_ALIASES.get(model_name, model_name) + if resolved_model in allowed_models: + return True, "Model access allowed", allowed_models + + return False, f"Your current subscription plan does not include access to {model_name}. Please upgrade your subscription or choose from your available models: {', '.join(allowed_models)}", allowed_models + +async def check_billing_status(client, user_id: str) -> Tuple[bool, str, Optional[Dict]]: + """ + Check if a user can run agents based on their subscription and usage. + + Returns: + Tuple[bool, str, Optional[Dict]]: (can_run, message, subscription_info) + """ + if config.ENV_MODE == EnvMode.LOCAL: + logger.info("Running in local development mode - billing checks are disabled") + return True, "Local development mode - billing disabled", { + "price_id": "local_dev", + "plan_name": "Local Development", + "minutes_limit": "no limit" + } + + # Get current subscription + subscription = await get_user_subscription(user_id) + # print("Current subscription:", subscription) + + # If no subscription, they can use free tier + if not subscription: + subscription = { + 'price_id': config.STRIPE_FREE_TIER_ID, # Free tier + 'plan_name': 'free' + } + + # Extract price ID from subscription items + price_id = None + if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0: + price_id = subscription['items']['data'][0]['price']['id'] + else: + price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID) + + # Get tier info - default to free tier if not found + tier_info = SUBSCRIPTION_TIERS.get(price_id) + if not tier_info: + logger.warning(f"Unknown subscription tier: {price_id}, defaulting to free tier") + tier_info = SUBSCRIPTION_TIERS[config.STRIPE_FREE_TIER_ID] + + # Calculate current month's usage + current_usage = await calculate_monthly_usage(client, user_id) + + # Check if within limits + if current_usage >= tier_info['minutes']: + return False, f"Monthly limit of {tier_info['minutes']} minutes reached. Please upgrade your plan or wait until next month.", subscription + + return True, "OK", subscription + +# API endpoints +@router.post("/create-checkout-session") +async def create_checkout_session( + request: CreateCheckoutSessionRequest, + current_user_id: str = Depends(get_current_user_id_from_jwt) +): + """Create a Stripe Checkout session or modify an existing subscription.""" + try: + # Get Supabase client + db = DBConnection() + client = await db.client + + # Get user email from auth.users + user_result = await client.auth.admin.get_user_by_id(current_user_id) + if not user_result: raise HTTPException(status_code=404, detail="User not found") + email = user_result.user.email + + # Get or create Stripe customer + customer_id = await get_stripe_customer_id(client, current_user_id) + if not customer_id: customer_id = await create_stripe_customer(client, current_user_id, email) + + # Get the target price and product ID + try: + price = stripe.Price.retrieve(request.price_id, expand=['product']) + product_id = price['product']['id'] + except stripe.error.InvalidRequestError: + raise HTTPException(status_code=400, detail=f"Invalid price ID: {request.price_id}") + + # Verify the price belongs to our product + if product_id != config.STRIPE_PRODUCT_ID: + raise HTTPException(status_code=400, detail="Price ID does not belong to the correct product.") + + # Check for existing subscription for our product + existing_subscription = await get_user_subscription(current_user_id) + # print("Existing subscription for product:", existing_subscription) + + if existing_subscription: + # --- Handle Subscription Change (Upgrade or Downgrade) --- + try: + subscription_id = existing_subscription['id'] + subscription_item = existing_subscription['items']['data'][0] + current_price_id = subscription_item['price']['id'] + + # Skip if already on this plan + if current_price_id == request.price_id: + return { + "subscription_id": subscription_id, + "status": "no_change", + "message": "Already subscribed to this plan.", + "details": { + "is_upgrade": None, + "effective_date": None, + "current_price": round(price['unit_amount'] / 100, 2) if price.get('unit_amount') else 0, + "new_price": round(price['unit_amount'] / 100, 2) if price.get('unit_amount') else 0, + } + } + + # Get current and new price details + current_price = stripe.Price.retrieve(current_price_id) + new_price = price # Already retrieved + is_upgrade = new_price['unit_amount'] > current_price['unit_amount'] + + if is_upgrade: + # --- Handle Upgrade --- Immediate modification + updated_subscription = stripe.Subscription.modify( + subscription_id, + items=[{ + 'id': subscription_item['id'], + 'price': request.price_id, + }], + proration_behavior='always_invoice', # Prorate and charge immediately + billing_cycle_anchor='now' # Reset billing cycle + ) + + # Update active status in database to true (customer has active subscription) + await client.schema('basejump').from_('billing_customers').update( + {'active': True} + ).eq('id', customer_id).execute() + logger.info(f"Updated customer {customer_id} active status to TRUE after subscription upgrade") + + latest_invoice = None + if updated_subscription.get('latest_invoice'): + latest_invoice = stripe.Invoice.retrieve(updated_subscription['latest_invoice']) + + return { + "subscription_id": updated_subscription['id'], + "status": "updated", + "message": "Subscription upgraded successfully", + "details": { + "is_upgrade": True, + "effective_date": "immediate", + "current_price": round(current_price['unit_amount'] / 100, 2) if current_price.get('unit_amount') else 0, + "new_price": round(new_price['unit_amount'] / 100, 2) if new_price.get('unit_amount') else 0, + "invoice": { + "id": latest_invoice['id'] if latest_invoice else None, + "status": latest_invoice['status'] if latest_invoice else None, + "amount_due": round(latest_invoice['amount_due'] / 100, 2) if latest_invoice else 0, + "amount_paid": round(latest_invoice['amount_paid'] / 100, 2) if latest_invoice else 0 + } if latest_invoice else None + } + } + else: + # --- Handle Downgrade --- Use Subscription Schedule + try: + current_period_end_ts = subscription_item['current_period_end'] + + # Retrieve the subscription again to get the schedule ID if it exists + # This ensures we have the latest state before creating/modifying schedule + sub_with_schedule = stripe.Subscription.retrieve(subscription_id) + schedule_id = sub_with_schedule.get('schedule') + + # Get the current phase configuration from the schedule or subscription + if schedule_id: + schedule = stripe.SubscriptionSchedule.retrieve(schedule_id) + # Find the current phase in the schedule + # This logic assumes simple schedules; might need refinement for complex ones + current_phase = None + for phase in reversed(schedule['phases']): + if phase['start_date'] <= datetime.now(timezone.utc).timestamp(): + current_phase = phase + break + if not current_phase: # Fallback if logic fails + current_phase = schedule['phases'][-1] + else: + # If no schedule, the current subscription state defines the current phase + current_phase = { + 'items': existing_subscription['items']['data'], # Use original items data + 'start_date': existing_subscription['current_period_start'], # Use sub start if no schedule + # Add other relevant fields if needed for create/modify + } + + # Prepare the current phase data for the update/create + # Ensure items is formatted correctly for the API + current_phase_items_for_api = [] + for item in current_phase.get('items', []): + price_data = item.get('price') + quantity = item.get('quantity') + price_id = None + + # Safely extract price ID whether it's an object or just the ID string + if isinstance(price_data, dict): + price_id = price_data.get('id') + elif isinstance(price_data, str): + price_id = price_data + + if price_id and quantity is not None: + current_phase_items_for_api.append({'price': price_id, 'quantity': quantity}) + else: + logger.warning(f"Skipping item in current phase due to missing price ID or quantity: {item}") + + if not current_phase_items_for_api: + raise ValueError("Could not determine valid items for the current phase.") + + current_phase_update_data = { + 'items': current_phase_items_for_api, + 'start_date': current_phase['start_date'], # Preserve original start date + 'end_date': current_period_end_ts, # End this phase at period end + 'proration_behavior': 'none' + # Include other necessary fields from current_phase if modifying? + # e.g., 'billing_cycle_anchor', 'collection_method'? Usually inherited. + } + + # Define the new (downgrade) phase + new_downgrade_phase_data = { + 'items': [{'price': request.price_id, 'quantity': 1}], + 'start_date': current_period_end_ts, # Start immediately after current phase ends + 'proration_behavior': 'none' + # iterations defaults to 1, meaning it runs for one billing cycle + # then schedule ends based on end_behavior + } + + # Update or Create Schedule + if schedule_id: + # Update existing schedule, replacing all future phases + # print(f"Updating existing schedule {schedule_id}") + logger.info(f"Updating existing schedule {schedule_id} for subscription {subscription_id}") + logger.debug(f"Current phase data: {current_phase_update_data}") + logger.debug(f"New phase data: {new_downgrade_phase_data}") + updated_schedule = stripe.SubscriptionSchedule.modify( + schedule_id, + phases=[current_phase_update_data, new_downgrade_phase_data], + end_behavior='release' + ) + logger.info(f"Successfully updated schedule {updated_schedule['id']}") + else: + # Create a new schedule using the defined phases + print(f"Creating new schedule for subscription {subscription_id}") + logger.info(f"Creating new schedule for subscription {subscription_id}") + # Deep debug logging - write subscription details to help diagnose issues + logger.debug(f"Subscription details: {subscription_id}, current_period_end_ts: {current_period_end_ts}") + logger.debug(f"Current price: {current_price_id}, New price: {request.price_id}") + + try: + updated_schedule = stripe.SubscriptionSchedule.create( + from_subscription=subscription_id, + phases=[ + { + 'start_date': current_phase['start_date'], + 'end_date': current_period_end_ts, + 'proration_behavior': 'none', + 'items': [ + { + 'price': current_price_id, + 'quantity': 1 + } + ] + }, + { + 'start_date': current_period_end_ts, + 'proration_behavior': 'none', + 'items': [ + { + 'price': request.price_id, + 'quantity': 1 + } + ] + } + ], + end_behavior='release' + ) + # Don't try to link the schedule - that's handled by from_subscription + logger.info(f"Created new schedule {updated_schedule['id']} from subscription {subscription_id}") + # print(f"Created new schedule {updated_schedule['id']} from subscription {subscription_id}") + + # Verify the schedule was created correctly + fetched_schedule = stripe.SubscriptionSchedule.retrieve(updated_schedule['id']) + logger.info(f"Schedule verification - Status: {fetched_schedule.get('status')}, Phase Count: {len(fetched_schedule.get('phases', []))}") + logger.debug(f"Schedule details: {fetched_schedule}") + except Exception as schedule_error: + logger.exception(f"Failed to create schedule: {str(schedule_error)}") + raise schedule_error # Re-raise to be caught by the outer try-except + + return { + "subscription_id": subscription_id, + "schedule_id": updated_schedule['id'], + "status": "scheduled", + "message": "Subscription downgrade scheduled", + "details": { + "is_upgrade": False, + "effective_date": "end_of_period", + "current_price": round(current_price['unit_amount'] / 100, 2) if current_price.get('unit_amount') else 0, + "new_price": round(new_price['unit_amount'] / 100, 2) if new_price.get('unit_amount') else 0, + "effective_at": datetime.fromtimestamp(current_period_end_ts, tz=timezone.utc).isoformat() + } + } + except Exception as e: + logger.exception(f"Error handling subscription schedule for sub {subscription_id}: {str(e)}") + raise HTTPException(status_code=500, detail=f"Error handling subscription schedule: {str(e)}") + except Exception as e: + logger.exception(f"Error updating subscription {existing_subscription.get('id') if existing_subscription else 'N/A'}: {str(e)}") + raise HTTPException(status_code=500, detail=f"Error updating subscription: {str(e)}") + else: + + session = stripe.checkout.Session.create( + customer=customer_id, + payment_method_types=['card'], + line_items=[{'price': request.price_id, 'quantity': 1}], + mode='subscription', + success_url=request.success_url, + cancel_url=request.cancel_url, + metadata={ + 'user_id': current_user_id, + 'product_id': product_id, + 'tolt_referral': request.tolt_referral + }, + allow_promotion_codes=True + ) + + # Update customer status to potentially active (will be confirmed by webhook) + # This ensures customer is marked as active once payment is completed + await client.schema('basejump').from_('billing_customers').update( + {'active': True} + ).eq('id', customer_id).execute() + logger.info(f"Updated customer {customer_id} active status to TRUE after creating checkout session") + + return {"session_id": session['id'], "url": session['url'], "status": "new"} + + except Exception as e: + logger.exception(f"Error creating checkout session: {str(e)}") + # Check if it's a Stripe error with more details + if hasattr(e, 'json_body') and e.json_body and 'error' in e.json_body: + error_detail = e.json_body['error'].get('message', str(e)) + else: + error_detail = str(e) + raise HTTPException(status_code=500, detail=f"Error creating checkout session: {error_detail}") + +@router.post("/create-portal-session") +async def create_portal_session( + request: CreatePortalSessionRequest, + current_user_id: str = Depends(get_current_user_id_from_jwt) +): + """Create a Stripe Customer Portal session for subscription management.""" + try: + # Get Supabase client + db = DBConnection() + client = await db.client + + # Get customer ID + customer_id = await get_stripe_customer_id(client, current_user_id) + if not customer_id: + raise HTTPException(status_code=404, detail="No billing customer found") + + # Ensure the portal configuration has subscription_update enabled + try: + # First, check if we have a configuration that already enables subscription update + configurations = stripe.billing_portal.Configuration.list(limit=100) + active_config = None + + # Look for a configuration with subscription_update enabled + for config in configurations.get('data', []): + features = config.get('features', {}) + subscription_update = features.get('subscription_update', {}) + if subscription_update.get('enabled', False): + active_config = config + logger.info(f"Found existing portal configuration with subscription_update enabled: {config['id']}") + break + + # If no config with subscription_update found, create one or update the active one + if not active_config: + # Find the active configuration or create a new one + if configurations.get('data', []): + default_config = configurations['data'][0] + logger.info(f"Updating default portal configuration: {default_config['id']} to enable subscription_update") + + active_config = stripe.billing_portal.Configuration.update( + default_config['id'], + features={ + 'subscription_update': { + 'enabled': True, + 'proration_behavior': 'create_prorations', + 'default_allowed_updates': ['price'] + }, + # Preserve other features that may already be enabled + 'customer_update': default_config.get('features', {}).get('customer_update', {'enabled': True, 'allowed_updates': ['email', 'address']}), + 'invoice_history': {'enabled': True}, + 'payment_method_update': {'enabled': True} + } + ) + else: + # Create a new configuration with subscription_update enabled + logger.info("Creating new portal configuration with subscription_update enabled") + active_config = stripe.billing_portal.Configuration.create( + business_profile={ + 'headline': 'Subscription Management', + 'privacy_policy_url': config.FRONTEND_URL + '/privacy', + 'terms_of_service_url': config.FRONTEND_URL + '/terms' + }, + features={ + 'subscription_update': { + 'enabled': True, + 'proration_behavior': 'create_prorations', + 'default_allowed_updates': ['price'] + }, + 'customer_update': { + 'enabled': True, + 'allowed_updates': ['email', 'address'] + }, + 'invoice_history': {'enabled': True}, + 'payment_method_update': {'enabled': True} + } + ) + + # Log the active configuration for debugging + logger.info(f"Using portal configuration: {active_config['id']} with subscription_update: {active_config.get('features', {}).get('subscription_update', {}).get('enabled', False)}") + + except Exception as config_error: + logger.warning(f"Error configuring portal: {config_error}. Continuing with default configuration.") + + # Create portal session using the proper configuration if available + portal_params = { + "customer": customer_id, + "return_url": request.return_url + } + + # Add configuration_id if we found or created one with subscription_update enabled + if active_config: + portal_params["configuration"] = active_config['id'] + + # Create the session + session = stripe.billing_portal.Session.create(**portal_params) + + return {"url": session.url} + + except Exception as e: + logger.error(f"Error creating portal session: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@router.get("/subscription") +async def get_subscription( + current_user_id: str = Depends(get_current_user_id_from_jwt) +): + """Get the current subscription status for the current user, including scheduled changes.""" + try: + # Get subscription from Stripe (this helper already handles filtering/cleanup) + subscription = await get_user_subscription(current_user_id) + # print("Subscription data for status:", subscription) + + if not subscription: + # Default to free tier status if no active subscription for our product + free_tier_id = config.STRIPE_FREE_TIER_ID + free_tier_info = SUBSCRIPTION_TIERS.get(free_tier_id) + return SubscriptionStatus( + status="no_subscription", + plan_name=free_tier_info.get('name', 'free') if free_tier_info else 'free', + price_id=free_tier_id, + minutes_limit=free_tier_info.get('minutes') if free_tier_info else 0 + ) + + # Extract current plan details + current_item = subscription['items']['data'][0] + current_price_id = current_item['price']['id'] + current_tier_info = SUBSCRIPTION_TIERS.get(current_price_id) + if not current_tier_info: + # Fallback if somehow subscribed to an unknown price within our product + logger.warning(f"User {current_user_id} subscribed to unknown price {current_price_id}. Defaulting info.") + current_tier_info = {'name': 'unknown', 'minutes': 0} + + # Calculate current usage + db = DBConnection() + client = await db.client + current_usage = await calculate_monthly_usage(client, current_user_id) + + status_response = SubscriptionStatus( + status=subscription['status'], # 'active', 'trialing', etc. + plan_name=subscription['plan'].get('nickname') or current_tier_info['name'], + price_id=current_price_id, + current_period_end=datetime.fromtimestamp(current_item['current_period_end'], tz=timezone.utc), + cancel_at_period_end=subscription['cancel_at_period_end'], + trial_end=datetime.fromtimestamp(subscription['trial_end'], tz=timezone.utc) if subscription.get('trial_end') else None, + minutes_limit=current_tier_info['minutes'], + current_usage=round(current_usage, 2), + has_schedule=False # Default + ) + + # Check for an attached schedule (indicates pending downgrade) + schedule_id = subscription.get('schedule') + if schedule_id: + try: + schedule = stripe.SubscriptionSchedule.retrieve(schedule_id) + # Find the *next* phase after the current one + next_phase = None + current_phase_end = current_item['current_period_end'] + + for phase in schedule.get('phases', []): + # Check if this phase starts exactly when the current one ends + if phase.get('start_date') == current_phase_end: + next_phase = phase + break # Found the immediate next phase + + if next_phase: + scheduled_item = next_phase['items'][0] # Assuming single item + scheduled_price_id = scheduled_item['price'] # Price ID might be string here + scheduled_tier_info = SUBSCRIPTION_TIERS.get(scheduled_price_id) + + status_response.has_schedule = True + status_response.status = 'scheduled_downgrade' # Override status + status_response.scheduled_plan_name = scheduled_tier_info.get('name', 'unknown') if scheduled_tier_info else 'unknown' + status_response.scheduled_price_id = scheduled_price_id + status_response.scheduled_change_date = datetime.fromtimestamp(next_phase['start_date'], tz=timezone.utc) + + except Exception as schedule_error: + logger.error(f"Error retrieving or parsing schedule {schedule_id} for sub {subscription['id']}: {schedule_error}") + # Proceed without schedule info if retrieval fails + + return status_response + + except Exception as e: + logger.exception(f"Error getting subscription status for user {current_user_id}: {str(e)}") # Use logger.exception + raise HTTPException(status_code=500, detail="Error retrieving subscription status.") + +@router.get("/check-status") +async def check_status( + current_user_id: str = Depends(get_current_user_id_from_jwt) +): + """Check if the user can run agents based on their subscription and usage.""" + try: + # Get Supabase client + db = DBConnection() + client = await db.client + + can_run, message, subscription = await check_billing_status(client, current_user_id) + + return { + "can_run": can_run, + "message": message, + "subscription": subscription + } + + except Exception as e: + logger.error(f"Error checking billing status: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@router.post("/webhook") +async def stripe_webhook(request: Request): + """Handle Stripe webhook events.""" + try: + # Get the webhook secret from config + webhook_secret = config.STRIPE_WEBHOOK_SECRET + + # Get the webhook payload + payload = await request.body() + sig_header = request.headers.get('stripe-signature') + + # Verify webhook signature + try: + event = stripe.Webhook.construct_event( + payload, sig_header, webhook_secret + ) + except ValueError as e: + raise HTTPException(status_code=400, detail="Invalid payload") + except stripe.error.SignatureVerificationError as e: + raise HTTPException(status_code=400, detail="Invalid signature") + + # Handle the event + if event.type in ['customer.subscription.created', 'customer.subscription.updated', 'customer.subscription.deleted']: + # Extract the subscription and customer information + subscription = event.data.object + customer_id = subscription.get('customer') + + if not customer_id: + logger.warning(f"No customer ID found in subscription event: {event.type}") + return {"status": "error", "message": "No customer ID found"} + + # Get database connection + db = DBConnection() + client = await db.client + + if event.type == 'customer.subscription.created' or event.type == 'customer.subscription.updated': + # Check if subscription is active + if subscription.get('status') in ['active', 'trialing']: + # Update customer's active status to true + await client.schema('basejump').from_('billing_customers').update( + {'active': True} + ).eq('id', customer_id).execute() + logger.info(f"Webhook: Updated customer {customer_id} active status to TRUE based on {event.type}") + else: + # Subscription is not active (e.g., past_due, canceled, etc.) + # Check if customer has any other active subscriptions before updating status + has_active = len(stripe.Subscription.list( + customer=customer_id, + status='active', + limit=1 + ).get('data', [])) > 0 + + if not has_active: + await client.schema('basejump').from_('billing_customers').update( + {'active': False} + ).eq('id', customer_id).execute() + logger.info(f"Webhook: Updated customer {customer_id} active status to FALSE based on {event.type}") + + elif event.type == 'customer.subscription.deleted': + # Check if customer has any other active subscriptions + has_active = len(stripe.Subscription.list( + customer=customer_id, + status='active', + limit=1 + ).get('data', [])) > 0 + + if not has_active: + # If no active subscriptions left, set active to false + await client.schema('basejump').from_('billing_customers').update( + {'active': False} + ).eq('id', customer_id).execute() + logger.info(f"Webhook: Updated customer {customer_id} active status to FALSE after subscription deletion") + + logger.info(f"Processed {event.type} event for customer {customer_id}") + + return {"status": "success"} + + except Exception as e: + logger.error(f"Error processing webhook: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@router.get("/available-models") +async def get_available_models( + current_user_id: str = Depends(get_current_user_id_from_jwt) +): + """Get the list of models available to the user based on their subscription tier.""" + try: + # Get Supabase client + db = DBConnection() + client = await db.client + + # Check if we're in local development mode + if config.ENV_MODE == EnvMode.LOCAL: + logger.info("Running in local development mode - billing checks are disabled") + + # In local mode, return all models from MODEL_NAME_ALIASES + model_info = [] + for short_name, full_name in MODEL_NAME_ALIASES.items(): + # Skip entries where the key is a full name to avoid duplicates + # if short_name == full_name or '/' in short_name: + # continue + + model_info.append({ + "id": full_name, + "display_name": short_name, + "short_name": short_name, + "requires_subscription": False # Always false in local dev mode + }) + + return { + "models": model_info, + "subscription_tier": "Local Development", + "total_models": len(model_info) + } + + # For non-local mode, get list of allowed models for this user + allowed_models = await get_allowed_models_for_user(client, current_user_id) + free_tier_models = MODEL_ACCESS_TIERS.get('free', []) + + # Get subscription info for context + subscription = await get_user_subscription(current_user_id) + + # Determine tier name from subscription + tier_name = 'free' + if subscription: + price_id = None + if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0: + price_id = subscription['items']['data'][0]['price']['id'] + else: + price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID) + + # Get tier info for this price_id + tier_info = SUBSCRIPTION_TIERS.get(price_id) + if tier_info: + tier_name = tier_info['name'] + + # Get all unique full model names from MODEL_NAME_ALIASES + all_models = set() + model_aliases = {} + + for short_name, full_name in MODEL_NAME_ALIASES.items(): + # Add all unique full model names + all_models.add(full_name) + + # Only include short names that don't match their full names for aliases + if short_name != full_name and not short_name.startswith("openai/") and not short_name.startswith("anthropic/") and not short_name.startswith("openrouter/") and not short_name.startswith("xai/"): + if full_name not in model_aliases: + model_aliases[full_name] = short_name + + # Create model info with display names for ALL models + model_info = [] + for model in all_models: + display_name = model_aliases.get(model, model.split('/')[-1] if '/' in model else model) + + # Check if model requires subscription (not in free tier) + requires_sub = model not in free_tier_models + + # Check if model is available with current subscription + is_available = model in allowed_models + + model_info.append({ + "id": model, + "display_name": display_name, + "short_name": model_aliases.get(model), + "requires_subscription": requires_sub, + "is_available": is_available + }) + + return { + "models": model_info, + "subscription_tier": tier_name, + "total_models": len(model_info) + } + + except Exception as e: + logger.error(f"Error getting available models: {str(e)}") + raise HTTPException(status_code=500, detail=f"Error getting available models: {str(e)}") \ No newline at end of file diff --git a/app/services/docker/redis.conf b/app/services/docker/redis.conf new file mode 100644 index 0000000..b8b4180 --- /dev/null +++ b/app/services/docker/redis.conf @@ -0,0 +1 @@ +timeout 120 diff --git a/app/services/email.py b/app/services/email.py new file mode 100644 index 0000000..76dcb47 --- /dev/null +++ b/app/services/email.py @@ -0,0 +1,192 @@ +import os +import logging +from typing import Optional +import mailtrap as mt +from utils.config import config + +logger = logging.getLogger(__name__) + +class EmailService: + def __init__(self): + self.api_token = os.getenv('MAILTRAP_API_TOKEN') + self.sender_email = os.getenv('MAILTRAP_SENDER_EMAIL', 'dom@kortix.ai') + self.sender_name = os.getenv('MAILTRAP_SENDER_NAME', 'Suna Team') + + if not self.api_token: + logger.warning("MAILTRAP_API_TOKEN not found in environment variables") + self.client = None + else: + self.client = mt.MailtrapClient(token=self.api_token) + + def send_welcome_email(self, user_email: str, user_name: Optional[str] = None) -> bool: + if not self.client: + logger.error("Cannot send email: MAILTRAP_API_TOKEN not configured") + return False + + if not user_name: + user_name = user_email.split('@')[0].title() + + subject = "🎉 Welcome to Suna — Let's Get Started " + html_content = self._get_welcome_email_template(user_name) + text_content = self._get_welcome_email_text(user_name) + + return self._send_email( + to_email=user_email, + to_name=user_name, + subject=subject, + html_content=html_content, + text_content=text_content + ) + + def _send_email( + self, + to_email: str, + to_name: str, + subject: str, + html_content: str, + text_content: str + ) -> bool: + try: + mail = mt.Mail( + sender=mt.Address(email=self.sender_email, name=self.sender_name), + to=[mt.Address(email=to_email, name=to_name)], + subject=subject, + text=text_content, + html=html_content, + category="welcome" + ) + + response = self.client.send(mail) + + logger.info(f"Welcome email sent to {to_email}. Response: {response}") + return True + + except Exception as e: + logger.error(f"Error sending email to {to_email}: {str(e)}") + return False + + def _get_welcome_email_template(self, user_name: str) -> str: + return f""" + + + + + Welcome to Kortix Suna + + + +
+
+ +
+

Welcome to Kortix Suna!

+ +

Hi {user_name},

+ +

Welcome to Kortix Suna — we're excited to have you on board!

+ +

To get started, we'd like to get to know you better: fill out this short form!

+ +

To celebrate your arrival, here's a 15% discount to try out the best version of Suna (1 month):

+ +

🎁 Use code WELCOME15 at checkout.

+ +

Let us know if you need help getting started or have questions — we're always here, and join our Discord community.

+ +

For your business: if you want to automate manual and ordinary tasks for your company, book a call with us here

+ +

Thanks again, and welcome to the Suna community 🌞

+ +

— The Suna Team

+ + Go to the platform +
+ +""" + + def _get_welcome_email_text(self, user_name: str) -> str: + return f"""Hi {user_name}, + +Welcome to Suna — we're excited to have you on board! + +To get started, we'd like to get to know you better: fill out this short form! +https://docs.google.com/forms/d/e/1FAIpQLSef1EHuqmIh_iQz-kwhjnzSC3Ml-V_5wIySDpMoMU9W_j24JQ/viewform + +To celebrate your arrival, here's a 15% discount to try out the best version of Suna (1 month): +🎁 Use code WELCOME15 at checkout. + +Let us know if you need help getting started or have questions — we're always here, and join our Discord community: https://discord.com/invite/FjD644cfcs + +For your business: if you want to automate manual and ordinary tasks for your company, book a call with us here: https://cal.com/team/kortix/enterprise-demo + +Thanks again, and welcome to the Suna community 🌞 + +— The Suna Team + +Go to the platform: https://www.suna.so/ + +--- +© 2024 Suna. All rights reserved. +You received this email because you signed up for a Suna account.""" + +email_service = EmailService() diff --git a/app/services/email_api.py b/app/services/email_api.py new file mode 100644 index 0000000..9834c7b --- /dev/null +++ b/app/services/email_api.py @@ -0,0 +1,70 @@ +from fastapi import APIRouter, HTTPException, Depends +from pydantic import BaseModel, EmailStr +from typing import Optional +import asyncio +from services.email import email_service +from utils.logger import logger + +router = APIRouter() + +class SendWelcomeEmailRequest(BaseModel): + email: EmailStr + name: Optional[str] = None + +class EmailResponse(BaseModel): + success: bool + message: str + +@router.post("/send-welcome-email", response_model=EmailResponse) +async def send_welcome_email(request: SendWelcomeEmailRequest): + try: + logger.info(f"Sending welcome email to {request.email}") + success = email_service.send_welcome_email( + user_email=request.email, + user_name=request.name + ) + + if success: + return EmailResponse( + success=True, + message="Welcome email sent successfully" + ) + else: + return EmailResponse( + success=False, + message="Failed to send welcome email" + ) + + except Exception as e: + logger.error(f"Error sending welcome email to {request.email}: {str(e)}") + raise HTTPException( + status_code=500, + detail="Internal server error while sending email" + ) + +@router.post("/send-welcome-email-background", response_model=EmailResponse) +async def send_welcome_email_background(request: SendWelcomeEmailRequest): + try: + logger.info(f"Queuing welcome email for {request.email}") + + def send_email(): + return email_service.send_welcome_email( + user_email=request.email, + user_name=request.name + ) + + import concurrent.futures + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(send_email) + + return EmailResponse( + success=True, + message="Welcome email queued for sending" + ) + + except Exception as e: + logger.error(f"Error queuing welcome email for {request.email}: {str(e)}") + raise HTTPException( + status_code=500, + detail="Internal server error while queuing email" + ) diff --git a/app/services/langfuse.py b/app/services/langfuse.py new file mode 100644 index 0000000..cf624bf --- /dev/null +++ b/app/services/langfuse.py @@ -0,0 +1,12 @@ +import os +from langfuse import Langfuse + +public_key = os.getenv("LANGFUSE_PUBLIC_KEY") +secret_key = os.getenv("LANGFUSE_SECRET_KEY") +host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + +enabled = False +if public_key and secret_key: + enabled = True + +langfuse = Langfuse(enabled=enabled) diff --git a/app/services/llm.py b/app/services/llm.py new file mode 100644 index 0000000..c749170 --- /dev/null +++ b/app/services/llm.py @@ -0,0 +1,411 @@ +""" +LLM API interface for making calls to various language models. + +This module provides a unified interface for making API calls to different LLM providers +(OpenAI, Anthropic, Groq, etc.) using LiteLLM. It includes support for: +- Streaming responses +- Tool calls and function calling +- Retry logic with exponential backoff +- Model-specific configurations +- Comprehensive error handling and logging +""" + +from typing import Union, Dict, Any, Optional, AsyncGenerator, List +import os +import json +import asyncio +from openai import OpenAIError +import litellm +from utils.logger import logger +from utils.config import config + +# litellm.set_verbose=True +litellm.modify_params=True + +# Constants +MAX_RETRIES = 2 +RATE_LIMIT_DELAY = 30 +RETRY_DELAY = 0.1 + +class LLMError(Exception): + """Base exception for LLM-related errors.""" + pass + +class LLMRetryError(LLMError): + """Exception raised when retries are exhausted.""" + pass + +def setup_api_keys() -> None: + """Set up API keys from environment variables.""" + providers = ['OPENAI', 'ANTHROPIC', 'GROQ', 'OPENROUTER'] + for provider in providers: + key = getattr(config, f'{provider}_API_KEY') + if key: + logger.debug(f"API key set for provider: {provider}") + else: + logger.warning(f"No API key found for provider: {provider}") + + # Set up OpenRouter API base if not already set + if config.OPENROUTER_API_KEY and config.OPENROUTER_API_BASE: + os.environ['OPENROUTER_API_BASE'] = config.OPENROUTER_API_BASE + logger.debug(f"Set OPENROUTER_API_BASE to {config.OPENROUTER_API_BASE}") + + # Set up AWS Bedrock credentials + aws_access_key = config.AWS_ACCESS_KEY_ID + aws_secret_key = config.AWS_SECRET_ACCESS_KEY + aws_region = config.AWS_REGION_NAME + + if aws_access_key and aws_secret_key and aws_region: + logger.debug(f"AWS credentials set for Bedrock in region: {aws_region}") + # Configure LiteLLM to use AWS credentials + os.environ['AWS_ACCESS_KEY_ID'] = aws_access_key + os.environ['AWS_SECRET_ACCESS_KEY'] = aws_secret_key + os.environ['AWS_REGION_NAME'] = aws_region + else: + logger.warning(f"Missing AWS credentials for Bedrock integration - access_key: {bool(aws_access_key)}, secret_key: {bool(aws_secret_key)}, region: {aws_region}") + +async def handle_error(error: Exception, attempt: int, max_attempts: int) -> None: + """Handle API errors with appropriate delays and logging.""" + delay = RATE_LIMIT_DELAY if isinstance(error, litellm.exceptions.RateLimitError) else RETRY_DELAY + logger.warning(f"Error on attempt {attempt + 1}/{max_attempts}: {str(error)}") + logger.debug(f"Waiting {delay} seconds before retry...") + await asyncio.sleep(delay) + +def prepare_params( + messages: List[Dict[str, Any]], + model_name: str, + temperature: float = 0, + max_tokens: Optional[int] = None, + response_format: Optional[Any] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: str = "auto", + api_key: Optional[str] = None, + api_base: Optional[str] = None, + stream: bool = False, + top_p: Optional[float] = None, + model_id: Optional[str] = None, + enable_thinking: Optional[bool] = False, + reasoning_effort: Optional[str] = 'low' +) -> Dict[str, Any]: + """Prepare parameters for the API call.""" + params = { + "model": model_name, + "messages": messages, + "temperature": temperature, + "response_format": response_format, + "top_p": top_p, + "stream": stream, + } + + if api_key: + params["api_key"] = api_key + if api_base: + params["api_base"] = api_base + if model_id: + params["model_id"] = model_id + + # Handle token limits + if max_tokens is not None: + # For Claude 3.7 in Bedrock, do not set max_tokens or max_tokens_to_sample + # as it causes errors with inference profiles + if model_name.startswith("bedrock/") and "claude-3-7" in model_name: + logger.debug(f"Skipping max_tokens for Claude 3.7 model: {model_name}") + # Do not add any max_tokens parameter for Claude 3.7 + else: + param_name = "max_completion_tokens" if 'o1' in model_name else "max_tokens" + params[param_name] = max_tokens + + # Add tools if provided + if tools: + params.update({ + "tools": tools, + "tool_choice": tool_choice + }) + logger.debug(f"Added {len(tools)} tools to API parameters") + + # # Add Claude-specific headers + if "claude" in model_name.lower() or "anthropic" in model_name.lower(): + params["extra_headers"] = { + # "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15" + "anthropic-beta": "output-128k-2025-02-19" + } + params["fallbacks"] = [{ + "model": "openrouter/anthropic/claude-sonnet-4", + "messages": messages, + }] + logger.debug("Added Claude-specific headers") + + # Add OpenRouter-specific parameters + if model_name.startswith("openrouter/"): + logger.debug(f"Preparing OpenRouter parameters for model: {model_name}") + + # Add optional site URL and app name from config + site_url = config.OR_SITE_URL + app_name = config.OR_APP_NAME + if site_url or app_name: + extra_headers = params.get("extra_headers", {}) + if site_url: + extra_headers["HTTP-Referer"] = site_url + if app_name: + extra_headers["X-Title"] = app_name + params["extra_headers"] = extra_headers + logger.debug(f"Added OpenRouter site URL and app name to headers") + + # Add Bedrock-specific parameters + if model_name.startswith("bedrock/"): + logger.debug(f"Preparing AWS Bedrock parameters for model: {model_name}") + + if not model_id and "anthropic.claude-3-7-sonnet" in model_name: + params["model_id"] = "arn:aws:bedrock:us-west-2:935064898258:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + logger.debug(f"Auto-set model_id for Claude 3.7 Sonnet: {params['model_id']}") + + # Apply Anthropic prompt caching (minimal implementation) + # Check model name *after* potential modifications (like adding bedrock/ prefix) + effective_model_name = params.get("model", model_name) # Use model from params if set, else original + if "claude" in effective_model_name.lower() or "anthropic" in effective_model_name.lower(): + messages = params["messages"] # Direct reference, modification affects params + + # Ensure messages is a list + if not isinstance(messages, list): + return params # Return early if messages format is unexpected + + # 1. Process the first message if it's a system prompt with string content + if messages and messages[0].get("role") == "system": + content = messages[0].get("content") + if isinstance(content, str): + # Wrap the string content in the required list structure + messages[0]["content"] = [ + {"type": "text", "text": content, "cache_control": {"type": "ephemeral"}} + ] + elif isinstance(content, list): + # If content is already a list, check if the first text block needs cache_control + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + if "cache_control" not in item: + item["cache_control"] = {"type": "ephemeral"} + break # Apply to the first text block only for system prompt + + # 2. Find and process relevant user and assistant messages (limit to 4 max) + last_user_idx = -1 + second_last_user_idx = -1 + last_assistant_idx = -1 + + for i in range(len(messages) - 1, -1, -1): + role = messages[i].get("role") + if role == "user": + if last_user_idx == -1: + last_user_idx = i + elif second_last_user_idx == -1: + second_last_user_idx = i + elif role == "assistant": + if last_assistant_idx == -1: + last_assistant_idx = i + + # Stop searching if we've found all needed messages (system, last user, second last user, last assistant) + found_count = sum(idx != -1 for idx in [last_user_idx, second_last_user_idx, last_assistant_idx]) + if found_count >= 3: + break + + # Helper function to apply cache control + def apply_cache_control(message_idx: int, message_role: str): + if message_idx == -1: + return + + message = messages[message_idx] + content = message.get("content") + + if isinstance(content, str): + message["content"] = [ + {"type": "text", "text": content, "cache_control": {"type": "ephemeral"}} + ] + elif isinstance(content, list): + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + if "cache_control" not in item: + item["cache_control"] = {"type": "ephemeral"} + + # Apply cache control to the identified messages (max 4: system, last user, second last user, last assistant) + # System message is always at index 0 if present + apply_cache_control(0, "system") + apply_cache_control(last_user_idx, "last user") + apply_cache_control(second_last_user_idx, "second last user") + apply_cache_control(last_assistant_idx, "last assistant") + + # Add reasoning_effort for Anthropic models if enabled + use_thinking = enable_thinking if enable_thinking is not None else False + is_anthropic = "anthropic" in effective_model_name.lower() or "claude" in effective_model_name.lower() + + if is_anthropic and use_thinking: + effort_level = reasoning_effort if reasoning_effort else 'low' + params["reasoning_effort"] = effort_level + params["temperature"] = 1.0 # Required by Anthropic when reasoning_effort is used + logger.info(f"Anthropic thinking enabled with reasoning_effort='{effort_level}'") + + return params + +async def make_llm_api_call( + messages: List[Dict[str, Any]], + model_name: str, + response_format: Optional[Any] = None, + temperature: float = 0, + max_tokens: Optional[int] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: str = "auto", + api_key: Optional[str] = None, + api_base: Optional[str] = None, + stream: bool = False, + top_p: Optional[float] = None, + model_id: Optional[str] = None, + enable_thinking: Optional[bool] = False, + reasoning_effort: Optional[str] = 'low' +) -> Union[Dict[str, Any], AsyncGenerator]: + """ + Make an API call to a language model using LiteLLM. + + Args: + messages: List of message dictionaries for the conversation + model_name: Name of the model to use (e.g., "gpt-4", "claude-3", "openrouter/openai/gpt-4", "bedrock/anthropic.claude-3-sonnet-20240229-v1:0") + response_format: Desired format for the response + temperature: Sampling temperature (0-1) + max_tokens: Maximum tokens in the response + tools: List of tool definitions for function calling + tool_choice: How to select tools ("auto" or "none") + api_key: Override default API key + api_base: Override default API base URL + stream: Whether to stream the response + top_p: Top-p sampling parameter + model_id: Optional ARN for Bedrock inference profiles + enable_thinking: Whether to enable thinking + reasoning_effort: Level of reasoning effort + + Returns: + Union[Dict[str, Any], AsyncGenerator]: API response or stream + + Raises: + LLMRetryError: If API call fails after retries + LLMError: For other API-related errors + """ + # debug .json messages + logger.info(f"Making LLM API call to model: {model_name} (Thinking: {enable_thinking}, Effort: {reasoning_effort})") + logger.info(f"📡 API Call: Using model {model_name}") + params = prepare_params( + messages=messages, + model_name=model_name, + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + tools=tools, + tool_choice=tool_choice, + api_key=api_key, + api_base=api_base, + stream=stream, + top_p=top_p, + model_id=model_id, + enable_thinking=enable_thinking, + reasoning_effort=reasoning_effort + ) + last_error = None + for attempt in range(MAX_RETRIES): + try: + logger.debug(f"Attempt {attempt + 1}/{MAX_RETRIES}") + # logger.debug(f"API request parameters: {json.dumps(params, indent=2)}") + + response = await litellm.acompletion(**params) + logger.debug(f"Successfully received API response from {model_name}") + logger.debug(f"Response: {response}") + return response + + except (litellm.exceptions.RateLimitError, OpenAIError, json.JSONDecodeError) as e: + last_error = e + await handle_error(e, attempt, MAX_RETRIES) + + except Exception as e: + logger.error(f"Unexpected error during API call: {str(e)}", exc_info=True) + raise LLMError(f"API call failed: {str(e)}") + + error_msg = f"Failed to make API call after {MAX_RETRIES} attempts" + if last_error: + error_msg += f". Last error: {str(last_error)}" + logger.error(error_msg, exc_info=True) + raise LLMRetryError(error_msg) + +# Initialize API keys on module import +setup_api_keys() + +# Test code for OpenRouter integration +async def test_openrouter(): + """Test the OpenRouter integration with a simple query.""" + test_messages = [ + {"role": "user", "content": "Hello, can you give me a quick test response?"} + ] + + try: + # Test with standard OpenRouter model + print("\n--- Testing standard OpenRouter model ---") + response = await make_llm_api_call( + model_name="openrouter/openai/gpt-4o-mini", + messages=test_messages, + temperature=0.7, + max_tokens=100 + ) + print(f"Response: {response.choices[0].message.content}") + + # Test with deepseek model + print("\n--- Testing deepseek model ---") + response = await make_llm_api_call( + model_name="openrouter/deepseek/deepseek-r1-distill-llama-70b", + messages=test_messages, + temperature=0.7, + max_tokens=100 + ) + print(f"Response: {response.choices[0].message.content}") + print(f"Model used: {response.model}") + + # Test with Mistral model + print("\n--- Testing Mistral model ---") + response = await make_llm_api_call( + model_name="openrouter/mistralai/mixtral-8x7b-instruct", + messages=test_messages, + temperature=0.7, + max_tokens=100 + ) + print(f"Response: {response.choices[0].message.content}") + print(f"Model used: {response.model}") + + return True + except Exception as e: + print(f"Error testing OpenRouter: {str(e)}") + return False + +async def test_bedrock(): + """Test the AWS Bedrock integration with a simple query.""" + test_messages = [ + {"role": "user", "content": "Hello, can you give me a quick test response?"} + ] + + try: + response = await make_llm_api_call( + model_name="bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0", + model_id="arn:aws:bedrock:us-west-2:935064898258:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + messages=test_messages, + temperature=0.7, + # Claude 3.7 has issues with max_tokens, so omit it + # max_tokens=100 + ) + print(f"Response: {response.choices[0].message.content}") + print(f"Model used: {response.model}") + + return True + except Exception as e: + print(f"Error testing Bedrock: {str(e)}") + return False + +if __name__ == "__main__": + import asyncio + + test_success = asyncio.run(test_bedrock()) + + if test_success: + print("\n✅ integration test completed successfully!") + else: + print("\n❌ Bedrock integration test failed!") diff --git a/app/services/mcp_custom.py b/app/services/mcp_custom.py new file mode 100644 index 0000000..1a9c245 --- /dev/null +++ b/app/services/mcp_custom.py @@ -0,0 +1,129 @@ +import os +import sys +import json +import asyncio +import subprocess +from typing import Dict, Any +from concurrent.futures import ThreadPoolExecutor +from fastapi import HTTPException # type: ignore +from utils.logger import logger +from mcp import ClientSession +from mcp.client.sse import sse_client # type: ignore +from mcp.client.streamable_http import streamablehttp_client # type: ignore + +async def connect_streamable_http_server(url): + async with streamablehttp_client(url) as ( + read_stream, + write_stream, + _, + ): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + tool_result = await session.list_tools() + print(f"Connected via HTTP ({len(tool_result.tools)} tools)") + + tools_info = [] + for tool in tool_result.tools: + tool_info = { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.inputSchema + } + tools_info.append(tool_info) + + return tools_info + +async def discover_custom_tools(request_type: str, config: Dict[str, Any]): + logger.info(f"Received custom MCP discovery request: type={request_type}") + logger.debug(f"Request config: {config}") + + tools = [] + server_name = None + + if request_type == 'http': + if 'url' not in config: + raise HTTPException(status_code=400, detail="HTTP configuration must include 'url' field") + url = config['url'] + + try: + async with asyncio.timeout(15): + tools_info = await connect_streamable_http_server(url) + for tool_info in tools_info: + tools.append({ + "name": tool_info["name"], + "description": tool_info["description"], + "inputSchema": tool_info["inputSchema"] + }) + except asyncio.TimeoutError: + raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond") + except Exception as e: + logger.error(f"Error connecting to HTTP MCP server: {e}") + raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}") + + elif request_type == 'sse': + if 'url' not in config: + raise HTTPException(status_code=400, detail="SSE configuration must include 'url' field") + + url = config['url'] + headers = config.get('headers', {}) + + try: + async with asyncio.timeout(15): + try: + async with sse_client(url, headers=headers) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + tools_info = [] + for tool in tools_result.tools: + tool_info = { + "name": tool.name, + "description": tool.description, + "input_schema": tool.inputSchema + } + tools_info.append(tool_info) + + for tool_info in tools_info: + tools.append({ + "name": tool_info["name"], + "description": tool_info["description"], + "inputSchema": tool_info["input_schema"] + }) + except TypeError as e: + if "unexpected keyword argument" in str(e): + async with sse_client(url) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + tools_result = await session.list_tools() + tools_info = [] + for tool in tools_result.tools: + tool_info = { + "name": tool.name, + "description": tool.description, + "input_schema": tool.inputSchema + } + tools_info.append(tool_info) + + for tool_info in tools_info: + tools.append({ + "name": tool_info["name"], + "description": tool_info["description"], + "inputSchema": tool_info["input_schema"] + }) + else: + raise + except asyncio.TimeoutError: + raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond") + except Exception as e: + logger.error(f"Error connecting to SSE MCP server: {e}") + raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}") + else: + raise HTTPException(status_code=400, detail="Invalid server type. Must be 'http' or 'sse'") + + response_data = {"tools": tools, "count": len(tools)} + + if server_name: + response_data["serverName"] = server_name + + logger.info(f"Returning {len(tools)} tools for server {server_name}") + return response_data diff --git a/app/services/mcp_temp.py b/app/services/mcp_temp.py new file mode 100644 index 0000000..71d3c01 --- /dev/null +++ b/app/services/mcp_temp.py @@ -0,0 +1,299 @@ +import os +import sys +import json +import asyncio +import subprocess +from typing import Dict, Any +from concurrent.futures import ThreadPoolExecutor +from fastapi import HTTPException # type: ignore +from utils.logger import logger +from mcp import ClientSession +from mcp.client.sse import sse_client # type: ignore +from mcp.client.streamable_http import streamablehttp_client # type: ignore + +windows_executor = ThreadPoolExecutor(max_workers=4) + +# def run_mcp_stdio_sync(command, args, env_vars, timeout=30): +# try: +# env = os.environ.copy() +# env.update(env_vars) + +# full_command = [command] + args + +# process = subprocess.Popen( +# full_command, +# stdin=subprocess.PIPE, +# stdout=subprocess.PIPE, +# stderr=subprocess.PIPE, +# env=env, +# text=True, +# bufsize=0, +# creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0 +# ) + +# init_request = { +# "jsonrpc": "2.0", +# "id": 1, +# "method": "initialize", +# "params": { +# "protocolVersion": "2024-11-05", +# "capabilities": {}, +# "clientInfo": {"name": "mcp-client", "version": "1.0.0"} +# } +# } + +# process.stdin.write(json.dumps(init_request) + "\n") +# process.stdin.flush() + +# init_response_line = process.stdout.readline().strip() +# if not init_response_line: +# raise Exception("No response from MCP server during initialization") + +# init_response = json.loads(init_response_line) + +# init_notification = { +# "jsonrpc": "2.0", +# "method": "notifications/initialized" +# } +# process.stdin.write(json.dumps(init_notification) + "\n") +# process.stdin.flush() + +# tools_request = { +# "jsonrpc": "2.0", +# "id": 2, +# "method": "tools/list", +# "params": {} +# } + +# process.stdin.write(json.dumps(tools_request) + "\n") +# process.stdin.flush() + +# tools_response_line = process.stdout.readline().strip() +# if not tools_response_line: +# raise Exception("No response from MCP server for tools list") + +# tools_response = json.loads(tools_response_line) + +# tools_info = [] +# if "result" in tools_response and "tools" in tools_response["result"]: +# for tool in tools_response["result"]["tools"]: +# tool_info = { +# "name": tool["name"], +# "description": tool.get("description", ""), +# "input_schema": tool.get("inputSchema", {}) +# } +# tools_info.append(tool_info) + +# return { +# "status": "connected", +# "transport": "stdio", +# "tools": tools_info +# } + +# except subprocess.TimeoutExpired: +# return { +# "status": "error", +# "error": f"Process timeout after {timeout} seconds", +# "tools": [] +# } +# except json.JSONDecodeError as e: +# return { +# "status": "error", +# "error": f"Invalid JSON response: {str(e)}", +# "tools": [] +# } +# except Exception as e: +# return { +# "status": "error", +# "error": str(e), +# "tools": [] +# } +# finally: +# try: +# if 'process' in locals(): +# process.terminate() +# process.wait(timeout=5) +# except: +# pass + + +# async def connect_stdio_server_windows(server_name, server_config, all_tools, timeout): +# """Windows-compatible stdio connection using subprocess""" + +# logger.info(f"Connecting to {server_name} using Windows subprocess method") + +# command = server_config["command"] +# args = server_config.get("args", []) +# env_vars = server_config.get("env", {}) + +# loop = asyncio.get_event_loop() +# result = await loop.run_in_executor( +# windows_executor, +# run_mcp_stdio_sync, +# command, +# args, +# env_vars, +# timeout +# ) + +# all_tools[server_name] = result + +# if result["status"] == "connected": +# logger.info(f" {server_name}: Connected via Windows subprocess ({len(result['tools'])} tools)") +# else: +# logger.error(f" {server_name}: Error - {result['error']}") + + +# async def list_mcp_tools_mixed_windows(config, timeout=15): +# all_tools = {} + +# if "mcpServers" not in config: +# return all_tools + +# mcp_servers = config["mcpServers"] + +# for server_name, server_config in mcp_servers.items(): +# logger.info(f"Connecting to MCP server: {server_name}") +# if server_config.get("disabled", False): +# all_tools[server_name] = {"status": "disabled", "tools": []} +# logger.info(f" {server_name}: Disabled") +# continue + +# try: +# await connect_stdio_server_windows(server_name, server_config, all_tools, timeout) + +# except asyncio.TimeoutError: +# all_tools[server_name] = { +# "status": "error", +# "error": f"Connection timeout after {timeout} seconds", +# "tools": [] +# } +# logger.error(f" {server_name}: Timeout after {timeout} seconds") +# except Exception as e: +# error_msg = str(e) +# all_tools[server_name] = { +# "status": "error", +# "error": error_msg, +# "tools": [] +# } +# logger.error(f" {server_name}: Error - {error_msg}") +# import traceback +# logger.debug(f"Full traceback for {server_name}: {traceback.format_exc()}") + +# return all_tools + + +async def discover_custom_tools(request_type: str, config: Dict[str, Any]): + logger.info(f"Received custom MCP discovery request: type={request_type}") + logger.debug(f"Request config: {config}") + + tools = [] + server_name = None + + # if request_type == 'json': + # try: + # all_tools = await list_mcp_tools_mixed_windows(config, timeout=30) + # if "mcpServers" in config and config["mcpServers"]: + # server_name = list(config["mcpServers"].keys())[0] + + # if server_name in all_tools: + # server_info = all_tools[server_name] + # if server_info["status"] == "connected": + # tools = server_info["tools"] + # logger.info(f"Found {len(tools)} tools for server {server_name}") + # else: + # error_msg = server_info.get("error", "Unknown error") + # logger.error(f"Server {server_name} failed: {error_msg}") + # raise HTTPException( + # status_code=400, + # detail=f"Failed to connect to MCP server '{server_name}': {error_msg}" + # ) + # else: + # logger.error(f"Server {server_name} not found in results") + # raise HTTPException(status_code=400, detail=f"Server '{server_name}' not found in results") + # else: + # logger.error("No MCP servers configured") + # raise HTTPException(status_code=400, detail="No MCP servers configured") + + # except HTTPException: + # raise + # except Exception as e: + # logger.error(f"Error connecting to stdio MCP server: {e}") + # import traceback + # logger.error(f"Full traceback: {traceback.format_exc()}") + # raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}") + + # if request_type == 'http': + # if 'url' not in config: + # raise HTTPException(status_code=400, detail="HTTP configuration must include 'url' field") + # url = config['url'] + # await connect_streamable_http_server(url) + # tools = await connect_streamable_http_server(url) + + # elif request_type == 'sse': + # if 'url' not in config: + # raise HTTPException(status_code=400, detail="SSE configuration must include 'url' field") + + # url = config['url'] + # headers = config.get('headers', {}) + + # try: + # async with asyncio.timeout(15): + # try: + # async with sse_client(url, headers=headers) as (read, write): + # async with ClientSession(read, write) as session: + # await session.initialize() + # tools_result = await session.list_tools() + # tools_info = [] + # for tool in tools_result.tools: + # tool_info = { + # "name": tool.name, + # "description": tool.description, + # "input_schema": tool.inputSchema + # } + # tools_info.append(tool_info) + + # for tool_info in tools_info: + # tools.append({ + # "name": tool_info["name"], + # "description": tool_info["description"], + # "inputSchema": tool_info["input_schema"] + # }) + # except TypeError as e: + # if "unexpected keyword argument" in str(e): + # async with sse_client(url) as (read, write): + # async with ClientSession(read, write) as session: + # await session.initialize() + # tools_result = await session.list_tools() + # tools_info = [] + # for tool in tools_result.tools: + # tool_info = { + # "name": tool.name, + # "description": tool.description, + # "input_schema": tool.inputSchema + # } + # tools_info.append(tool_info) + + # for tool_info in tools_info: + # tools.append({ + # "name": tool_info["name"], + # "description": tool_info["description"], + # "inputSchema": tool_info["input_schema"] + # }) + # else: + # raise + # except asyncio.TimeoutError: + # raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond") + # except Exception as e: + # logger.error(f"Error connecting to SSE MCP server: {e}") + # raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}") + # else: + # raise HTTPException(status_code=400, detail="Invalid server type. Must be 'json' or 'sse'") + + # response_data = {"tools": tools, "count": len(tools)} + + # if server_name: + # response_data["serverName"] = server_name + + # logger.info(f"Returning {len(tools)} tools for server {server_name}") + # return response_data diff --git a/app/services/redis.py b/app/services/redis.py new file mode 100644 index 0000000..1259879 --- /dev/null +++ b/app/services/redis.py @@ -0,0 +1,153 @@ +import redis.asyncio as redis +import os +from dotenv import load_dotenv +import asyncio +from utils.logger import logger +from typing import List, Any +from utils.retry import retry + +# Redis client +client: redis.Redis | None = None +_initialized = False +_init_lock = asyncio.Lock() + +# Constants +REDIS_KEY_TTL = 3600 * 24 # 24 hour TTL as safety mechanism + + +def initialize(): + """Initialize Redis connection using environment variables.""" + global client + + # Load environment variables if not already loaded + load_dotenv() + + # Get Redis configuration + redis_host = os.getenv("REDIS_HOST", "redis") + redis_port = int(os.getenv("REDIS_PORT", 6379)) + redis_password = os.getenv("REDIS_PASSWORD", "") + # Convert string 'True'/'False' to boolean + redis_ssl_str = os.getenv("REDIS_SSL", "False") + redis_ssl = redis_ssl_str.lower() == "true" + + logger.info(f"Initializing Redis connection to {redis_host}:{redis_port}") + + # Create Redis client with basic configuration + client = redis.Redis( + host=redis_host, + port=redis_port, + password=redis_password, + ssl=redis_ssl, + decode_responses=True, + socket_timeout=5.0, + socket_connect_timeout=5.0, + retry_on_timeout=True, + health_check_interval=30, + ) + + return client + + +async def initialize_async(): + """Initialize Redis connection asynchronously.""" + global client, _initialized + + async with _init_lock: + if not _initialized: + logger.info("Initializing Redis connection") + initialize() + + try: + await client.ping() + logger.info("Successfully connected to Redis") + _initialized = True + except Exception as e: + logger.error(f"Failed to connect to Redis: {e}") + client = None + _initialized = False + raise + + return client + + +async def close(): + """Close Redis connection.""" + global client, _initialized + if client: + logger.info("Closing Redis connection") + await client.aclose() + client = None + _initialized = False + logger.info("Redis connection closed") + + +async def get_client(): + """Get the Redis client, initializing if necessary.""" + global client, _initialized + if client is None or not _initialized: + await retry(lambda: initialize_async()) + return client + + +# Basic Redis operations +async def set(key: str, value: str, ex: int = None, nx: bool = False): + """Set a Redis key.""" + redis_client = await get_client() + return await redis_client.set(key, value, ex=ex, nx=nx) + + +async def get(key: str, default: str = None): + """Get a Redis key.""" + redis_client = await get_client() + result = await redis_client.get(key) + return result if result is not None else default + + +async def delete(key: str): + """Delete a Redis key.""" + redis_client = await get_client() + return await redis_client.delete(key) + + +async def publish(channel: str, message: str): + """Publish a message to a Redis channel.""" + redis_client = await get_client() + return await redis_client.publish(channel, message) + + +async def create_pubsub(): + """Create a Redis pubsub object.""" + redis_client = await get_client() + return redis_client.pubsub() + + +# List operations +async def rpush(key: str, *values: Any): + """Append one or more values to a list.""" + redis_client = await get_client() + return await redis_client.rpush(key, *values) + + +async def lrange(key: str, start: int, end: int) -> List[str]: + """Get a range of elements from a list.""" + redis_client = await get_client() + return await redis_client.lrange(key, start, end) + + +async def llen(key: str) -> int: + """Get the length of a list.""" + redis_client = await get_client() + return await redis_client.llen(key) + + +# Key management +async def expire(key: str, time: int): + """Set a key's time to live in seconds.""" + redis_client = await get_client() + return await redis_client.expire(key, time) + + +async def keys(pattern: str) -> List[str]: + """Get keys matching a pattern.""" + redis_client = await get_client() + return await redis_client.keys(pattern) diff --git a/app/services/supabase.py b/app/services/supabase.py new file mode 100644 index 0000000..0a3f855 --- /dev/null +++ b/app/services/supabase.py @@ -0,0 +1,113 @@ +""" +Centralized database connection management for AgentPress using Supabase. +""" + +from typing import Optional +from supabase import create_async_client, AsyncClient +from utils.logger import logger +from utils.config import config +import base64 +import uuid +from datetime import datetime + +class DBConnection: + """Singleton database connection manager using Supabase.""" + + _instance: Optional['DBConnection'] = None + _initialized = False + _client: Optional[AsyncClient] = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + """No initialization needed in __init__ as it's handled in __new__""" + pass + + async def initialize(self): + """Initialize the database connection.""" + if self._initialized: + return + + try: + supabase_url = config.SUPABASE_URL + # Use service role key preferentially for backend operations + supabase_key = config.SUPABASE_SERVICE_ROLE_KEY or config.SUPABASE_ANON_KEY + + if not supabase_url or not supabase_key: + logger.error("Missing required environment variables for Supabase connection") + raise RuntimeError("SUPABASE_URL and a key (SERVICE_ROLE_KEY or ANON_KEY) environment variables must be set.") + + logger.debug("Initializing Supabase connection") + self._client = await create_async_client(supabase_url, supabase_key) + self._initialized = True + key_type = "SERVICE_ROLE_KEY" if config.SUPABASE_SERVICE_ROLE_KEY else "ANON_KEY" + logger.debug(f"Database connection initialized with Supabase using {key_type}") + except Exception as e: + logger.error(f"Database initialization error: {e}") + raise RuntimeError(f"Failed to initialize database connection: {str(e)}") + + @classmethod + async def disconnect(cls): + """Disconnect from the database.""" + if cls._client: + logger.info("Disconnecting from Supabase database") + await cls._client.close() + cls._initialized = False + logger.info("Database disconnected successfully") + + @property + async def client(self) -> AsyncClient: + """Get the Supabase client instance.""" + if not self._initialized: + logger.debug("Supabase client not initialized, initializing now") + await self.initialize() + if not self._client: + logger.error("Database client is None after initialization") + raise RuntimeError("Database not initialized") + return self._client + + async def upload_base64_image(self, base64_data: str, bucket_name: str = "browser-screenshots") -> str: + """Upload a base64 encoded image to Supabase storage and return the URL. + + Args: + base64_data (str): Base64 encoded image data (with or without data URL prefix) + bucket_name (str): Name of the storage bucket to upload to + + Returns: + str: Public URL of the uploaded image + """ + try: + # Remove data URL prefix if present + if base64_data.startswith('data:'): + base64_data = base64_data.split(',')[1] + + # Decode base64 data + image_data = base64.b64decode(base64_data) + + # Generate unique filename + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') + unique_id = str(uuid.uuid4())[:8] + filename = f"image_{timestamp}_{unique_id}.png" + + # Upload to Supabase storage + client = await self.client + storage_response = await client.storage.from_(bucket_name).upload( + filename, + image_data, + {"content-type": "image/png"} + ) + + # Get public URL + public_url = await client.storage.from_(bucket_name).get_public_url(filename) + + logger.debug(f"Successfully uploaded image to {public_url}") + return public_url + + except Exception as e: + logger.error(f"Error uploading base64 image: {e}") + raise RuntimeError(f"Failed to upload image: {str(e)}") + + diff --git a/app/services/transcription.py b/app/services/transcription.py new file mode 100644 index 0000000..2a40eec --- /dev/null +++ b/app/services/transcription.py @@ -0,0 +1,76 @@ +import os +import openai +import tempfile +from fastapi import APIRouter, UploadFile, File, HTTPException, Depends +from pydantic import BaseModel +from typing import Optional +from utils.logger import logger +from utils.auth_utils import get_current_user_id_from_jwt + +router = APIRouter(tags=["transcription"]) + +class TranscriptionResponse(BaseModel): + text: str + +@router.post("/transcription", response_model=TranscriptionResponse) +async def transcribe_audio( + audio_file: UploadFile = File(...), + user_id: str = Depends(get_current_user_id_from_jwt) +): + """Transcribe audio file to text using OpenAI Whisper.""" + try: + # Validate file type - OpenAI supports these formats + allowed_types = [ + 'audio/mp3', 'audio/mpeg', 'audio/mp4', 'audio/m4a', + 'audio/wav', 'audio/webm', 'audio/mpga' + ] + + logger.info(f"Received audio file: {audio_file.filename}, content_type: {audio_file.content_type}") + + if audio_file.content_type not in allowed_types: + raise HTTPException( + status_code=400, + detail=f"Unsupported file type: {audio_file.content_type}. Supported types: {', '.join(allowed_types)}" + ) + + # Check file size (25MB limit) + content = await audio_file.read() + if len(content) > 25 * 1024 * 1024: # 25MB + raise HTTPException(status_code=400, detail="File size exceeds 25MB limit") + + # Reset file pointer + await audio_file.seek(0) + + # Initialize OpenAI client + client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")) + + # Create a temporary file with the correct extension + file_extension = audio_file.filename.split('.')[-1] if audio_file.filename and '.' in audio_file.filename else 'webm' + + with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{file_extension}') as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + try: + # Transcribe audio using the temporary file + # OpenAI Whisper API has built-in limits: 25MB file size and handles duration limits internally + with open(temp_file_path, 'rb') as f: + transcription = client.audio.transcriptions.create( + model="gpt-4o-mini-transcribe", + file=f, + response_format="text" + ) + + logger.info(f"Successfully transcribed audio for user {user_id}") + return TranscriptionResponse(text=transcription) + + finally: + # Clean up temporary file + try: + os.unlink(temp_file_path) + except Exception as e: + logger.warning(f"Failed to delete temporary file {temp_file_path}: {e}") + + except Exception as e: + logger.error(f"Error transcribing audio for user {user_id}: {str(e)}") + raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}") \ No newline at end of file diff --git a/app/tool/base.py b/app/tool/base.py index ba4084d..f9574c1 100644 --- a/app/tool/base.py +++ b/app/tool/base.py @@ -1,16 +1,75 @@ from abc import ABC, abstractmethod -from typing import Any, Dict, Optional - from pydantic import BaseModel, Field +from typing import Any, Dict, Optional, Union +import inspect +import json +from utils.logger import logger + +# class BaseTool(ABC, BaseModel): +# name: str +# description: str +# parameters: Optional[dict] = None + +# class Config: +# arbitrary_types_allowed = True + +# async def __call__(self, **kwargs) -> Any: +# """Execute the tool with given parameters.""" +# return await self.execute(**kwargs) + +# @abstractmethod +# async def execute(self, **kwargs) -> Any: +# """Execute the tool with given parameters.""" + +# def to_param(self) -> Dict: +# """Convert tool to function call format.""" +# return { +# "type": "function", +# "function": { +# "name": self.name, +# "description": self.description, +# "parameters": self.parameters, +# }, +# } class BaseTool(ABC, BaseModel): + """Consolidated base class for all tools combining BaseModel and Tool functionality. + + Provides: + - Pydantic model validation + - Schema registration + - Standardized result handling + - Abstract execution interface + + Attributes: + name (str): Tool name + description (str): Tool description + parameters (dict): Tool parameters schema + _schemas (Dict[str, List[ToolSchema]]): Registered method schemas + """ + name: str description: str parameters: Optional[dict] = None + # _schemas: Dict[str, List[ToolSchema]] = {} class Config: arbitrary_types_allowed = True + underscore_attrs_are_private = False + + # def __init__(self, **data): + # """Initialize tool with model validation and schema registration.""" + # super().__init__(**data) + # logger.debug(f"Initializing tool class: {self.__class__.__name__}") + # self._register_schemas() + + # def _register_schemas(self): + # """Register schemas from all decorated methods.""" + # for name, method in inspect.getmembers(self, predicate=inspect.ismethod): + # if hasattr(method, 'tool_schemas'): + # self._schemas[name] = method.tool_schemas + # logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}") async def __call__(self, **kwargs) -> Any: """Execute the tool with given parameters.""" @@ -21,7 +80,11 @@ class BaseTool(ABC, BaseModel): """Execute the tool with given parameters.""" def to_param(self) -> Dict: - """Convert tool to function call format.""" + """Convert tool to function call format. + + Returns: + Dictionary with tool metadata in OpenAI function calling format + """ return { "type": "function", "function": { @@ -31,6 +94,42 @@ class BaseTool(ABC, BaseModel): }, } + # def get_schemas(self) -> Dict[str, List[ToolSchema]]: + # """Get all registered tool schemas. + + # Returns: + # Dict mapping method names to their schema definitions + # """ + # return self._schemas + + def success_response(self, data: Union[Dict[str, Any], str]) -> ToolResult: + """Create a successful tool result. + + Args: + data: Result data (dictionary or string) + + Returns: + ToolResult with success=True and formatted output + """ + if isinstance(data, str): + text = data + else: + text = json.dumps(data, indent=2) + logger.debug(f"Created success response for {self.__class__.__name__}") + return ToolResult(output=text) + + def fail_response(self, msg: str) -> ToolResult: + """Create a failed tool result. + + Args: + msg: Error message describing the failure + + Returns: + ToolResult with success=False and error message + """ + logger.debug(f"Tool {self.__class__.__name__} returned failed result: {msg}") + return ToolResult(error=msg) + class ToolResult(BaseModel): """Represents the result of a tool execution.""" diff --git a/app/tool/data_providers_tool.py b/app/tool/data_providers_tool.py deleted file mode 100644 index f86ab28..0000000 --- a/app/tool/data_providers_tool.py +++ /dev/null @@ -1,188 +0,0 @@ -import json -from typing import Union, Dict, Any - -from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema -from agent.tools.data_providers.LinkedinProvider import LinkedinProvider -from agent.tools.data_providers.YahooFinanceProvider import YahooFinanceProvider -from agent.tools.data_providers.AmazonProvider import AmazonProvider -from agent.tools.data_providers.ZillowProvider import ZillowProvider -from agent.tools.data_providers.TwitterProvider import TwitterProvider - -class DataProvidersTool(Tool): - """Tool for making requests to various data providers.""" - - def __init__(self): - super().__init__() - - self.register_data_providers = { - "linkedin": LinkedinProvider(), - "yahoo_finance": YahooFinanceProvider(), - "amazon": AmazonProvider(), - "zillow": ZillowProvider(), - "twitter": TwitterProvider() - } - - @openapi_schema({ - "type": "function", - "function": { - "name": "get_data_provider_endpoints", - "description": "Get available endpoints for a specific data provider", - "parameters": { - "type": "object", - "properties": { - "service_name": { - "type": "string", - "description": "The name of the data provider (e.g., 'linkedin', 'twitter', 'zillow', 'amazon', 'yahoo_finance')" - } - }, - "required": ["service_name"] - } - } - }) - @xml_schema( - tag_name="get-data-provider-endpoints", - mappings=[ - {"param_name": "service_name", "node_type": "attribute", "path": "."} - ], - example=''' - - - - - -linkedin - - - ''' - ) - async def get_data_provider_endpoints( - self, - service_name: str - ) -> ToolResult: - """ - Get available endpoints for a specific data provider. - - Parameters: - - service_name: The name of the data provider (e.g., 'linkedin') - """ - try: - if not service_name: - return self.fail_response("Data provider name is required.") - - if service_name not in self.register_data_providers: - return self.fail_response(f"Data provider '{service_name}' not found. Available data providers: {list(self.register_data_providers.keys())}") - - endpoints = self.register_data_providers[service_name].get_endpoints() - return self.success_response(endpoints) - - except Exception as e: - error_message = str(e) - simplified_message = f"Error getting data provider endpoints: {error_message[:200]}" - if len(error_message) > 200: - simplified_message += "..." - return self.fail_response(simplified_message) - - @openapi_schema({ - "type": "function", - "function": { - "name": "execute_data_provider_call", - "description": "Execute a call to a specific data provider endpoint", - "parameters": { - "type": "object", - "properties": { - "service_name": { - "type": "string", - "description": "The name of the API service (e.g., 'linkedin')" - }, - "route": { - "type": "string", - "description": "The key of the endpoint to call" - }, - "payload": { - "type": "object", - "description": "The payload to send with the API call" - } - }, - "required": ["service_name", "route"] - } - } - }) - @xml_schema( - tag_name="execute-data-provider-call", - mappings=[ - {"param_name": "service_name", "node_type": "attribute", "path": "service_name"}, - {"param_name": "route", "node_type": "attribute", "path": "route"}, - {"param_name": "payload", "node_type": "content", "path": "."} - ], - example=''' - - - - - - linkedin - person - {"link": "https://www.linkedin.com/in/johndoe/"} - - - ''' - ) - async def execute_data_provider_call( - self, - service_name: str, - route: str, - payload: Union[Dict[str, Any], str, None] = None - ) -> ToolResult: - """ - Execute a call to a specific data provider endpoint. - - Parameters: - - service_name: The name of the data provider (e.g., 'linkedin') - - route: The key of the endpoint to call - - payload: The payload to send with the data provider call (dict or JSON string) - """ - try: - # Handle payload - it can be either a dict or a JSON string - if isinstance(payload, str): - try: - payload = json.loads(payload) - except json.JSONDecodeError as e: - return self.fail_response(f"Invalid JSON in payload: {str(e)}") - elif payload is None: - payload = {} - # If payload is already a dict, use it as-is - - if not service_name: - return self.fail_response("service_name is required.") - - if not route: - return self.fail_response("route is required.") - - if service_name not in self.register_data_providers: - return self.fail_response(f"API '{service_name}' not found. Available APIs: {list(self.register_data_providers.keys())}") - - data_provider = self.register_data_providers[service_name] - if route == service_name: - return self.fail_response(f"route '{route}' is the same as service_name '{service_name}'. YOU FUCKING IDIOT!") - - if route not in data_provider.get_endpoints().keys(): - return self.fail_response(f"Endpoint '{route}' not found in {service_name} data provider.") - - - result = data_provider.call_endpoint(route, payload) - return self.success_response(result) - - except Exception as e: - error_message = str(e) - print(error_message) - simplified_message = f"Error executing data provider call: {error_message[:200]}" - if len(error_message) > 200: - simplified_message += "..." - return self.fail_response(simplified_message) diff --git a/app/tool/expand_msg_tool.py b/app/tool/expand_msg_tool.py deleted file mode 100644 index 9d65326..0000000 --- a/app/tool/expand_msg_tool.py +++ /dev/null @@ -1,103 +0,0 @@ -from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema -from agentpress.thread_manager import ThreadManager -import json - -class ExpandMessageTool(Tool): - """Tool for expanding a previous message to the user.""" - - def __init__(self, thread_id: str, thread_manager: ThreadManager): - super().__init__() - self.thread_manager = thread_manager - self.thread_id = thread_id - - @openapi_schema({ - "type": "function", - "function": { - "name": "expand_message", - "description": "Expand a message from the previous conversation with the user. Use this tool to expand a message that was truncated in the earlier conversation.", - "parameters": { - "type": "object", - "properties": { - "message_id": { - "type": "string", - "description": "The ID of the message to expand. Must be a UUID." - } - }, - "required": ["message_id"] - } - } - }) - @xml_schema( - tag_name="expand-message", - mappings=[ - {"param_name": "message_id", "node_type": "attribute", "path": "."} - ], - example=''' - - - - ecde3a4c-c7dc-4776-ae5c-8209517c5576 - - - - - - - f47ac10b-58cc-4372-a567-0e02b2c3d479 - - - - - - - 550e8400-e29b-41d4-a716-446655440000 - - - ''' - ) - async def expand_message(self, message_id: str) -> ToolResult: - """Expand a message from the previous conversation with the user. - - Args: - message_id: The ID of the message to expand - - Returns: - ToolResult indicating the message was successfully expanded - """ - try: - client = await self.thread_manager.db.client - message = await client.table('messages').select('*').eq('message_id', message_id).eq('thread_id', self.thread_id).execute() - - if not message.data or len(message.data) == 0: - return self.fail_response(f"Message with ID {message_id} not found in thread {self.thread_id}") - - message_data = message.data[0] - message_content = message_data['content'] - final_content = message_content - if isinstance(message_content, dict) and 'content' in message_content: - final_content = message_content['content'] - elif isinstance(message_content, str): - try: - parsed_content = json.loads(message_content) - if isinstance(parsed_content, dict) and 'content' in parsed_content: - final_content = parsed_content['content'] - except json.JSONDecodeError: - pass - - return self.success_response({"status": "Message expanded successfully.", "message": final_content}) - except Exception as e: - return self.fail_response(f"Error expanding message: {str(e)}") - -if __name__ == "__main__": - import asyncio - - async def test_expand_message_tool(): - expand_message_tool = ExpandMessageTool() - - # Test expand message - expand_message_result = await expand_message_tool.expand_message( - message_id="004ab969-ef9a-4656-8aba-e392345227cd" - ) - print("Expand message result:", expand_message_result) - - asyncio.run(test_expand_message_tool()) \ No newline at end of file diff --git a/app/tool/mcp_tool_wrapper.py b/app/tool/mcp_tool_wrapper.py deleted file mode 100644 index 31f3312..0000000 --- a/app/tool/mcp_tool_wrapper.py +++ /dev/null @@ -1,715 +0,0 @@ -""" -MCP Tool Wrapper for AgentPress - -This module provides a generic tool wrapper that handles all MCP (Model Context Protocol) -server tool calls through dynamically generated individual function methods. -""" - -import json -from typing import Any, Dict, List, Optional -from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema, ToolSchema, SchemaType -from mcp_local.client import MCPManager -from utils.logger import logger -import inspect -from mcp import ClientSession -from mcp.client.sse import sse_client -from mcp.client.stdio import stdio_client -from mcp.client.streamable_http import streamablehttp_client -from mcp import StdioServerParameters -import asyncio - - -class MCPToolWrapper(Tool): - """ - A generic tool wrapper that dynamically creates individual methods for each MCP tool. - - This tool creates separate function calls for each MCP tool while routing them all - through the same underlying implementation. - """ - - def __init__(self, mcp_configs: Optional[List[Dict[str, Any]]] = None): - """ - Initialize the MCP tool wrapper. - - Args: - mcp_configs: List of MCP configurations from agent's configured_mcps - """ - # Don't call super().__init__() yet - we need to set up dynamic methods first - self.mcp_manager = MCPManager() - self.mcp_configs = mcp_configs or [] - self._initialized = False - self._dynamic_tools = {} - self._schemas: Dict[str, List[ToolSchema]] = {} - self._custom_tools = {} # Store custom MCP tools separately - - # Now initialize the parent class which will call _register_schemas - super().__init__() - - async def _ensure_initialized(self): - """Ensure MCP servers are initialized.""" - if not self._initialized: - # Initialize standard MCP servers from Smithery - standard_configs = [cfg for cfg in self.mcp_configs if not cfg.get('isCustom', False)] - custom_configs = [cfg for cfg in self.mcp_configs if cfg.get('isCustom', False)] - - # Initialize standard MCPs through MCPManager - if standard_configs: - for config in standard_configs: - try: - logger.info(f"Attempting to connect to MCP server: {config['qualifiedName']}") - await self.mcp_manager.connect_server(config) - logger.info(f"Successfully connected to MCP server: {config['qualifiedName']}") - except Exception as e: - logger.error(f"Failed to connect to MCP server {config['qualifiedName']}: {e}") - import traceback - logger.error(f"Full traceback: {traceback.format_exc()}") - - # Initialize custom MCPs directly - if custom_configs: - await self._initialize_custom_mcps(custom_configs) - - # Create dynamic tools for all connected servers - await self._create_dynamic_tools() - self._initialized = True - - async def _connect_sse_server(self, server_name, server_config, all_tools, timeout): - url = server_config["url"] - headers = server_config.get("headers", {}) - - async with asyncio.timeout(timeout): - try: - async with sse_client(url, headers=headers) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - tools_result = await session.list_tools() - tools_info = [] - for tool in tools_result.tools: - tool_info = { - "name": tool.name, - "description": tool.description, - "input_schema": tool.inputSchema - } - tools_info.append(tool_info) - - all_tools[server_name] = { - "status": "connected", - "transport": "sse", - "url": url, - "tools": tools_info - } - - logger.info(f" {server_name}: Connected via SSE ({len(tools_info)} tools)") - except TypeError as e: - if "unexpected keyword argument" in str(e): - async with sse_client(url) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - tools_result = await session.list_tools() - tools_info = [] - for tool in tools_result.tools: - tool_info = { - "name": tool.name, - "description": tool.description, - "input_schema": tool.inputSchema - } - tools_info.append(tool_info) - - all_tools[server_name] = { - "status": "connected", - "transport": "sse", - "url": url, - "tools": tools_info - } - logger.info(f" {server_name}: Connected via SSE ({len(tools_info)} tools)") - else: - raise - - async def _connect_streamable_http_server(self, url): - async with streamablehttp_client(url) as ( - read_stream, - write_stream, - _, - ): - async with ClientSession(read_stream, write_stream) as session: - await session.initialize() - tool_result = await session.list_tools() - print(f"Connected via HTTP ({len(tool_result.tools)} tools)") - - tools_info = [] - for tool in tool_result.tools: - tool_info = { - "name": tool.name, - "description": tool.description, - "inputSchema": tool.inputSchema - } - tools_info.append(tool_info) - - return tools_info - - async def _connect_stdio_server(self, server_name, server_config, all_tools, timeout): - """Connect to a stdio-based MCP server.""" - server_params = StdioServerParameters( - command=server_config["command"], - args=server_config.get("args", []), - env=server_config.get("env", {}) - ) - - async with asyncio.timeout(timeout): - async with stdio_client(server_params) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - tools_result = await session.list_tools() - tools_info = [] - for tool in tools_result.tools: - tool_info = { - "name": tool.name, - "description": tool.description, - "input_schema": tool.inputSchema - } - tools_info.append(tool_info) - - all_tools[server_name] = { - "status": "connected", - "transport": "stdio", - "tools": tools_info - } - - logger.info(f" {server_name}: Connected via stdio ({len(tools_info)} tools)") - - async def _initialize_custom_mcps(self, custom_configs): - """Initialize custom MCP servers.""" - for config in custom_configs: - try: - logger.info(f"Initializing custom MCP: {config}") - custom_type = config.get('customType', 'sse') - server_config = config.get('config', {}) - enabled_tools = config.get('enabledTools', []) - server_name = config.get('name', 'Unknown') - - logger.info(f"Initializing custom MCP: {server_name} (type: {custom_type})") - - if custom_type == 'sse': - if 'url' not in server_config: - logger.error(f"Custom MCP {server_name}: Missing 'url' in config") - continue - - url = server_config['url'] - logger.info(f"Initializing custom MCP {url} with SSE type") - - try: - # Use the working connect_sse_server method - all_tools = {} - await self._connect_sse_server(server_name, server_config, all_tools, 15) - - # Process the results - if server_name in all_tools and all_tools[server_name].get('status') == 'connected': - tools_info = all_tools[server_name].get('tools', []) - tools_registered = 0 - - for tool_info in tools_info: - tool_name_from_server = tool_info['name'] - if not enabled_tools or tool_name_from_server in enabled_tools: - tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}" - self._custom_tools[tool_name] = { - 'name': tool_name, - 'description': tool_info['description'], - 'parameters': tool_info['input_schema'], - 'server': server_name, - 'original_name': tool_name_from_server, - 'is_custom': True, - 'custom_type': custom_type, - 'custom_config': server_config - } - tools_registered += 1 - logger.debug(f"Registered custom tool: {tool_name}") - - logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools") - else: - logger.error(f"Failed to connect to custom MCP {server_name}") - - except Exception as e: - logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}") - continue - - elif custom_type == 'http': - if 'url' not in server_config: - logger.error(f"Custom MCP {server_name}: Missing 'url' in config") - continue - - url = server_config['url'] - logger.info(f"Initializing custom MCP {url} with HTTP type") - - try: - - tools_info = await self._connect_streamable_http_server(url) - tools_registered = 0 - - for tool_info in tools_info: - tool_name_from_server = tool_info['name'] - if not enabled_tools or tool_name_from_server in enabled_tools: - tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}" - self._custom_tools[tool_name] = { - 'name': tool_name, - 'description': tool_info['description'], - 'parameters': tool_info['inputSchema'], - 'server': server_name, - 'original_name': tool_name_from_server, - 'is_custom': True, - 'custom_type': custom_type, - 'custom_config': server_config - } - tools_registered += 1 - logger.debug(f"Registered custom tool: {tool_name}") - - logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools") - - except Exception as e: - logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}") - continue - - elif custom_type == 'json': - if 'command' not in server_config: - logger.error(f"Custom MCP {server_name}: Missing 'command' in config") - continue - - logger.info(f"Initializing custom MCP {server_name} with JSON/stdio type") - - try: - # Use the stdio connection method - all_tools = {} - await self._connect_stdio_server(server_name, server_config, all_tools, 15) - - # Process the results - if server_name in all_tools and all_tools[server_name].get('status') == 'connected': - tools_info = all_tools[server_name].get('tools', []) - tools_registered = 0 - - for tool_info in tools_info: - tool_name_from_server = tool_info['name'] - if not enabled_tools or tool_name_from_server in enabled_tools: - tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}" - self._custom_tools[tool_name] = { - 'name': tool_name, - 'description': tool_info['description'], - 'parameters': tool_info['input_schema'], - 'server': server_name, - 'original_name': tool_name_from_server, - 'is_custom': True, - 'custom_type': custom_type, - 'custom_config': server_config - } - tools_registered += 1 - logger.debug(f"Registered custom tool: {tool_name}") - - logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools") - else: - logger.error(f"Failed to connect to custom MCP {server_name}") - - except Exception as e: - logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}") - continue - - else: - logger.error(f"Custom MCP {server_name}: Unsupported type '{custom_type}', supported types are 'sse', 'http' and 'json'") - continue - - except Exception as e: - logger.error(f"Failed to initialize custom MCP {config.get('name', 'Unknown')}: {e}") - continue - - async def initialize_and_register_tools(self, tool_registry=None): - """Initialize MCP tools and optionally update the tool registry. - - This method should be called after the tool has been registered to dynamically - add the MCP tool schemas to the registry. - - Args: - tool_registry: Optional ToolRegistry instance to update with new schemas - """ - await self._ensure_initialized() - - if tool_registry and self._dynamic_tools: - logger.info(f"Updating tool registry with {len(self._dynamic_tools)} MCP tools") - for method_name, schemas in self._schemas.items(): - if method_name not in ['call_mcp_tool']: # Skip the fallback method - pass - - async def _create_dynamic_tools(self): - """Create dynamic tool methods for each available MCP tool.""" - try: - # Get standard MCP tools - available_tools = self.mcp_manager.get_all_tools_openapi() - logger.info(f"MCPManager returned {len(available_tools)} tools") - - for tool_info in available_tools: - tool_name = tool_info.get('name', '') - logger.info(f"Processing tool: {tool_name}") - if tool_name: - # Create a dynamic method for this tool with proper OpenAI schema - self._create_dynamic_method(tool_name, tool_info) - - # Get custom MCP tools - logger.info(f"Processing {len(self._custom_tools)} custom MCP tools") - for tool_name, tool_info in self._custom_tools.items(): - logger.info(f"Processing custom tool: {tool_name}") - # Convert custom tool info to the expected format - openapi_tool_info = { - "name": tool_name, - "description": tool_info['description'], - "parameters": tool_info['parameters'] - } - self._create_dynamic_method(tool_name, openapi_tool_info) - - logger.info(f"Created {len(self._dynamic_tools)} dynamic MCP tool methods") - - except Exception as e: - logger.error(f"Error creating dynamic MCP tools: {e}") - - def _create_dynamic_method(self, tool_name: str, tool_info: Dict[str, Any]): - """Create a dynamic method for a specific MCP tool with proper OpenAI schema.""" - if tool_name.startswith("custom_"): - if tool_name in self._custom_tools: - clean_tool_name = self._custom_tools[tool_name]['original_name'] - server_name = self._custom_tools[tool_name]['server'] - else: - parts = tool_name.split("_") - if len(parts) >= 3: - clean_tool_name = "_".join(parts[2:]) - server_name = parts[1] if len(parts) > 1 else "unknown" - else: - clean_tool_name = tool_name - server_name = "unknown" - else: - parts = tool_name.split("_", 2) - clean_tool_name = parts[2] if len(parts) > 2 else tool_name - server_name = parts[1] if len(parts) > 1 else "unknown" - - method_name = clean_tool_name.replace('-', '_') - - logger.info(f"Creating dynamic method for tool '{tool_name}': clean_tool_name='{clean_tool_name}', method_name='{method_name}', server='{server_name}'") - - original_full_name = tool_name - - # Create the dynamic method - async def dynamic_tool_method(**kwargs) -> ToolResult: - """Dynamically created method for MCP tool.""" - # Use the original full tool name for execution - return await self._execute_mcp_tool(original_full_name, kwargs) - - # Set the method name to match the tool name - dynamic_tool_method.__name__ = method_name - dynamic_tool_method.__qualname__ = f"{self.__class__.__name__}.{method_name}" - - # Build a more descriptive description - base_description = tool_info.get("description", f"MCP tool from {server_name}") - full_description = f"{base_description} (MCP Server: {server_name})" - - # Create the OpenAI schema for this tool - openapi_function_schema = { - "type": "function", - "function": { - "name": method_name, # Use the clean method name for function calling - "description": full_description, - "parameters": tool_info.get("parameters", { - "type": "object", - "properties": {}, - "required": [] - }) - } - } - - # Create a ToolSchema object - tool_schema = ToolSchema( - schema_type=SchemaType.OPENAPI, - schema=openapi_function_schema - ) - - # Add the schema to our schemas dict - self._schemas[method_name] = [tool_schema] - - # Also add the schema to the method itself (for compatibility) - dynamic_tool_method.tool_schemas = [tool_schema] - - # Store the method and its info - self._dynamic_tools[tool_name] = { - 'method': dynamic_tool_method, - 'method_name': method_name, - 'original_tool_name': tool_name, - 'clean_tool_name': clean_tool_name, - 'server_name': server_name, - 'info': tool_info, - 'schema': tool_schema - } - - # Add the method to this instance - setattr(self, method_name, dynamic_tool_method) - - logger.debug(f"Created dynamic method '{method_name}' for MCP tool '{tool_name}' from server '{server_name}'") - - def _register_schemas(self): - """Register schemas from all decorated methods and dynamic tools.""" - # First register static schemas from decorated methods - for name, method in inspect.getmembers(self, predicate=inspect.ismethod): - if hasattr(method, 'tool_schemas'): - self._schemas[name] = method.tool_schemas - logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}") - - # Note: Dynamic schemas will be added after async initialization - logger.debug(f"Initial registration complete for MCPToolWrapper") - - def get_schemas(self) -> Dict[str, List[ToolSchema]]: - """Get all registered tool schemas including dynamic ones.""" - # Return all schemas including dynamically added ones - return self._schemas - - def __getattr__(self, name: str): - """Handle calls to dynamically created MCP tool methods.""" - # Look for exact method name match first - for tool_data in self._dynamic_tools.values(): - if tool_data['method_name'] == name: - return tool_data['method'] - - # Try with underscore/hyphen conversion - name_with_hyphens = name.replace('_', '-') - for tool_name, tool_data in self._dynamic_tools.items(): - if tool_data['method_name'] == name or tool_name == name_with_hyphens: - return tool_data['method'] - - raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'") - - async def get_available_tools(self) -> List[Dict[str, Any]]: - """Get all available MCP tools in OpenAPI format.""" - await self._ensure_initialized() - return self.mcp_manager.get_all_tools_openapi() - - async def _execute_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> ToolResult: - """Execute an MCP tool call.""" - await self._ensure_initialized() - logger.info(f"Executing MCP tool {tool_name} with arguments {arguments}") - try: - # Check if it's a custom MCP tool first - if tool_name in self._custom_tools: - tool_info = self._custom_tools[tool_name] - return await self._execute_custom_mcp_tool(tool_name, arguments, tool_info) - else: - # Use standard MCP manager for Smithery servers - result = await self.mcp_manager.execute_tool(tool_name, arguments) - - if isinstance(result, dict): - if result.get('isError', False): - return self.fail_response(result.get('content', 'Tool execution failed')) - else: - return self.success_response(result.get('content', result)) - else: - return self.success_response(result) - - except Exception as e: - logger.error(f"Error executing MCP tool {tool_name}: {str(e)}") - return self.fail_response(f"Error executing tool: {str(e)}") - - async def _execute_custom_mcp_tool(self, tool_name: str, arguments: Dict[str, Any], tool_info: Dict[str, Any]) -> ToolResult: - """Execute a custom MCP tool call.""" - try: - custom_type = tool_info['custom_type'] - custom_config = tool_info['custom_config'] - original_tool_name = tool_info['original_name'] - - if custom_type == 'sse': - # Execute SSE-based custom MCP using the same pattern as _connect_sse_server - url = custom_config['url'] - headers = custom_config.get('headers', {}) - - async with asyncio.timeout(30): # 30 second timeout for tool execution - try: - # Try with headers first (same pattern as _connect_sse_server) - async with sse_client(url, headers=headers) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - result = await session.call_tool(original_tool_name, arguments) - - # Handle the result properly - if hasattr(result, 'content'): - content = result.content - if isinstance(content, list): - # Extract text from content list - text_parts = [] - for item in content: - if hasattr(item, 'text'): - text_parts.append(item.text) - else: - text_parts.append(str(item)) - content_str = "\n".join(text_parts) - elif hasattr(content, 'text'): - content_str = content.text - else: - content_str = str(content) - - return self.success_response(content_str) - else: - return self.success_response(str(result)) - - except TypeError as e: - if "unexpected keyword argument" in str(e): - # Fallback: try without headers (exact pattern from _connect_sse_server) - async with sse_client(url) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - result = await session.call_tool(original_tool_name, arguments) - - # Handle the result properly - if hasattr(result, 'content'): - content = result.content - if isinstance(content, list): - # Extract text from content list - text_parts = [] - for item in content: - if hasattr(item, 'text'): - text_parts.append(item.text) - else: - text_parts.append(str(item)) - content_str = "\n".join(text_parts) - elif hasattr(content, 'text'): - content_str = content.text - else: - content_str = str(content) - - return self.success_response(content_str) - else: - return self.success_response(str(result)) - else: - raise - - elif custom_type == 'http': - # Execute HTTP-based custom MCP - url = custom_config['url'] - - async with asyncio.timeout(30): # 30 second timeout for tool execution - async with streamablehttp_client(url) as (read, write, _): - async with ClientSession(read, write) as session: - await session.initialize() - result = await session.call_tool(original_tool_name, arguments) - - # Handle the result properly - if hasattr(result, 'content'): - content = result.content - if isinstance(content, list): - # Extract text from content list - text_parts = [] - for item in content: - if hasattr(item, 'text'): - text_parts.append(item.text) - else: - text_parts.append(str(item)) - content_str = "\n".join(text_parts) - elif hasattr(content, 'text'): - content_str = content.text - else: - content_str = str(content) - - return self.success_response(content_str) - else: - return self.success_response(str(result)) - - elif custom_type == 'json': - # Execute stdio-based custom MCP using the same pattern as _connect_stdio_server - server_params = StdioServerParameters( - command=custom_config["command"], - args=custom_config.get("args", []), - env=custom_config.get("env", {}) - ) - - async with asyncio.timeout(30): # 30 second timeout for tool execution - async with stdio_client(server_params) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - result = await session.call_tool(original_tool_name, arguments) - - # Handle the result properly - if hasattr(result, 'content'): - content = result.content - if isinstance(content, list): - # Extract text from content list - text_parts = [] - for item in content: - if hasattr(item, 'text'): - text_parts.append(item.text) - else: - text_parts.append(str(item)) - content_str = "\n".join(text_parts) - elif hasattr(content, 'text'): - content_str = content.text - else: - content_str = str(content) - - return self.success_response(content_str) - else: - return self.success_response(str(result)) - else: - return self.fail_response(f"Unsupported custom MCP type: {custom_type}") - - except asyncio.TimeoutError: - return self.fail_response(f"Tool execution timeout for {tool_name}") - except Exception as e: - logger.error(f"Error executing custom MCP tool {tool_name}: {str(e)}") - return self.fail_response(f"Error executing custom tool: {str(e)}") - - # Keep the original call_mcp_tool method as a fallback - @openapi_schema({ - "type": "function", - "function": { - "name": "call_mcp_tool", - "description": "Execute a tool from any connected MCP server. This is a fallback wrapper that forwards calls to MCP tools. The tool_name should be in the format 'mcp_{server}_{tool}' where {server} is the MCP server's qualified name and {tool} is the specific tool name.", - "parameters": { - "type": "object", - "properties": { - "tool_name": { - "type": "string", - "description": "The full MCP tool name in format 'mcp_{server}_{tool}', e.g., 'mcp_exa_web_search_exa'" - }, - "arguments": { - "type": "object", - "description": "The arguments to pass to the MCP tool, as a JSON object. The required arguments depend on the specific tool being called.", - "additionalProperties": True - } - }, - "required": ["tool_name", "arguments"] - } - } - }) - @xml_schema( - tag_name="call-mcp-tool", - mappings=[ - {"param_name": "tool_name", "node_type": "attribute", "path": "."}, - {"param_name": "arguments", "node_type": "content", "path": "."} - ], - example=''' - - - mcp_exa_web_search_exa - {"query": "latest developments in AI", "num_results": 10} - - - ''' - ) - async def call_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> ToolResult: - """ - Execute an MCP tool call (fallback method). - - Args: - tool_name: The full MCP tool name (e.g., "mcp_exa_web_search_exa") - arguments: The arguments to pass to the tool - - Returns: - ToolResult with the tool execution result - """ - return await self._execute_mcp_tool(tool_name, arguments) - - async def cleanup(self): - """Disconnect all MCP servers.""" - if self._initialized: - try: - await self.mcp_manager.disconnect_all() - except Exception as e: - logger.error(f"Error during MCP cleanup: {str(e)}") - finally: - self._initialized = False \ No newline at end of file diff --git a/app/tool/message_tool.py b/app/tool/message_tool.py deleted file mode 100644 index eef3ef5..0000000 --- a/app/tool/message_tool.py +++ /dev/null @@ -1,270 +0,0 @@ -from typing import List, Optional, Union -from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema -from utils.logger import logger - -class MessageTool(Tool): - """Tool for user communication and interaction. - - This tool provides methods for asking questions, with support for - attachments and user takeover suggestions. - """ - - def __init__(self): - super().__init__() - - # Commented out as we are just doing this via prompt as there is no need to call it as a tool - - @openapi_schema({ - "type": "function", - "function": { - "name": "ask", - "description": "Ask user a question and wait for response. Use for: 1) Requesting clarification on ambiguous requirements, 2) Seeking confirmation before proceeding with high-impact changes, 3) Gathering additional information needed to complete a task, 4) Offering options and requesting user preference, 5) Validating assumptions when critical to task success. IMPORTANT: Use this tool only when user input is essential to proceed. Always provide clear context and options when applicable. Include relevant attachments when the question relates to specific files or resources.", - "parameters": { - "type": "object", - "properties": { - "text": { - "type": "string", - "description": "Question text to present to user - should be specific and clearly indicate what information you need. Include: 1) Clear question or request, 2) Context about why the input is needed, 3) Available options if applicable, 4) Impact of different choices, 5) Any relevant constraints or considerations." - }, - "attachments": { - "anyOf": [ - {"type": "string"}, - {"items": {"type": "string"}, "type": "array"} - ], - "description": "(Optional) List of files or URLs to attach to the question. Include when: 1) Question relates to specific files or configurations, 2) User needs to review content before answering, 3) Options or choices are documented in files, 4) Supporting evidence or context is needed. Always use relative paths to /workspace directory." - } - }, - "required": ["text"] - } - } - }) - @xml_schema( - tag_name="ask", - mappings=[ - {"param_name": "text", "node_type": "content", "path": "."}, - {"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False} - ], - example=''' - - - I'm planning to bake the chocolate cake for your birthday party. The recipe mentions "rich frosting" but doesn't specify what type. Could you clarify your preferences? For example: -1. Would you prefer buttercream or cream cheese frosting? -2. Do you want any specific flavor added to the frosting (vanilla, coffee, etc.)? -3. Should I add any decorative toppings like sprinkles or fruit? -4. Do you have any dietary restrictions I should be aware of? - -This information will help me make sure the cake meets your expectations for the celebration. - recipes/chocolate_cake.txt,photos/cake_examples.jpg - - - ''' - ) - async def ask(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult: - """Ask the user a question and wait for a response. - - Args: - text: The question to present to the user - attachments: Optional file paths or URLs to attach to the question - - Returns: - ToolResult indicating the question was successfully sent - """ - try: - # Convert single attachment to list for consistent handling - if attachments and isinstance(attachments, str): - attachments = [attachments] - - return self.success_response({"status": "Awaiting user response..."}) - except Exception as e: - return self.fail_response(f"Error asking user: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "web_browser_takeover", - "description": "Request user takeover of browser interaction. Use this tool when: 1) The page requires complex human interaction that automated tools cannot handle, 2) Authentication or verification steps require human input, 3) The page has anti-bot measures that prevent automated access, 4) Complex form filling or navigation is needed, 5) The page requires human verification (CAPTCHA, etc.). IMPORTANT: This tool should be used as a last resort after web-search and crawl-webpage have failed, and when direct browser tools are insufficient. Always provide clear context about why takeover is needed and what actions the user should take.", - "parameters": { - "type": "object", - "properties": { - "text": { - "type": "string", - "description": "Instructions for the user about what actions to take in the browser. Include: 1) Clear explanation of why takeover is needed, 2) Specific steps the user should take, 3) What information to look for or extract, 4) How to indicate when they're done, 5) Any important context about the current page state." - }, - "attachments": { - "anyOf": [ - {"type": "string"}, - {"items": {"type": "string"}, "type": "array"} - ], - "description": "(Optional) List of files or URLs to attach to the takeover request. Include when: 1) Screenshots or visual references are needed, 2) Previous search results or crawled content is relevant, 3) Supporting documentation is required. Always use relative paths to /workspace directory." - } - }, - "required": ["text"] - } - } - }) - @xml_schema( - tag_name="web-browser-takeover", - mappings=[ - {"param_name": "text", "node_type": "content", "path": "."}, - {"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False} - ], - example=''' - - - I've encountered a CAPTCHA verification on the page. Please: -1. Solve the CAPTCHA puzzle -2. Let me know once you've completed it -3. I'll then continue with the automated process - -If you encounter any issues or need to take additional steps, please let me know. - - - ''' - ) - async def web_browser_takeover(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult: - """Request user takeover of browser interaction. - - Args: - text: Instructions for the user about what actions to take - attachments: Optional file paths or URLs to attach to the request - - Returns: - ToolResult indicating the takeover request was successfully sent - """ - try: - # Convert single attachment to list for consistent handling - if attachments and isinstance(attachments, str): - attachments = [attachments] - - return self.success_response({"status": "Awaiting user browser takeover..."}) - except Exception as e: - return self.fail_response(f"Error requesting browser takeover: {str(e)}") - -# @openapi_schema({ -# "type": "function", -# "function": { -# "name": "inform", -# "description": "Inform the user about progress, completion of a major step, or important context. Use this tool: 1) To provide updates between major sections of work, 2) After accomplishing significant milestones, 3) When transitioning to a new phase of work, 4) To confirm actions were completed successfully, 5) To provide context about upcoming steps. IMPORTANT: Use FREQUENTLY throughout execution to provide UI context to the user. The user CANNOT respond to this tool - they can only respond to the 'ask' tool. Use this tool to keep the user informed without requiring their input.", -# "parameters": { -# "type": "object", -# "properties": { -# "text": { -# "type": "string", -# "description": "Information to present to the user. Include: 1) Clear statement of what has been accomplished or what is happening, 2) Relevant context or impact, 3) Brief indication of next steps if applicable." -# }, -# "attachments": { -# "anyOf": [ -# {"type": "string"}, -# {"items": {"type": "string"}, "type": "array"} -# ], -# "description": "(Optional) List of files or URLs to attach to the information. Include when: 1) Information relates to specific files or resources, 2) Showing intermediate results or outputs, 3) Providing supporting documentation. Always use relative paths to /workspace directory." -# } -# }, -# "required": ["text"] -# } -# } -# }) -# @xml_schema( -# tag_name="inform", -# mappings=[ -# {"param_name": "text", "node_type": "content", "path": "."}, -# {"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False} -# ], -# example=''' - -# Inform the user about progress, completion of a major step, or important context. Use this tool: 1) To provide updates between major sections of work, 2) After accomplishing significant milestones, 3) When transitioning to a new phase of work, 4) To confirm actions were completed successfully, 5) To provide context about upcoming steps. IMPORTANT: Use FREQUENTLY throughout execution to provide UI context to the user. The user CANNOT respond to this tool - they can only respond to the 'ask' tool. Use this tool to keep the user informed without requiring their input." - -# -# -# -# -# -# -# -# -# - -# -# I've completed the data analysis of the sales figures. Key findings include: -# - Q4 sales were 28% higher than Q3 -# - Product line A showed the strongest performance -# - Three regions missed their targets - -# I'll now proceed with creating the executive summary report based on these findings. -# -# ''' -# ) -# async def inform(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult: -# """Inform the user about progress or important updates without requiring a response. - -# Args: -# text: The information to present to the user -# attachments: Optional file paths or URLs to attach - -# Returns: -# ToolResult indicating the information was successfully sent -# """ -# try: -# # Convert single attachment to list for consistent handling -# if attachments and isinstance(attachments, str): -# attachments = [attachments] - -# return self.success_response({"status": "Information sent"}) -# except Exception as e: -# return self.fail_response(f"Error informing user: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "complete", - "description": "A special tool to indicate you have completed all tasks and are about to enter complete state. Use ONLY when: 1) All tasks in todo.md are marked complete [x], 2) The user's original request has been fully addressed, 3) There are no pending actions or follow-ups required, 4) You've delivered all final outputs and results to the user. IMPORTANT: This is the ONLY way to properly terminate execution. Never use this tool unless ALL tasks are complete and verified. Always ensure you've provided all necessary outputs and references before using this tool.", - "parameters": { - "type": "object", - "properties": {}, - "required": [] - } - } - }) - @xml_schema( - tag_name="complete", - mappings=[], - example=''' - - - - - ''' - ) - async def complete(self) -> ToolResult: - """Indicate that the agent has completed all tasks and is entering complete state. - - Returns: - ToolResult indicating successful transition to complete state - """ - try: - return self.success_response({"status": "complete"}) - except Exception as e: - return self.fail_response(f"Error entering complete state: {str(e)}") - - -if __name__ == "__main__": - import asyncio - - async def test_message_tool(): - message_tool = MessageTool() - - # Test question - ask_result = await message_tool.ask( - text="Would you like to proceed with the next phase?", - attachments="summary.pdf" - ) - print("Question result:", ask_result) - - # Test inform - inform_result = await message_tool.inform( - text="Completed analysis of data. Processing results now.", - attachments="analysis.pdf" - ) - print("Inform result:", inform_result) - - asyncio.run(test_message_tool()) diff --git a/app/tool/sb_browser_tool.py b/app/tool/sb_browser_tool.py index 74bcfe6..7ebfe86 100644 --- a/app/tool/sb_browser_tool.py +++ b/app/tool/sb_browser_tool.py @@ -6,14 +6,14 @@ from PIL import Image from agentpress.tool import ToolResult, openapi_schema, xml_schema from agentpress.thread_manager import ThreadManager -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from utils.logger import logger from utils.s3_upload_utils import upload_base64_image class SandboxBrowserTool(SandboxToolsBase): """Tool for executing tasks in a Daytona sandbox with browser-use capabilities.""" - + def __init__(self, project_id: str, thread_id: str, thread_manager: ThreadManager): super().__init__(project_id, thread_manager) self.thread_id = thread_id @@ -21,11 +21,11 @@ class SandboxBrowserTool(SandboxToolsBase): def _validate_base64_image(self, base64_string: str, max_size_mb: int = 10) -> tuple[bool, str]: """ Comprehensive validation of base64 image data. - + Args: base64_string (str): The base64 encoded image data max_size_mb (int): Maximum allowed image size in megabytes - + Returns: tuple[bool, str]: (is_valid, error_message) """ @@ -33,94 +33,94 @@ class SandboxBrowserTool(SandboxToolsBase): # Check if data exists and has reasonable length if not base64_string or len(base64_string) < 10: return False, "Base64 string is empty or too short" - + # Remove data URL prefix if present (data:image/jpeg;base64,...) if base64_string.startswith('data:'): try: base64_string = base64_string.split(',', 1)[1] except (IndexError, ValueError): return False, "Invalid data URL format" - + # Check if string contains only valid base64 characters # Base64 alphabet: A-Z, a-z, 0-9, +, /, = (padding) import re if not re.match(r'^[A-Za-z0-9+/]*={0,2}$', base64_string): return False, "Invalid base64 characters detected" - + # Check if base64 string length is valid (must be multiple of 4) if len(base64_string) % 4 != 0: return False, "Invalid base64 string length" - + # Attempt to decode base64 try: image_data = base64.b64decode(base64_string, validate=True) except Exception as e: return False, f"Base64 decoding failed: {str(e)}" - + # Check decoded data size if len(image_data) == 0: return False, "Decoded image data is empty" - + # Check if decoded data size exceeds limit max_size_bytes = max_size_mb * 1024 * 1024 if len(image_data) > max_size_bytes: return False, f"Image size ({len(image_data)} bytes) exceeds limit ({max_size_bytes} bytes)" - + # Validate that decoded data is actually a valid image using PIL try: image_stream = io.BytesIO(image_data) with Image.open(image_stream) as img: # Verify the image by attempting to load it img.verify() - + # Check if image format is supported supported_formats = {'JPEG', 'PNG', 'GIF', 'BMP', 'WEBP', 'TIFF'} if img.format not in supported_formats: return False, f"Unsupported image format: {img.format}" - + # Re-open for dimension checks (verify() closes the image) image_stream.seek(0) with Image.open(image_stream) as img_check: width, height = img_check.size - + # Check reasonable dimension limits max_dimension = 8192 # 8K resolution limit if width > max_dimension or height > max_dimension: return False, f"Image dimensions ({width}x{height}) exceed limit ({max_dimension}x{max_dimension})" - + # Check minimum dimensions if width < 1 or height < 1: return False, f"Invalid image dimensions: {width}x{height}" - + logger.debug(f"Valid image detected: {img.format}, {width}x{height}, {len(image_data)} bytes") - + except Exception as e: return False, f"Invalid image data: {str(e)}" - + return True, "Valid image" - + except Exception as e: logger.error(f"Unexpected error during base64 image validation: {e}") return False, f"Validation error: {str(e)}" async def _execute_browser_action(self, endpoint: str, params: dict = None, method: str = "POST") -> ToolResult: """Execute a browser automation action through the API - + Args: endpoint (str): The API endpoint to call params (dict, optional): Parameters to send. Defaults to None. method (str, optional): HTTP method to use. Defaults to "POST". - + Returns: ToolResult: Result of the execution """ try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # Build the curl command url = f"http://localhost:8003/api/automation/{endpoint}" - + if method == "GET" and params: query_params = "&".join([f"{k}={v}" for k, v in params.items()]) url = f"{url}?{query_params}" @@ -130,19 +130,19 @@ class SandboxBrowserTool(SandboxToolsBase): if params: json_data = json.dumps(params) curl_cmd += f" -d '{json_data}'" - + logger.debug("\033[95mExecuting curl command:\033[0m") logger.debug(f"{curl_cmd}") - + response = self.sandbox.process.exec(curl_cmd, timeout=30) - + if response.exit_code == 0: try: result = json.loads(response.result) if not "content" in result: result["content"] = "" - + if not "role" in result: result["role"] = "assistant" @@ -153,7 +153,7 @@ class SandboxBrowserTool(SandboxToolsBase): # Comprehensive validation of the base64 image data screenshot_data = result["screenshot_base64"] is_valid, validation_message = self._validate_base64_image(screenshot_data) - + if is_valid: logger.debug(f"Screenshot validation passed: {validation_message}") image_url = await upload_base64_image(screenshot_data) @@ -162,10 +162,10 @@ class SandboxBrowserTool(SandboxToolsBase): else: logger.warning(f"Screenshot validation failed: {validation_message}") result["image_validation_error"] = validation_message - + # Remove base64 data from result to keep it clean del result["screenshot_base64"] - + except Exception as e: logger.error(f"Failed to process screenshot: {e}") result["image_upload_error"] = str(e) @@ -251,10 +251,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_navigate_to(self, url: str) -> ToolResult: """Navigate to a specific url - + Args: url (str): The url to navigate to - + Returns: dict: Result of the execution """ @@ -290,10 +290,10 @@ class SandboxBrowserTool(SandboxToolsBase): # ) # async def browser_search_google(self, query: str) -> ToolResult: # """Search Google with the provided query - + # Args: # query (str): The search query to use - + # Returns: # dict: Result of the execution # """ @@ -323,7 +323,7 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_go_back(self) -> ToolResult: """Navigate back in browser history - + Returns: dict: Result of the execution """ @@ -361,10 +361,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_wait(self, seconds: int = 3) -> ToolResult: """Wait for the specified number of seconds - + Args: seconds (int, optional): Number of seconds to wait. Defaults to 3. - + Returns: dict: Result of the execution """ @@ -403,10 +403,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_click_element(self, index: int) -> ToolResult: """Click on an element by index - + Args: index (int): The index of the element to click - + Returns: dict: Result of the execution """ @@ -451,11 +451,11 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_input_text(self, index: int, text: str) -> ToolResult: """Input text into an element - + Args: index (int): The index of the element to input text into text (str): The text to input - + Returns: dict: Result of the execution """ @@ -494,10 +494,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_send_keys(self, keys: str) -> ToolResult: """Send keyboard keys - + Args: keys (str): The keys to send (e.g., 'Enter', 'Escape', 'Control+a') - + Returns: dict: Result of the execution """ @@ -536,10 +536,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_switch_tab(self, page_id: int) -> ToolResult: """Switch to a different browser tab - + Args: page_id (int): The ID of the tab to switch to - + Returns: dict: Result of the execution """ @@ -576,10 +576,10 @@ class SandboxBrowserTool(SandboxToolsBase): # ) # async def browser_open_tab(self, url: str) -> ToolResult: # """Open a new browser tab with the specified URL - + # Args: # url (str): The URL to open in the new tab - + # Returns: # dict: Result of the execution # """ @@ -618,10 +618,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_close_tab(self, page_id: int) -> ToolResult: """Close a browser tab - + Args: page_id (int): The ID of the tab to close - + Returns: dict: Result of the execution """ @@ -658,23 +658,23 @@ class SandboxBrowserTool(SandboxToolsBase): # ) # async def browser_extract_content(self, goal: str) -> ToolResult: # """Extract content from the current page based on the provided goal - + # Args: # goal (str): The extraction goal - + # Returns: # dict: Result of the execution # """ # logger.debug(f"\033[95mExtracting content with goal: {goal}\033[0m") # result = await self._execute_browser_action("extract_content", {"goal": goal}) - + # # Format content for better readability # if result.get("success"): # logger.debug(f"\033[92mContent extraction successful\033[0m") # content = result.data.get("content", "") # url = result.data.get("url", "") # title = result.data.get("title", "") - + # if content: # content_preview = content[:200] + "..." if len(content) > 200 else content # logger.debug(f"\033[95mExtracted content from {title} ({url}):\033[0m") @@ -684,7 +684,7 @@ class SandboxBrowserTool(SandboxToolsBase): # logger.debug(f"\033[93mNo content extracted from {url}\033[0m") # else: # logger.debug(f"\033[91mFailed to extract content: {result.data.get('error', 'Unknown error')}\033[0m") - + # return result @openapi_schema({ @@ -718,10 +718,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_scroll_down(self, amount: int = None) -> ToolResult: """Scroll down the page - + Args: amount (int, optional): Pixel amount to scroll. If None, scrolls one page. - + Returns: dict: Result of the execution """ @@ -731,7 +731,7 @@ class SandboxBrowserTool(SandboxToolsBase): logger.debug(f"\033[95mScrolling down by {amount} pixels\033[0m") else: logger.debug(f"\033[95mScrolling down one page\033[0m") - + return await self._execute_browser_action("scroll_down", params) @openapi_schema({ @@ -765,10 +765,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_scroll_up(self, amount: int = None) -> ToolResult: """Scroll up the page - + Args: amount (int, optional): Pixel amount to scroll. If None, scrolls one page. - + Returns: dict: Result of the execution """ @@ -778,7 +778,7 @@ class SandboxBrowserTool(SandboxToolsBase): logger.debug(f"\033[95mScrolling up by {amount} pixels\033[0m") else: logger.debug(f"\033[95mScrolling up one page\033[0m") - + return await self._execute_browser_action("scroll_up", params) @openapi_schema({ @@ -813,10 +813,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_scroll_to_text(self, text: str) -> ToolResult: """Scroll to specific text on the page - + Args: text (str): The text to scroll to - + Returns: dict: Result of the execution """ @@ -855,10 +855,10 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_get_dropdown_options(self, index: int) -> ToolResult: """Get all options from a dropdown element - + Args: index (int): The index of the dropdown element - + Returns: dict: Result of the execution with the dropdown options """ @@ -903,11 +903,11 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_select_dropdown_option(self, index: int, text: str) -> ToolResult: """Select an option from a dropdown by text - + Args: index (int): The index of the dropdown element text (str): The text of the option to select - + Returns: dict: Result of the execution """ @@ -969,11 +969,11 @@ class SandboxBrowserTool(SandboxToolsBase): ''' ) - async def browser_drag_drop(self, element_source: str = None, element_target: str = None, + async def browser_drag_drop(self, element_source: str = None, element_target: str = None, coord_source_x: int = None, coord_source_y: int = None, coord_target_x: int = None, coord_target_y: int = None) -> ToolResult: """Perform drag and drop operation between elements or coordinates - + Args: element_source (str, optional): The source element selector element_target (str, optional): The target element selector @@ -981,12 +981,12 @@ class SandboxBrowserTool(SandboxToolsBase): coord_source_y (int, optional): The source Y coordinate coord_target_x (int, optional): The target X coordinate coord_target_y (int, optional): The target Y coordinate - + Returns: dict: Result of the execution """ params = {} - + if element_source and element_target: params["element_source"] = element_source params["element_target"] = element_target @@ -999,7 +999,7 @@ class SandboxBrowserTool(SandboxToolsBase): logger.debug(f"\033[95mDragging from coordinates ({coord_source_x}, {coord_source_y}) to ({coord_target_x}, {coord_target_y})\033[0m") else: return self.fail_response("Must provide either element selectors or coordinates for drag and drop") - + return await self._execute_browser_action("drag_drop", params) @openapi_schema({ @@ -1040,13 +1040,13 @@ class SandboxBrowserTool(SandboxToolsBase): ) async def browser_click_coordinates(self, x: int, y: int) -> ToolResult: """Click at specific X,Y coordinates on the page - + Args: x (int): The X coordinate to click y (int): The Y coordinate to click - + Returns: dict: Result of the execution """ logger.debug(f"\033[95mClicking at coordinates: ({x}, {y})\033[0m") - return await self._execute_browser_action("click_coordinates", {"x": x, "y": y}) \ No newline at end of file + return await self._execute_browser_action("click_coordinates", {"x": x, "y": y}) diff --git a/app/tool/sb_deploy_tool.py b/app/tool/sb_deploy_tool.py index 5c52f39..7be0929 100644 --- a/app/tool/sb_deploy_tool.py +++ b/app/tool/sb_deploy_tool.py @@ -1,7 +1,7 @@ import os from dotenv import load_dotenv from agentpress.tool import ToolResult, openapi_schema, xml_schema -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from utils.files_utils import clean_path from agentpress.thread_manager import ThreadManager @@ -48,11 +48,11 @@ class SandboxDeployTool(SandboxToolsBase): {"param_name": "directory_path", "node_type": "attribute", "path": "directory_path"} ], example=''' - @@ -68,11 +68,11 @@ class SandboxDeployTool(SandboxToolsBase): """ Deploy a static website (HTML+CSS+JS) from the sandbox to Cloudflare Pages. Only use this tool when permanent deployment to a production environment is needed. - + Args: name: Name for the deployment, will be used in the URL as {name}.kortix.cloud directory_path: Path to the directory to deploy, relative to /workspace - + Returns: ToolResult containing: - Success: Deployment information including URL @@ -81,10 +81,10 @@ class SandboxDeployTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + directory_path = self.clean_path(directory_path) full_path = f"{self.workspace_path}/{directory_path}" - + # Verify the directory exists try: dir_info = self.sandbox.fs.get_file_info(full_path) @@ -92,26 +92,26 @@ class SandboxDeployTool(SandboxToolsBase): return self.fail_response(f"'{directory_path}' is not a directory") except Exception as e: return self.fail_response(f"Directory '{directory_path}' does not exist: {str(e)}") - + # Deploy to Cloudflare Pages directly from the container try: # Get Cloudflare API token from environment if not self.cloudflare_api_token: return self.fail_response("CLOUDFLARE_API_TOKEN environment variable not set") - + # Single command that creates the project if it doesn't exist and then deploys project_name = f"{self.sandbox_id}-{name}" - deploy_cmd = f'''cd {self.workspace_path} && export CLOUDFLARE_API_TOKEN={self.cloudflare_api_token} && - (npx wrangler pages deploy {full_path} --project-name {project_name} || - (npx wrangler pages project create {project_name} --production-branch production && + deploy_cmd = f'''cd {self.workspace_path} && export CLOUDFLARE_API_TOKEN={self.cloudflare_api_token} && + (npx wrangler pages deploy {full_path} --project-name {project_name} || + (npx wrangler pages project create {project_name} --production-branch production && npx wrangler pages deploy {full_path} --project-name {project_name}))''' # Execute the command directly using the sandbox's process.exec method response = self.sandbox.process.exec(f"/bin/sh -c \"{deploy_cmd}\"", timeout=300) - + print(f"Deployment command output: {response.result}") - + if response.exit_code == 0: return self.success_response({ "message": f"Website deployed successfully", @@ -127,21 +127,21 @@ class SandboxDeployTool(SandboxToolsBase): if __name__ == "__main__": import asyncio import sys - + async def test_deploy(): # Replace these with actual values for testing sandbox_id = "sandbox-ccb30b35" password = "test-password" - + # Initialize the deploy tool deploy_tool = SandboxDeployTool(sandbox_id, password) - + # Test deployment - replace with actual directory path and site name result = await deploy_tool.deploy( name="test-site-1x", directory_path="website" # Directory containing static site files ) print(f"Deployment result: {result}") - + asyncio.run(test_deploy()) diff --git a/app/tool/sb_expose_tool.py b/app/tool/sb_expose_tool.py index ff833b3..dc64a6f 100644 --- a/app/tool/sb_expose_tool.py +++ b/app/tool/sb_expose_tool.py @@ -1,5 +1,5 @@ from agentpress.tool import ToolResult, openapi_schema, xml_schema -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from agentpress.thread_manager import ThreadManager import asyncio import time @@ -13,25 +13,25 @@ class SandboxExposeTool(SandboxToolsBase): async def _wait_for_sandbox_services(self, timeout: int = 30) -> bool: """Wait for sandbox services to be fully started before exposing ports.""" start_time = time.time() - + while time.time() - start_time < timeout: try: # Check if supervisord is running and managing services result = self.sandbox.process.exec("supervisorctl status", timeout=10) - + if result.exit_code == 0: # Check if key services are running status_output = result.output if "http_server" in status_output and "RUNNING" in status_output: return True - + # If services aren't ready, wait a bit await asyncio.sleep(2) - + except Exception as e: # If we can't check status, wait a bit and try again await asyncio.sleep(2) - + return False @openapi_schema({ @@ -85,10 +85,10 @@ class SandboxExposeTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # Convert port to integer if it's a string port = int(port) - + # Validate port number if not 1 <= port <= 65535: return self.fail_response(f"Invalid port number: {port}. Must be between 1 and 65535.") @@ -110,16 +110,16 @@ class SandboxExposeTool(SandboxToolsBase): # Get the preview link for the specified port preview_link = self.sandbox.get_preview_link(port) - + # Extract the actual URL from the preview link object url = preview_link.url if hasattr(preview_link, 'url') else str(preview_link) - + return self.success_response({ "url": url, "port": port, "message": f"Successfully exposed port {port} to the public. Users can now access this service at: {url}" }) - + except ValueError: return self.fail_response(f"Invalid port number: {port}. Must be a valid integer between 1 and 65535.") except Exception as e: diff --git a/app/tool/sb_files_tool.py b/app/tool/sb_files_tool.py index 34f3349..b16235b 100644 --- a/app/tool/sb_files_tool.py +++ b/app/tool/sb_files_tool.py @@ -1,5 +1,5 @@ from agentpress.tool import ToolResult, openapi_schema, xml_schema -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from utils.files_utils import should_exclude_file, clean_path from agentpress.thread_manager import ThreadManager from utils.logger import logger @@ -35,11 +35,11 @@ class SandboxFilesTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + files = self.sandbox.fs.list_files(self.workspace_path) for file_info in files: rel_path = file_info.name - + # Skip excluded files and directories if self._should_exclude_file(rel_path) or file_info.is_dir: continue @@ -59,7 +59,7 @@ class SandboxFilesTool(SandboxToolsBase): print(f"Skipping binary file: {rel_path}") return files_state - + except Exception as e: print(f"Error getting workspace state: {str(e)}") return {} @@ -111,7 +111,7 @@ class SandboxFilesTool(SandboxToolsBase): # This is the file content def main(): print("Hello, World!") - + if __name__ == "__main__": main() @@ -123,23 +123,23 @@ class SandboxFilesTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + file_path = self.clean_path(file_path) full_path = f"{self.workspace_path}/{file_path}" if self._file_exists(full_path): return self.fail_response(f"File '{file_path}' already exists. Use update_file to modify existing files.") - + # Create parent directories if needed parent_dir = '/'.join(full_path.split('/')[:-1]) if parent_dir: self.sandbox.fs.create_folder(parent_dir, "755") - + # Write the file content self.sandbox.fs.upload_file(file_contents.encode(), full_path) self.sandbox.fs.set_file_permissions(full_path, permissions) - + message = f"File '{file_path}' created successfully." - + # Check if index.html was created and add 8080 server info (only in root workspace) if file_path.lower() == 'index.html': try: @@ -149,7 +149,7 @@ class SandboxFilesTool(SandboxToolsBase): message += "\n[Note: Use the provided HTTP server URL above instead of starting a new server]" except Exception as e: logger.warning(f"Failed to get website URL for index.html: {str(e)}") - + return self.success_response(message) except Exception as e: return self.fail_response(f"Error creating file: {str(e)}") @@ -200,41 +200,41 @@ class SandboxFilesTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + file_path = self.clean_path(file_path) full_path = f"{self.workspace_path}/{file_path}" if not self._file_exists(full_path): return self.fail_response(f"File '{file_path}' does not exist") - + content = self.sandbox.fs.download_file(full_path).decode() old_str = old_str.expandtabs() new_str = new_str.expandtabs() - + occurrences = content.count(old_str) if occurrences == 0: return self.fail_response(f"String '{old_str}' not found in file") if occurrences > 1: lines = [i+1 for i, line in enumerate(content.split('\n')) if old_str in line] return self.fail_response(f"Multiple occurrences found in lines {lines}. Please ensure string is unique") - + # Perform replacement new_content = content.replace(old_str, new_str) self.sandbox.fs.upload_file(new_content.encode(), full_path) - + # Show snippet around the edit replacement_line = content.split(old_str)[0].count('\n') start_line = max(0, replacement_line - self.SNIPPET_LINES) end_line = replacement_line + self.SNIPPET_LINES + new_str.count('\n') snippet = '\n'.join(new_content.split('\n')[start_line:end_line + 1]) - + # Get preview URL if it's an HTML file # preview_url = self._get_preview_url(file_path) message = f"Replacement successful." # if preview_url: # message += f"\n\nYou can preview this HTML file at: {preview_url}" - + return self.success_response(message) - + except Exception as e: return self.fail_response(f"Error replacing string: {str(e)}") @@ -288,17 +288,17 @@ class SandboxFilesTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + file_path = self.clean_path(file_path) full_path = f"{self.workspace_path}/{file_path}" if not self._file_exists(full_path): return self.fail_response(f"File '{file_path}' does not exist. Use create_file to create a new file.") - + self.sandbox.fs.upload_file(file_contents.encode(), full_path) self.sandbox.fs.set_file_permissions(full_path, permissions) - + message = f"File '{file_path}' completely rewritten successfully." - + # Check if index.html was rewritten and add 8080 server info (only in root workspace) if file_path.lower() == 'index.html': try: @@ -308,7 +308,7 @@ class SandboxFilesTool(SandboxToolsBase): message += "\n[Note: Use the provided HTTP server URL above instead of starting a new server]" except Exception as e: logger.warning(f"Failed to get website URL for index.html: {str(e)}") - + return self.success_response(message) except Exception as e: return self.fail_response(f"Error rewriting file: {str(e)}") @@ -347,12 +347,12 @@ class SandboxFilesTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + file_path = self.clean_path(file_path) full_path = f"{self.workspace_path}/{file_path}" if not self._file_exists(full_path): return self.fail_response(f"File '{file_path}' does not exist") - + self.sandbox.fs.delete_file(full_path) return self.success_response(f"File '{file_path}' deleted successfully.") except Exception as e: @@ -412,12 +412,12 @@ class SandboxFilesTool(SandboxToolsBase): # ) # async def read_file(self, file_path: str, start_line: int = 1, end_line: Optional[int] = None) -> ToolResult: # """Read file content with optional line range specification. - + # Args: # file_path: Path to the file relative to /workspace # start_line: Starting line number (1-based), defaults to 1 # end_line: Ending line number (inclusive), defaults to None (end of file) - + # Returns: # ToolResult containing: # - Success: File content and metadata @@ -426,27 +426,27 @@ class SandboxFilesTool(SandboxToolsBase): # try: # file_path = self.clean_path(file_path) # full_path = f"{self.workspace_path}/{file_path}" - + # if not self._file_exists(full_path): # return self.fail_response(f"File '{file_path}' does not exist") - + # # Download and decode file content # content = self.sandbox.fs.download_file(full_path).decode() - + # # Split content into lines # lines = content.split('\n') # total_lines = len(lines) - + # # Handle line range if specified # if start_line > 1 or end_line is not None: # # Convert to 0-based indices # start_idx = max(0, start_line - 1) # end_idx = end_line if end_line is not None else total_lines # end_idx = min(end_idx, total_lines) # Ensure we don't exceed file length - + # # Extract the requested lines # content = '\n'.join(lines[start_idx:end_idx]) - + # return self.success_response({ # "content": content, # "file_path": file_path, @@ -454,7 +454,7 @@ class SandboxFilesTool(SandboxToolsBase): # "end_line": end_line if end_line is not None else total_lines, # "total_lines": total_lines # }) - + # except UnicodeDecodeError: # return self.fail_response(f"File '{file_path}' appears to be binary and cannot be read as text") # except Exception as e: diff --git a/app/tool/sb_shell_tool.py b/app/tool/sb_shell_tool.py index b165472..74acc6c 100644 --- a/app/tool/sb_shell_tool.py +++ b/app/tool/sb_shell_tool.py @@ -2,11 +2,11 @@ from typing import Optional, Dict, Any import time from uuid import uuid4 from agentpress.tool import ToolResult, openapi_schema, xml_schema -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from agentpress.thread_manager import ThreadManager class SandboxShellTool(SandboxToolsBase): - """Tool for executing tasks in a Daytona sandbox with browser-use capabilities. + """Tool for executing tasks in a Daytona sandbox with browser-use capabilities. Uses sessions for maintaining state between commands and provides comprehensive process management.""" def __init__(self, project_id: str, thread_manager: ThreadManager): @@ -108,8 +108,8 @@ class SandboxShellTool(SandboxToolsBase): ''' ) async def execute_command( - self, - command: str, + self, + command: str, folder: Optional[str] = None, session_name: Optional[str] = None, blocking: bool = False, @@ -118,61 +118,61 @@ class SandboxShellTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # Set up working directory cwd = self.workspace_path if folder: folder = folder.strip('/') cwd = f"{self.workspace_path}/{folder}" - + # Generate a session name if not provided if not session_name: session_name = f"session_{str(uuid4())[:8]}" - + # Check if tmux session already exists check_session = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'") session_exists = "not_exists" not in check_session.get("output", "") - + if not session_exists: # Create a new tmux session await self._execute_raw_command(f"tmux new-session -d -s {session_name}") - + # Ensure we're in the correct directory and send command to tmux full_command = f"cd {cwd} && {command}" wrapped_command = full_command.replace('"', '\\"') # Escape double quotes - + # Send command to tmux session await self._execute_raw_command(f'tmux send-keys -t {session_name} "{wrapped_command}" Enter') - + if blocking: # For blocking execution, wait and capture output start_time = time.time() while (time.time() - start_time) < timeout: # Wait a bit before checking time.sleep(2) - + # Check if session still exists (command might have exited) check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'ended'") if "ended" in check_result.get("output", ""): break - + # Get current output and check for common completion indicators output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -") current_output = output_result.get("output", "") - + # Check for prompt indicators that suggest command completion last_lines = current_output.split('\n')[-3:] completion_indicators = ['$', '#', '>', 'Done', 'Completed', 'Finished', '✓'] if any(indicator in line for indicator in completion_indicators for line in last_lines): break - + # Capture final output output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -") final_output = output_result.get("output", "") - + # Kill the session after capture await self._execute_raw_command(f"tmux kill-session -t {session_name}") - + return self.success_response({ "output": final_output, "session_name": session_name, @@ -187,7 +187,7 @@ class SandboxShellTool(SandboxToolsBase): "message": f"Command sent to tmux session '{session_name}'. Use check_command_output to view results.", "completed": False }) - + except Exception as e: # Attempt to clean up session in case of error if session_name: @@ -201,7 +201,7 @@ class SandboxShellTool(SandboxToolsBase): """Execute a raw command directly in the sandbox.""" # Ensure session exists for raw commands session_id = await self._ensure_session("raw_commands") - + # Execute command in session from sandbox.sandbox import SessionExecuteRequest req = SessionExecuteRequest( @@ -209,18 +209,18 @@ class SandboxShellTool(SandboxToolsBase): var_async=False, cwd=self.workspace_path ) - + response = self.sandbox.process.execute_session_command( session_id=session_id, req=req, timeout=30 # Short timeout for utility commands ) - + logs = self.sandbox.process.get_session_command_logs( session_id=session_id, command_id=response.cmd_id ) - + return { "output": logs, "exit_code": response.exit_code @@ -260,7 +260,7 @@ class SandboxShellTool(SandboxToolsBase): dev_server - + @@ -278,29 +278,29 @@ class SandboxShellTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # Check if session exists check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'") if "not_exists" in check_result.get("output", ""): return self.fail_response(f"Tmux session '{session_name}' does not exist.") - + # Get output from tmux pane output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -") output = output_result.get("output", "") - + # Kill session if requested if kill_session: await self._execute_raw_command(f"tmux kill-session -t {session_name}") termination_status = "Session terminated." else: termination_status = "Session still running." - + return self.success_response({ "output": output, "session_name": session_name, "status": termination_status }) - + except Exception as e: return self.fail_response(f"Error checking command output: {str(e)}") @@ -341,19 +341,19 @@ class SandboxShellTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # Check if session exists check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'") if "not_exists" in check_result.get("output", ""): return self.fail_response(f"Tmux session '{session_name}' does not exist.") - + # Kill the session await self._execute_raw_command(f"tmux kill-session -t {session_name}") - + return self.success_response({ "message": f"Tmux session '{session_name}' terminated successfully." }) - + except Exception as e: return self.fail_response(f"Error terminating command: {str(e)}") @@ -382,17 +382,17 @@ class SandboxShellTool(SandboxToolsBase): try: # Ensure sandbox is initialized await self._ensure_sandbox() - + # List all tmux sessions result = await self._execute_raw_command("tmux list-sessions 2>/dev/null || echo 'No sessions'") output = result.get("output", "") - + if "No sessions" in output or not output.strip(): return self.success_response({ "message": "No active tmux sessions found.", "sessions": [] }) - + # Parse session list sessions = [] for line in output.split('\n'): @@ -401,12 +401,12 @@ class SandboxShellTool(SandboxToolsBase): if parts: session_name = parts[0].strip() sessions.append(session_name) - + return self.success_response({ "message": f"Found {len(sessions)} active sessions.", "sessions": sessions }) - + except Exception as e: return self.fail_response(f"Error listing commands: {str(e)}") @@ -414,10 +414,10 @@ class SandboxShellTool(SandboxToolsBase): """Clean up all sessions.""" for session_name in list(self._sessions.keys()): await self._cleanup_session(session_name) - + # Also clean up any tmux sessions try: await self._ensure_sandbox() await self._execute_raw_command("tmux kill-server 2>/dev/null || true") except: - pass \ No newline at end of file + pass diff --git a/app/tool/sb_vision_tool.py b/app/tool/sb_vision_tool.py index 6459b2d..2cd356e 100644 --- a/app/tool/sb_vision_tool.py +++ b/app/tool/sb_vision_tool.py @@ -6,7 +6,7 @@ from io import BytesIO from PIL import Image from agentpress.tool import ToolResult, openapi_schema, xml_schema -from sandbox.tool_base import SandboxToolsBase +from daytona.tool_base import SandboxToolsBase from agentpress.thread_manager import ThreadManager import json @@ -38,19 +38,19 @@ class SandboxVisionTool(SandboxToolsBase): def compress_image(self, image_bytes: bytes, mime_type: str, file_path: str) -> Tuple[bytes, str]: """Compress an image to reduce its size while maintaining reasonable quality. - + Args: image_bytes: Original image bytes mime_type: MIME type of the image file_path: Path to the image file (for logging) - + Returns: Tuple of (compressed_bytes, new_mime_type) """ try: # Open image from bytes img = Image.open(BytesIO(image_bytes)) - + # Convert RGBA to RGB if necessary (for JPEG) if img.mode in ('RGBA', 'LA', 'P'): # Create a white background @@ -59,7 +59,7 @@ class SandboxVisionTool(SandboxToolsBase): img = img.convert('RGBA') background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None) img = background - + # Calculate new dimensions while maintaining aspect ratio width, height = img.size if width > DEFAULT_MAX_WIDTH or height > DEFAULT_MAX_HEIGHT: @@ -68,10 +68,10 @@ class SandboxVisionTool(SandboxToolsBase): new_height = int(height * ratio) img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) print(f"[SeeImage] Resized image from {width}x{height} to {new_width}x{new_height}") - + # Save to bytes with compression output = BytesIO() - + # Determine output format based on original mime type if mime_type == 'image/gif': # Keep GIFs as GIFs to preserve animation @@ -85,17 +85,17 @@ class SandboxVisionTool(SandboxToolsBase): # Convert everything else to JPEG for better compression img.save(output, format='JPEG', quality=DEFAULT_JPEG_QUALITY, optimize=True) output_mime = 'image/jpeg' - + compressed_bytes = output.getvalue() - + # Log compression results original_size = len(image_bytes) compressed_size = len(compressed_bytes) compression_ratio = (1 - compressed_size / original_size) * 100 print(f"[SeeImage] Compressed '{file_path}' from {original_size / 1024:.1f}KB to {compressed_size / 1024:.1f}KB ({compression_ratio:.1f}% reduction)") - + return compressed_bytes, output_mime - + except Exception as e: print(f"[SeeImage] Failed to compress image: {str(e)}. Using original.") return image_bytes, mime_type @@ -173,7 +173,7 @@ class SandboxVisionTool(SandboxToolsBase): # Compress the image compressed_bytes, compressed_mime_type = self.compress_image(image_bytes, mime_type, cleaned_path) - + # Check if compressed image is still too large if len(compressed_bytes) > MAX_COMPRESSED_SIZE: return self.fail_response(f"Image file '{cleaned_path}' is still too large after compression ({len(compressed_bytes) / (1024*1024):.2f}MB). Maximum compressed size is {MAX_COMPRESSED_SIZE / (1024*1024)}MB.") @@ -203,4 +203,4 @@ class SandboxVisionTool(SandboxToolsBase): return self.success_response(f"Successfully loaded and compressed the image '{cleaned_path}' (reduced from {file_info.size / 1024:.1f}KB to {len(compressed_bytes) / 1024:.1f}KB).") except Exception as e: - return self.fail_response(f"An unexpected error occurred while trying to see the image: {str(e)}") \ No newline at end of file + return self.fail_response(f"An unexpected error occurred while trying to see the image: {str(e)}") diff --git a/app/tool/update_agent_tool.py b/app/tool/update_agent_tool.py deleted file mode 100644 index 5113e56..0000000 --- a/app/tool/update_agent_tool.py +++ /dev/null @@ -1,889 +0,0 @@ -import json -import httpx -from typing import Optional, Dict, Any, List -from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema -from agentpress.thread_manager import ThreadManager - -class UpdateAgentTool(Tool): - """Tool for updating agent configuration. - - This tool is used by the agent builder to update agent properties - based on user requirements. - """ - - def __init__(self, thread_manager: ThreadManager, db_connection, agent_id: str): - super().__init__() - self.thread_manager = thread_manager - self.db = db_connection - self.agent_id = agent_id - # Smithery API configuration - self.smithery_api_base_url = "https://registry.smithery.ai" - import os - self.smithery_api_key = os.getenv("SMITHERY_API_KEY") - - @openapi_schema({ - "type": "function", - "function": { - "name": "update_agent", - "description": "Update the agent's configuration including name, description, system prompt, tools, and MCP servers. Call this whenever the user wants to modify any aspect of the agent.", - "parameters": { - "type": "object", - "properties": { - "name": { - "type": "string", - "description": "The name of the agent. Should be descriptive and indicate the agent's purpose." - }, - "description": { - "type": "string", - "description": "A brief description of what the agent does and its capabilities." - }, - "system_prompt": { - "type": "string", - "description": "The system instructions that define the agent's behavior, expertise, and approach. This should be comprehensive and well-structured." - }, - "agentpress_tools": { - "type": "object", - "description": "Configuration for AgentPress tools. Each key is a tool name, and the value is an object with 'enabled' (boolean) and 'description' (string) properties.", - "additionalProperties": { - "type": "object", - "properties": { - "enabled": {"type": "boolean"}, - "description": {"type": "string"} - } - } - }, - "configured_mcps": { - "type": "array", - "description": "List of configured MCP servers for external integrations.", - "items": { - "type": "object", - "properties": { - "name": {"type": "string"}, - "qualifiedName": {"type": "string"}, - "config": {"type": "object"}, - "enabledTools": { - "type": "array", - "items": {"type": "string"} - } - } - } - }, - "avatar": { - "type": "string", - "description": "Emoji to use as the agent's avatar." - }, - "avatar_color": { - "type": "string", - "description": "Hex color code for the agent's avatar background." - } - }, - "required": [] - } - } - }) - @xml_schema( - tag_name="update-agent", - mappings=[ - {"param_name": "name", "node_type": "attribute", "path": ".", "required": False}, - {"param_name": "description", "node_type": "element", "path": "description", "required": False}, - {"param_name": "system_prompt", "node_type": "element", "path": "system_prompt", "required": False}, - {"param_name": "agentpress_tools", "node_type": "element", "path": "agentpress_tools", "required": False}, - {"param_name": "configured_mcps", "node_type": "element", "path": "configured_mcps", "required": False}, - {"param_name": "avatar", "node_type": "attribute", "path": ".", "required": False}, - {"param_name": "avatar_color", "node_type": "attribute", "path": ".", "required": False} - ], - example=''' - - - Research Assistant - An AI assistant specialized in conducting research and providing comprehensive analysis - You are a research assistant with expertise in gathering, analyzing, and synthesizing information. Your approach is thorough and methodical... - {"web_search": {"enabled": true, "description": "Search the web for information"}, "sb_files": {"enabled": true, "description": "Read and write files"}} - 🔬 - #4F46E5 - - - ''' - ) - async def update_agent( - self, - name: Optional[str] = None, - description: Optional[str] = None, - system_prompt: Optional[str] = None, - agentpress_tools: Optional[Dict[str, Dict[str, Any]]] = None, - configured_mcps: Optional[list] = None, - avatar: Optional[str] = None, - avatar_color: Optional[str] = None - ) -> ToolResult: - """Update agent configuration with provided fields. - - Args: - name: Agent name - description: Agent description - system_prompt: System instructions for the agent - agentpress_tools: AgentPress tools configuration - configured_mcps: MCP servers configuration - avatar: Emoji avatar - avatar_color: Avatar background color - - Returns: - ToolResult with updated agent data or error - """ - try: - client = await self.db.client - - update_data = {} - if name is not None: - update_data["name"] = name - if description is not None: - update_data["description"] = description - if system_prompt is not None: - update_data["system_prompt"] = system_prompt - if agentpress_tools is not None: - formatted_tools = {} - for tool_name, tool_config in agentpress_tools.items(): - if isinstance(tool_config, dict): - formatted_tools[tool_name] = { - "enabled": tool_config.get("enabled", False), - "description": tool_config.get("description", "") - } - update_data["agentpress_tools"] = formatted_tools - if configured_mcps is not None: - if isinstance(configured_mcps, str): - configured_mcps = json.loads(configured_mcps) - update_data["configured_mcps"] = configured_mcps - if avatar is not None: - update_data["avatar"] = avatar - if avatar_color is not None: - update_data["avatar_color"] = avatar_color - - if not update_data: - return self.fail_response("No fields provided to update") - - result = await client.table('agents').update(update_data).eq('agent_id', self.agent_id).execute() - - if not result.data: - return self.fail_response("Failed to update agent") - - return self.success_response({ - "message": "Agent updated successfully", - "updated_fields": list(update_data.keys()), - "agent": result.data[0] - }) - - except Exception as e: - return self.fail_response(f"Error updating agent: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "get_current_agent_config", - "description": "Get the current configuration of the agent being edited. Use this to check what's already configured before making updates.", - "parameters": { - "type": "object", - "properties": {}, - "required": [] - } - } - }) - @xml_schema( - tag_name="get-current-agent-config", - mappings=[], - example=''' - - - - - ''' - ) - async def get_current_agent_config(self) -> ToolResult: - """Get the current agent configuration. - - Returns: - ToolResult with current agent configuration - """ - try: - client = await self.db.client - - result = await client.table('agents').select('*').eq('agent_id', self.agent_id).execute() - - if not result.data: - return self.fail_response("Agent not found") - - agent = result.data[0] - - config_summary = { - "agent_id": agent["agent_id"], - "name": agent.get("name", "Untitled Agent"), - "description": agent.get("description", "No description set"), - "system_prompt": agent.get("system_prompt", "No system prompt set"), - "avatar": agent.get("avatar", "🤖"), - "avatar_color": agent.get("avatar_color", "#6B7280"), - "agentpress_tools": agent.get("agentpress_tools", {}), - "configured_mcps": agent.get("configured_mcps", []), - "created_at": agent.get("created_at"), - "updated_at": agent.get("updated_at") - } - - tools_count = len([t for t, cfg in config_summary["agentpress_tools"].items() if cfg.get("enabled")]) - mcps_count = len(config_summary["configured_mcps"]) - - return self.success_response({ - "summary": f"Agent '{config_summary['name']}' has {tools_count} tools enabled and {mcps_count} MCP servers configured.", - "configuration": config_summary - }) - - except Exception as e: - return self.fail_response(f"Error getting agent configuration: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "search_mcp_servers", - "description": "Search for MCP servers from the Smithery registry based on user requirements. Use this when the user wants to add MCP tools to their agent.", - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query for finding relevant MCP servers (e.g., 'linear', 'github', 'database', 'search')" - }, - "category": { - "type": "string", - "description": "Optional category filter", - "enum": ["AI & Search", "Development & Version Control", "Project Management", "Communication & Collaboration", "Data & Analytics", "Cloud & Infrastructure", "File Storage", "Marketing & Sales", "Customer Support", "Finance", "Automation & Productivity", "Utilities"] - }, - "limit": { - "type": "integer", - "description": "Maximum number of servers to return (default: 10)", - "default": 10 - } - }, - "required": ["query"] - } - } - }) - @xml_schema( - tag_name="search-mcp-servers", - mappings=[ - {"param_name": "query", "node_type": "attribute", "path": "."}, - {"param_name": "category", "node_type": "attribute", "path": "."}, - {"param_name": "limit", "node_type": "attribute", "path": "."} - ], - example=''' - - - linear - 5 - - - ''' - ) - async def search_mcp_servers( - self, - query: str, - category: Optional[str] = None, - limit: int = 10 - ) -> ToolResult: - """Search for MCP servers based on user requirements. - - Args: - query: Search query for finding relevant MCP servers - category: Optional category filter - limit: Maximum number of servers to return - - Returns: - ToolResult with matching MCP servers - """ - try: - async with httpx.AsyncClient() as client: - headers = { - "Accept": "application/json", - "User-Agent": "Suna-MCP-Integration/1.0" - } - - if self.smithery_api_key: - headers["Authorization"] = f"Bearer {self.smithery_api_key}" - - params = { - "q": query, - "page": 1, - "pageSize": min(limit * 2, 50) # Get more results to filter - } - - response = await client.get( - f"{self.smithery_api_base_url}/servers", - headers=headers, - params=params, - timeout=30.0 - ) - - response.raise_for_status() - data = response.json() - servers = data.get("servers", []) - - # Filter by category if specified - if category: - filtered_servers = [] - for server in servers: - server_category = self._categorize_server(server) - if server_category == category: - filtered_servers.append(server) - servers = filtered_servers - - # Sort by useCount and limit results - servers = sorted(servers, key=lambda x: x.get("useCount", 0), reverse=True)[:limit] - - # Format results for user-friendly display - formatted_servers = [] - for server in servers: - formatted_servers.append({ - "name": server.get("displayName", server.get("qualifiedName", "Unknown")), - "qualifiedName": server.get("qualifiedName"), - "description": server.get("description", "No description available"), - "useCount": server.get("useCount", 0), - "category": self._categorize_server(server), - "homepage": server.get("homepage", ""), - "isDeployed": server.get("isDeployed", False) - }) - - if not formatted_servers: - return ToolResult( - success=False, - output=json.dumps([], ensure_ascii=False) - ) - - return ToolResult( - success=True, - output=json.dumps(formatted_servers, ensure_ascii=False) - ) - - except Exception as e: - return self.fail_response(f"Error searching MCP servers: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "get_mcp_server_tools", - "description": "Get detailed information about a specific MCP server including its available tools. Use this after the user selects a server they want to connect to.", - "parameters": { - "type": "object", - "properties": { - "qualified_name": { - "type": "string", - "description": "The qualified name of the MCP server (e.g., 'exa', '@smithery-ai/github')" - } - }, - "required": ["qualified_name"] - } - } - }) - @xml_schema( - tag_name="get-mcp-server-tools", - mappings=[ - {"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True} - ], - example=''' - - - exa - - - ''' - ) - async def get_mcp_server_tools(self, qualified_name: str) -> ToolResult: - """Get detailed information about a specific MCP server and its tools. - - Args: - qualified_name: The qualified name of the MCP server - - Returns: - ToolResult with server details and available tools - """ - try: - # First get server metadata from registry - async with httpx.AsyncClient() as client: - headers = { - "Accept": "application/json", - "User-Agent": "Suna-MCP-Integration/1.0" - } - - if self.smithery_api_key: - headers["Authorization"] = f"Bearer {self.smithery_api_key}" - - # URL encode the qualified name if it contains special characters - from urllib.parse import quote - if '@' in qualified_name or '/' in qualified_name: - encoded_name = quote(qualified_name, safe='') - else: - encoded_name = qualified_name - - url = f"{self.smithery_api_base_url}/servers/{encoded_name}" - - response = await client.get( - url, - headers=headers, - timeout=30.0 - ) - - response.raise_for_status() - server_data = response.json() - - # Now connect to the MCP server to get actual tools using ClientSession - try: - # Import MCP components - from mcp import ClientSession - from mcp.client.streamable_http import streamablehttp_client - import base64 - import os - - # Check if Smithery API key is available - smithery_api_key = os.getenv("SMITHERY_API_KEY") - if not smithery_api_key: - raise ValueError("SMITHERY_API_KEY environment variable is not set") - - # Create server URL with empty config for testing - config_json = json.dumps({}) - config_b64 = base64.b64encode(config_json.encode()).decode() - server_url = f"https://server.smithery.ai/{qualified_name}/mcp?config={config_b64}&api_key={smithery_api_key}" - - # Connect and get tools - async with streamablehttp_client(server_url) as (read_stream, write_stream, _): - async with ClientSession(read_stream, write_stream) as session: - # Initialize the connection - await session.initialize() - - # List available tools - tools_result = await session.list_tools() - tools = tools_result.tools if hasattr(tools_result, 'tools') else tools_result - - # Format tools for user-friendly display - formatted_tools = [] - for tool in tools: - tool_info = { - "name": tool.name, - "description": getattr(tool, 'description', 'No description available'), - } - - # Extract parameters from inputSchema if available - if hasattr(tool, 'inputSchema') and tool.inputSchema: - schema = tool.inputSchema - if isinstance(schema, dict): - tool_info["parameters"] = schema.get("properties", {}) - tool_info["required_params"] = schema.get("required", []) - else: - tool_info["parameters"] = {} - tool_info["required_params"] = [] - else: - tool_info["parameters"] = {} - tool_info["required_params"] = [] - - formatted_tools.append(tool_info) - - # Extract configuration requirements from server metadata - config_requirements = [] - security = server_data.get("security", {}) - if security: - for key, value in security.items(): - if isinstance(value, dict): - config_requirements.append({ - "name": key, - "description": value.get("description", f"Configuration for {key}"), - "required": value.get("required", False), - "type": value.get("type", "string") - }) - - server_info = { - "name": server_data.get("displayName", qualified_name), - "qualifiedName": qualified_name, - "description": server_data.get("description", "No description available"), - "homepage": server_data.get("homepage", ""), - "iconUrl": server_data.get("iconUrl", ""), - "isDeployed": server_data.get("isDeployed", False), - "tools": formatted_tools, - "config_requirements": config_requirements, - "total_tools": len(formatted_tools) - } - - return self.success_response({ - "message": f"Found {len(formatted_tools)} tools for {server_info['name']}", - "server": server_info - }) - - except Exception as mcp_error: - # If MCP connection fails, fall back to registry data - tools = server_data.get("tools", []) - formatted_tools = [] - for tool in tools: - formatted_tools.append({ - "name": tool.get("name", "Unknown"), - "description": tool.get("description", "No description available"), - "parameters": tool.get("inputSchema", {}).get("properties", {}), - "required_params": tool.get("inputSchema", {}).get("required", []) - }) - - config_requirements = [] - security = server_data.get("security", {}) - if security: - for key, value in security.items(): - if isinstance(value, dict): - config_requirements.append({ - "name": key, - "description": value.get("description", f"Configuration for {key}"), - "required": value.get("required", False), - "type": value.get("type", "string") - }) - - server_info = { - "name": server_data.get("displayName", qualified_name), - "qualifiedName": qualified_name, - "description": server_data.get("description", "No description available"), - "homepage": server_data.get("homepage", ""), - "iconUrl": server_data.get("iconUrl", ""), - "isDeployed": server_data.get("isDeployed", False), - "tools": formatted_tools, - "config_requirements": config_requirements, - "total_tools": len(formatted_tools), - "note": "Tools listed from registry metadata (MCP connection failed - may need configuration)" - } - - return self.success_response({ - "message": f"Found {len(formatted_tools)} tools for {server_info['name']} (from registry)", - "server": server_info - }) - - except Exception as e: - return self.fail_response(f"Error getting MCP server tools: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "configure_mcp_server", - "description": "Configure and add an MCP server to the agent with selected tools. Use this after the user has chosen which tools they want from a server.", - "parameters": { - "type": "object", - "properties": { - "qualified_name": { - "type": "string", - "description": "The qualified name of the MCP server" - }, - "display_name": { - "type": "string", - "description": "Display name for the server" - }, - "enabled_tools": { - "type": "array", - "description": "List of tool names to enable for this server", - "items": {"type": "string"} - }, - "config": { - "type": "object", - "description": "Configuration object with API keys and other settings", - "additionalProperties": True - } - }, - "required": ["qualified_name", "display_name", "enabled_tools"] - } - } - }) - @xml_schema( - tag_name="configure-mcp-server", - mappings=[ - {"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True}, - {"param_name": "display_name", "node_type": "attribute", "path": ".", "required": True}, - {"param_name": "enabled_tools", "node_type": "element", "path": "enabled_tools", "required": True}, - {"param_name": "config", "node_type": "element", "path": "config", "required": False} - ], - example=''' - - - exa - Exa Search - ["search", "find_similar"] - {"exaApiKey": "user-api-key"} - - - ''' - ) - async def configure_mcp_server( - self, - qualified_name: str, - display_name: str, - enabled_tools: List[str], - config: Optional[Dict[str, Any]] = None - ) -> ToolResult: - """Configure and add an MCP server to the agent. - - Args: - qualified_name: The qualified name of the MCP server - display_name: Display name for the server - enabled_tools: List of tool names to enable - config: Configuration object with API keys and settings - - Returns: - ToolResult with configuration status - """ - try: - client = await self.db.client - - # Get current agent configuration - result = await client.table('agents').select('configured_mcps').eq('agent_id', self.agent_id).execute() - - if not result.data: - return self.fail_response("Agent not found") - - current_mcps = result.data[0].get('configured_mcps', []) - - # Check if server is already configured - existing_server_index = None - for i, mcp in enumerate(current_mcps): - if mcp.get('qualifiedName') == qualified_name: - existing_server_index = i - break - - # Create new MCP configuration - new_mcp_config = { - "name": display_name, - "qualifiedName": qualified_name, - "config": config or {}, - "enabledTools": enabled_tools - } - - # Update or add the configuration - if existing_server_index is not None: - current_mcps[existing_server_index] = new_mcp_config - action = "updated" - else: - current_mcps.append(new_mcp_config) - action = "added" - - # Save to database - update_result = await client.table('agents').update({ - 'configured_mcps': current_mcps - }).eq('agent_id', self.agent_id).execute() - - if not update_result.data: - return self.fail_response("Failed to save MCP configuration") - - return self.success_response({ - "message": f"Successfully {action} MCP server '{display_name}' with {len(enabled_tools)} tools", - "server": new_mcp_config, - "total_mcp_servers": len(current_mcps), - "action": action - }) - - except Exception as e: - return self.fail_response(f"Error configuring MCP server: {str(e)}") - - @openapi_schema({ - "type": "function", - "function": { - "name": "get_popular_mcp_servers", - "description": "Get a list of popular and recommended MCP servers organized by category. Use this to show users popular options when they want to add MCP tools.", - "parameters": { - "type": "object", - "properties": { - "category": { - "type": "string", - "description": "Optional category filter to show only servers from a specific category", - "enum": ["AI & Search", "Development & Version Control", "Project Management", "Communication & Collaboration", "Data & Analytics", "Cloud & Infrastructure", "File Storage", "Marketing & Sales", "Customer Support", "Finance", "Automation & Productivity", "Utilities"] - } - }, - "required": [] - } - } - }) - @xml_schema( - tag_name="get-popular-mcp-servers", - mappings=[ - {"param_name": "category", "node_type": "attribute", "path": ".", "required": False} - ], - example=''' - - - AI & Search - - - ''' - ) - async def get_popular_mcp_servers(self, category: Optional[str] = None) -> ToolResult: - """Get popular MCP servers organized by category. - - Args: - category: Optional category filter - - Returns: - ToolResult with popular MCP servers - """ - try: - async with httpx.AsyncClient() as client: - headers = { - "Accept": "application/json", - "User-Agent": "Suna-MCP-Integration/1.0" - } - - if self.smithery_api_key: - headers["Authorization"] = f"Bearer {self.smithery_api_key}" - - response = await client.get( - f"{self.smithery_api_base_url}/servers", - headers=headers, - params={"page": 1, "pageSize": 50}, - timeout=30.0 - ) - - response.raise_for_status() - data = response.json() - servers = data.get("servers", []) - - # Categorize servers - categorized = {} - for server in servers: - server_category = self._categorize_server(server) - if category and server_category != category: - continue - - if server_category not in categorized: - categorized[server_category] = [] - - categorized[server_category].append({ - "name": server.get("displayName", server.get("qualifiedName", "Unknown")), - "qualifiedName": server.get("qualifiedName"), - "description": server.get("description", "No description available"), - "useCount": server.get("useCount", 0), - "homepage": server.get("homepage", ""), - "isDeployed": server.get("isDeployed", False) - }) - - # Sort categories and servers within each category - for cat in categorized: - categorized[cat] = sorted(categorized[cat], key=lambda x: x["useCount"], reverse=True)[:5] - - return self.success_response({ - "message": f"Found popular MCP servers" + (f" in category '{category}'" if category else ""), - "categorized_servers": categorized, - "total_categories": len(categorized) - }) - - except Exception as e: - return self.fail_response(f"Error getting popular MCP servers: {str(e)}") - - def _categorize_server(self, server: Dict[str, Any]) -> str: - """Categorize a server based on its qualified name and description.""" - qualified_name = server.get("qualifiedName", "").lower() - description = server.get("description", "").lower() - - # Category mappings - category_mappings = { - "AI & Search": ["exa", "perplexity", "openai", "anthropic", "duckduckgo", "brave", "google", "search"], - "Development & Version Control": ["github", "gitlab", "bitbucket", "git"], - "Project Management": ["linear", "jira", "asana", "notion", "trello", "monday", "clickup"], - "Communication & Collaboration": ["slack", "discord", "teams", "zoom", "telegram"], - "Data & Analytics": ["postgres", "mysql", "mongodb", "bigquery", "snowflake", "sqlite", "redis", "database"], - "Cloud & Infrastructure": ["aws", "gcp", "azure", "vercel", "netlify", "cloudflare", "docker"], - "File Storage": ["gdrive", "google-drive", "dropbox", "box", "onedrive", "s3", "drive"], - "Marketing & Sales": ["hubspot", "salesforce", "mailchimp", "sendgrid"], - "Customer Support": ["zendesk", "intercom", "freshdesk", "helpscout"], - "Finance": ["stripe", "quickbooks", "xero", "plaid"], - "Automation & Productivity": ["playwright", "puppeteer", "selenium", "desktop-commander", "sequential-thinking", "automation"], - "Utilities": ["filesystem", "memory", "fetch", "time", "weather", "currency", "file"] - } - - # Check qualified name and description for category keywords - for category, keywords in category_mappings.items(): - for keyword in keywords: - if keyword in qualified_name or keyword in description: - return category - - return "Other" - - @openapi_schema({ - "type": "function", - "function": { - "name": "test_mcp_server_connection", - "description": "Test connectivity to an MCP server with provided configuration. Use this to validate that a server can be connected to before adding it to the agent.", - "parameters": { - "type": "object", - "properties": { - "qualified_name": { - "type": "string", - "description": "The qualified name of the MCP server" - }, - "config": { - "type": "object", - "description": "Configuration object with API keys and other settings", - "additionalProperties": True - } - }, - "required": ["qualified_name"] - } - } - }) - @xml_schema( - tag_name="test-mcp-server-connection", - mappings=[ - {"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True}, - {"param_name": "config", "node_type": "element", "path": "config", "required": False} - ], - example=''' - - - exa - {"exaApiKey": "user-api-key"} - - - ''' - ) - async def test_mcp_server_connection( - self, - qualified_name: str, - config: Optional[Dict[str, Any]] = None - ) -> ToolResult: - """Test connectivity to an MCP server with provided configuration. - - Args: - qualified_name: The qualified name of the MCP server - config: Configuration object with API keys and settings - - Returns: - ToolResult with connection test results - """ - try: - # Import MCP components - from mcp import ClientSession - from mcp.client.streamable_http import streamablehttp_client - import base64 - import os - - # Check if Smithery API key is available - smithery_api_key = os.getenv("SMITHERY_API_KEY") - if not smithery_api_key: - return self.fail_response("SMITHERY_API_KEY environment variable is not set") - - # Create server URL with provided config - config_json = json.dumps(config or {}) - config_b64 = base64.b64encode(config_json.encode()).decode() - server_url = f"https://server.smithery.ai/{qualified_name}/mcp?config={config_b64}&api_key={smithery_api_key}" - - # Test connection - async with streamablehttp_client(server_url) as (read_stream, write_stream, _): - async with ClientSession(read_stream, write_stream) as session: - # Initialize the connection - await session.initialize() - - # List available tools to verify connection - tools_result = await session.list_tools() - tools = tools_result.tools if hasattr(tools_result, 'tools') else tools_result - - tool_names = [tool.name for tool in tools] - - return self.success_response({ - "message": f"Successfully connected to {qualified_name}", - "qualified_name": qualified_name, - "connection_status": "success", - "available_tools": tool_names, - "total_tools": len(tool_names) - }) - - except Exception as e: - return self.fail_response(f"Failed to connect to {qualified_name}: {str(e)}") \ No newline at end of file diff --git a/app/utils/files_utils.py b/app/utils/files_utils.py index 508b1bb..fb3f567 100644 --- a/app/utils/files_utils.py +++ b/app/utils/files_utils.py @@ -40,10 +40,10 @@ EXCLUDED_EXT = { def should_exclude_file(rel_path: str) -> bool: """Check if a file should be excluded based on path, name, or extension - + Args: rel_path: Relative path of the file to check - + Returns: True if the file should be excluded, False otherwise """ @@ -62,30 +62,30 @@ def should_exclude_file(rel_path: str) -> bool: if ext.lower() in EXCLUDED_EXT: return True - return False + return False def clean_path(path: str, workspace_path: str = "/workspace") -> str: """Clean and normalize a path to be relative to the workspace - + Args: path: The path to clean workspace_path: The base workspace path to remove (default: "/workspace") - + Returns: The cleaned path, relative to the workspace """ # Remove any leading slash path = path.lstrip('/') - + # Remove workspace prefix if present if path.startswith(workspace_path.lstrip('/')): path = path[len(workspace_path.lstrip('/')):] - + # Remove workspace/ prefix if present if path.startswith('workspace/'): path = path[9:] - + # Remove any remaining leading slash path = path.lstrip('/') - - return path \ No newline at end of file + + return path