diff --git a/app/agent/sandbox_agent.py b/app/agent/sandbox_agent.py index 7480f87..e1716b2 100644 --- a/app/agent/sandbox_agent.py +++ b/app/agent/sandbox_agent.py @@ -229,9 +229,3 @@ class SandboxManus(ToolCallAgent): self.next_step_prompt = original_prompt return result - return result - return result - return result - return result - return result - return result diff --git a/app/agentpress/__init__.py b/app/agentpress/__init__.py deleted file mode 100644 index 8ab9008..0000000 --- a/app/agentpress/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Utility functions and constants for agent tools \ No newline at end of file diff --git a/app/agentpress/response_processor.py b/app/agentpress/response_processor.py deleted file mode 100644 index e54f973..0000000 --- a/app/agentpress/response_processor.py +++ /dev/null @@ -1,1890 +0,0 @@ -""" -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 app.utils.logger import logger -from app.agentpress.tool import ToolResult -from app.agentpress.tool_registry import ToolRegistry -from app.agentpress.xml_tool_parser import XMLToolParser -from langfuse.client import StatefulTraceClient -from app.services.langfuse import langfuse -from app.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/tool.py b/app/agentpress/tool.py deleted file mode 100644 index 8bcb902..0000000 --- a/app/agentpress/tool.py +++ /dev/null @@ -1,240 +0,0 @@ -""" -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 app.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 deleted file mode 100644 index 481410b..0000000 --- a/app/agentpress/tool_registry.py +++ /dev/null @@ -1,152 +0,0 @@ -from typing import Dict, Type, Any, List, Optional, Callable -from app.agentpress.tool import Tool, SchemaType -from app.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 deleted file mode 100644 index e0dabab..0000000 --- a/app/agentpress/utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# 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 deleted file mode 100644 index 3bda909..0000000 --- a/app/agentpress/utils/json_helpers.py +++ /dev/null @@ -1,174 +0,0 @@ -""" -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 deleted file mode 100644 index e882a66..0000000 --- a/app/agentpress/xml_tool_parser.py +++ /dev/null @@ -1,300 +0,0 @@ -""" -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/api.py b/app/daytona/api.py index ccc5926..8f90bb5 100644 --- a/app/daytona/api.py +++ b/app/daytona/api.py @@ -2,7 +2,16 @@ import os import urllib.parse from typing import Optional -from fastapi import FastAPI, UploadFile, File, HTTPException, APIRouter, Form, Depends, Request +from fastapi import ( + FastAPI, + UploadFile, + File, + HTTPException, + APIRouter, + Form, + Depends, + Request, +) from fastapi.responses import Response from pydantic import BaseModel @@ -15,14 +24,17 @@ from services.supabase import DBConnection router = APIRouter(tags=["sandbox"]) db = None + def initialize(_db: DBConnection): """Initialize the sandbox API with resources from the main API.""" global db db = _db logger.info("Initialized sandbox API with database connection") + class FileInfo(BaseModel): """Model for file information""" + name: str path: str is_dir: bool @@ -30,243 +42,285 @@ class FileInfo(BaseModel): mod_time: str permissions: Optional[str] = None + def normalize_path(path: str) -> str: """ Normalize a path to ensure proper UTF-8 encoding and handling. - + Args: path: The file path, potentially containing URL-encoded characters - + Returns: Normalized path with proper UTF-8 encoding """ try: # First, ensure the path is properly URL-decoded decoded_path = urllib.parse.unquote(path) - + # Handle Unicode escape sequences like \u0308 try: # Replace Python-style Unicode escapes (\u0308) with actual characters # This handles cases where the Unicode escape sequence is part of the URL import re - unicode_pattern = re.compile(r'\\u([0-9a-fA-F]{4})') - + + unicode_pattern = re.compile(r"\\u([0-9a-fA-F]{4})") + def replace_unicode(match): hex_val = match.group(1) return chr(int(hex_val, 16)) - + decoded_path = unicode_pattern.sub(replace_unicode, decoded_path) except Exception as unicode_err: - logger.warning(f"Error processing Unicode escapes in path '{path}': {str(unicode_err)}") - + logger.warning( + f"Error processing Unicode escapes in path '{path}': {str(unicode_err)}" + ) + logger.debug(f"Normalized path from '{path}' to '{decoded_path}'") return decoded_path except Exception as e: logger.error(f"Error normalizing path '{path}': {str(e)}") return path # Return original path if decoding fails + async def verify_sandbox_access(client, sandbox_id: str, user_id: Optional[str] = None): """ Verify that a user has access to a specific sandbox based on account membership. - + Args: client: The Supabase client sandbox_id: The sandbox ID to check access for user_id: The user ID to check permissions for. Can be None for public resource access. - + Returns: dict: Project data containing sandbox information - + Raises: HTTPException: If the user doesn't have access to the sandbox or sandbox doesn't exist """ # Find the project that owns this sandbox - project_result = await client.table('projects').select('*').filter('sandbox->>id', 'eq', sandbox_id).execute() - + project_result = ( + await client.table("projects") + .select("*") + .filter("sandbox->>id", "eq", sandbox_id) + .execute() + ) + if not project_result.data or len(project_result.data) == 0: raise HTTPException(status_code=404, detail="Sandbox not found") - + project_data = project_result.data[0] - if project_data.get('is_public'): + if project_data.get("is_public"): return project_data - + # For private projects, we must have a user_id if not user_id: - raise HTTPException(status_code=401, detail="Authentication required for this resource") - - account_id = project_data.get('account_id') - + raise HTTPException( + status_code=401, detail="Authentication required for this resource" + ) + + account_id = project_data.get("account_id") + # Verify account membership if account_id: - account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute() + account_user_result = ( + await client.schema("basejump") + .from_("account_user") + .select("account_role") + .eq("user_id", user_id) + .eq("account_id", account_id) + .execute() + ) if account_user_result.data and len(account_user_result.data) > 0: return project_data - + raise HTTPException(status_code=403, detail="Not authorized to access this sandbox") + async def get_sandbox_by_id_safely(client, sandbox_id: str): """ Safely retrieve a sandbox object by its ID, using the project that owns it. - + Args: client: The Supabase client sandbox_id: The sandbox ID to retrieve - + Returns: Sandbox: The sandbox object - + Raises: HTTPException: If the sandbox doesn't exist or can't be retrieved """ # Find the project that owns this sandbox - project_result = await client.table('projects').select('project_id').filter('sandbox->>id', 'eq', sandbox_id).execute() - + project_result = ( + await client.table("projects") + .select("project_id") + .filter("sandbox->>id", "eq", sandbox_id) + .execute() + ) + if not project_result.data or len(project_result.data) == 0: logger.error(f"No project found for sandbox ID: {sandbox_id}") - raise HTTPException(status_code=404, detail="Sandbox not found - no project owns this sandbox ID") - + raise HTTPException( + status_code=404, + detail="Sandbox not found - no project owns this sandbox ID", + ) + # project_id = project_result.data[0]['project_id'] # logger.debug(f"Found project {project_id} for sandbox {sandbox_id}") - + try: # Get the sandbox sandbox = await get_or_start_sandbox(sandbox_id) # Extract just the sandbox object from the tuple (sandbox, sandbox_id, sandbox_pass) # sandbox = sandbox_tuple[0] - + return sandbox except Exception as e: logger.error(f"Error retrieving sandbox {sandbox_id}: {str(e)}") - raise HTTPException(status_code=500, detail=f"Failed to retrieve sandbox: {str(e)}") + raise HTTPException( + status_code=500, detail=f"Failed to retrieve sandbox: {str(e)}" + ) + @router.post("/sandboxes/{sandbox_id}/files") async def create_file( - sandbox_id: str, + sandbox_id: str, path: str = Form(...), file: UploadFile = File(...), request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """Create a file in the sandbox using direct file upload""" # Normalize the path to handle UTF-8 encoding correctly path = normalize_path(path) - - logger.info(f"Received file upload request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") + + logger.info( + f"Received file upload request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}" + ) client = await db.client - + # Verify the user has access to this sandbox await verify_sandbox_access(client, sandbox_id, user_id) - + try: # Get sandbox using the safer method sandbox = await get_sandbox_by_id_safely(client, sandbox_id) - + # Read file content directly from the uploaded file content = await file.read() - + # Create file using raw binary content sandbox.fs.upload_file(content, path) logger.info(f"File created at {path} in sandbox {sandbox_id}") - + return {"status": "success", "created": True, "path": path} except Exception as e: logger.error(f"Error creating file in sandbox {sandbox_id}: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) + @router.get("/sandboxes/{sandbox_id}/files") async def list_files( - sandbox_id: str, + sandbox_id: str, path: str, request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """List files and directories at the specified path""" # Normalize the path to handle UTF-8 encoding correctly path = normalize_path(path) - - logger.info(f"Received list files request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") + + logger.info( + f"Received list files request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}" + ) client = await db.client - + # Verify the user has access to this sandbox await verify_sandbox_access(client, sandbox_id, user_id) - + try: # Get sandbox using the safer method sandbox = await get_sandbox_by_id_safely(client, sandbox_id) - + # List files files = sandbox.fs.list_files(path) result = [] - + for file in files: # Convert file information to our model # Ensure forward slashes are used for paths, regardless of OS - full_path = f"{path.rstrip('/')}/{file.name}" if path != '/' else f"/{file.name}" + full_path = ( + f"{path.rstrip('/')}/{file.name}" if path != "/" else f"/{file.name}" + ) file_info = FileInfo( name=file.name, - path=full_path, # Use the constructed path + path=full_path, # Use the constructed path is_dir=file.is_dir, size=file.size, mod_time=str(file.mod_time), - permissions=getattr(file, 'permissions', None) + permissions=getattr(file, "permissions", None), ) result.append(file_info) - + logger.info(f"Successfully listed {len(result)} files in sandbox {sandbox_id}") return {"files": [file.dict() for file in result]} except Exception as e: logger.error(f"Error listing files in sandbox {sandbox_id}: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) + @router.get("/sandboxes/{sandbox_id}/files/content") async def read_file( - sandbox_id: str, + sandbox_id: str, path: str, request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """Read a file from the sandbox""" # Normalize the path to handle UTF-8 encoding correctly original_path = path path = normalize_path(path) - - logger.info(f"Received file read request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") + + logger.info( + f"Received file read request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}" + ) if original_path != path: logger.info(f"Normalized path from '{original_path}' to '{path}'") - + client = await db.client - + # Verify the user has access to this sandbox await verify_sandbox_access(client, sandbox_id, user_id) - + try: # Get sandbox using the safer method sandbox = await get_sandbox_by_id_safely(client, sandbox_id) - + # Read file directly - don't check existence first with a separate call try: content = sandbox.fs.download_file(path) except Exception as download_err: - logger.error(f"Error downloading file {path} from sandbox {sandbox_id}: {str(download_err)}") - raise HTTPException( - status_code=404, - detail=f"Failed to download file: {str(download_err)}" + logger.error( + f"Error downloading file {path} from sandbox {sandbox_id}: {str(download_err)}" ) - + raise HTTPException( + status_code=404, detail=f"Failed to download file: {str(download_err)}" + ) + # Return a Response object with the content directly filename = os.path.basename(path) logger.info(f"Successfully read file {filename} from sandbox {sandbox_id}") - + # Ensure proper encoding by explicitly using UTF-8 for the filename in Content-Disposition header # This applies RFC 5987 encoding for the filename to support non-ASCII characters - encoded_filename = filename.encode('utf-8').decode('latin-1') + encoded_filename = filename.encode("utf-8").decode("latin-1") content_disposition = f"attachment; filename*=UTF-8''{encoded_filename}" - + return Response( content=content, media_type="application/octet-stream", - headers={"Content-Disposition": content_disposition} + headers={"Content-Disposition": content_disposition}, ) except HTTPException: # Re-raise HTTP exceptions without wrapping @@ -275,116 +329,149 @@ async def read_file( logger.error(f"Error reading file in sandbox {sandbox_id}: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) + @router.delete("/sandboxes/{sandbox_id}/files") async def delete_file( - sandbox_id: str, + sandbox_id: str, path: str, request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """Delete a file from the sandbox""" # Normalize the path to handle UTF-8 encoding correctly path = normalize_path(path) - - logger.info(f"Received file delete request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}") + + logger.info( + f"Received file delete request for sandbox {sandbox_id}, path: {path}, user_id: {user_id}" + ) client = await db.client - + # Verify the user has access to this sandbox await verify_sandbox_access(client, sandbox_id, user_id) - + try: # Get sandbox using the safer method sandbox = await get_sandbox_by_id_safely(client, sandbox_id) - + # Delete file sandbox.fs.delete_file(path) logger.info(f"File deleted at {path} in sandbox {sandbox_id}") - + return {"status": "success", "deleted": True, "path": path} except Exception as e: logger.error(f"Error deleting file in sandbox {sandbox_id}: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) + @router.delete("/sandboxes/{sandbox_id}") async def delete_sandbox_route( sandbox_id: str, request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """Delete an entire sandbox""" - logger.info(f"Received sandbox delete request for sandbox {sandbox_id}, user_id: {user_id}") + logger.info( + f"Received sandbox delete request for sandbox {sandbox_id}, user_id: {user_id}" + ) client = await db.client - + # Verify the user has access to this sandbox await verify_sandbox_access(client, sandbox_id, user_id) - + try: # Delete the sandbox using the sandbox module function await delete_sandbox(sandbox_id) - + return {"status": "success", "deleted": True, "sandbox_id": sandbox_id} except Exception as e: logger.error(f"Error deleting sandbox {sandbox_id}: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) + # Should happen on server-side fully @router.post("/project/{project_id}/sandbox/ensure-active") async def ensure_project_sandbox_active( project_id: str, request: Request = None, - user_id: Optional[str] = Depends(get_optional_user_id) + user_id: Optional[str] = Depends(get_optional_user_id), ): """ Ensure that a project's sandbox is active and running. Checks the sandbox status and starts it if it's not running. """ - logger.info(f"Received ensure sandbox active request for project {project_id}, user_id: {user_id}") + logger.info( + f"Received ensure sandbox active request for project {project_id}, user_id: {user_id}" + ) client = await db.client - + # Find the project and sandbox information - project_result = await client.table('projects').select('*').eq('project_id', project_id).execute() - + project_result = ( + await client.table("projects") + .select("*") + .eq("project_id", project_id) + .execute() + ) + if not project_result.data or len(project_result.data) == 0: logger.error(f"Project not found: {project_id}") raise HTTPException(status_code=404, detail="Project not found") - + project_data = project_result.data[0] - + # For public projects, no authentication is needed - if not project_data.get('is_public'): + if not project_data.get("is_public"): # For private projects, we must have a user_id if not user_id: logger.error(f"Authentication required for private project {project_id}") - raise HTTPException(status_code=401, detail="Authentication required for this resource") - - account_id = project_data.get('account_id') - + raise HTTPException( + status_code=401, detail="Authentication required for this resource" + ) + + account_id = project_data.get("account_id") + # Verify account membership if account_id: - account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute() + account_user_result = ( + await client.schema("basejump") + .from_("account_user") + .select("account_role") + .eq("user_id", user_id) + .eq("account_id", account_id) + .execute() + ) if not (account_user_result.data and len(account_user_result.data) > 0): - logger.error(f"User {user_id} not authorized to access project {project_id}") - raise HTTPException(status_code=403, detail="Not authorized to access this project") - + logger.error( + f"User {user_id} not authorized to access project {project_id}" + ) + raise HTTPException( + status_code=403, detail="Not authorized to access this project" + ) + try: # Get sandbox ID from project data - sandbox_info = project_data.get('sandbox', {}) - if not sandbox_info.get('id'): - raise HTTPException(status_code=404, detail="No sandbox found for this project") - - sandbox_id = sandbox_info['id'] - + sandbox_info = project_data.get("sandbox", {}) + if not sandbox_info.get("id"): + raise HTTPException( + status_code=404, detail="No sandbox found for this project" + ) + + sandbox_id = sandbox_info["id"] + # Get or start the sandbox logger.info(f"Ensuring sandbox is active for project {project_id}") sandbox = await get_or_start_sandbox(sandbox_id) - - logger.info(f"Successfully ensured sandbox {sandbox_id} is active for project {project_id}") - + + logger.info( + f"Successfully ensured sandbox {sandbox_id} is active for project {project_id}" + ) + return { - "status": "success", + "status": "success", "sandbox_id": sandbox_id, - "message": "Sandbox is active" + "message": "Sandbox is active", } except Exception as e: - logger.error(f"Error ensuring sandbox is active for project {project_id}: {str(e)}") + logger.error( + f"Error ensuring sandbox is active for project {project_id}: {str(e)}" + ) raise HTTPException(status_code=500, detail=str(e)) diff --git a/app/daytona/docker/Dockerfile b/app/daytona/docker/Dockerfile deleted file mode 100644 index d2f12ff..0000000 --- a/app/daytona/docker/Dockerfile +++ /dev/null @@ -1,133 +0,0 @@ -FROM python:3.11-slim - -# Install system dependencies -RUN apt-get update && apt-get install -y \ - wget \ - netcat-traditional \ - gnupg \ - curl \ - unzip \ - zip \ - xvfb \ - libgconf-2-4 \ - libxss1 \ - libnss3 \ - libnspr4 \ - libasound2 \ - libatk1.0-0 \ - libatk-bridge2.0-0 \ - libcups2 \ - libdbus-1-3 \ - libdrm2 \ - libgbm1 \ - libgtk-3-0 \ - libxcomposite1 \ - libxdamage1 \ - libxfixes3 \ - libxrandr2 \ - xdg-utils \ - fonts-liberation \ - dbus \ - xauth \ - xvfb \ - x11vnc \ - tigervnc-tools \ - supervisor \ - net-tools \ - procps \ - git \ - python3-numpy \ - fontconfig \ - fonts-dejavu \ - fonts-dejavu-core \ - fonts-dejavu-extra \ - tmux \ - # PDF Processing Tools - poppler-utils \ - wkhtmltopdf \ - # Document Processing Tools - antiword \ - unrtf \ - catdoc \ - # Text Processing Tools - grep \ - gawk \ - sed \ - # File Analysis Tools - file \ - # Data Processing Tools - jq \ - csvkit \ - xmlstarlet \ - # Additional Utilities - less \ - vim \ - tree \ - rsync \ - lsof \ - iputils-ping \ - dnsutils \ - sudo \ - # OCR Tools - tesseract-ocr \ - tesseract-ocr-eng \ - && rm -rf /var/lib/apt/lists/* - -# Install Node.js and npm -RUN curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \ - && apt-get install -y nodejs \ - && npm install -g npm@latest - -# Install Cloudflare Wrangler CLI globally -RUN npm install -g wrangler - -# Install noVNC -RUN git clone https://github.com/novnc/noVNC.git /opt/novnc \ - && git clone https://github.com/novnc/websockify /opt/novnc/utils/websockify \ - && ln -s /opt/novnc/vnc.html /opt/novnc/index.html - -# Set platform for ARM64 compatibility -ARG TARGETPLATFORM=linux/amd64 - -# Set up working directory -WORKDIR /app - -# Copy requirements and install Python dependencies -COPY requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt - -# Install Playwright and browsers with system dependencies -ENV PLAYWRIGHT_BROWSERS_PATH=/ms-playwright -# Install Playwright package first -RUN pip install playwright -# Then install dependencies and browsers -RUN playwright install-deps -RUN playwright install chromium -# Verify installation -RUN python -c "from playwright.sync_api import sync_playwright; print('Playwright installation verified')" - -# Copy server script -COPY . /app -COPY server.py /app/server.py -COPY browser_api.py /app/browser_api.py - -# Set environment variables -ENV PYTHONUNBUFFERED=1 -ENV CHROME_PATH=/ms-playwright/chromium-*/chrome-linux/chrome -ENV ANONYMIZED_TELEMETRY=false -ENV DISPLAY=:99 -ENV RESOLUTION=1024x768x24 -ENV VNC_PASSWORD=vncpassword -ENV CHROME_PERSISTENT_SESSION=true -ENV RESOLUTION_WIDTH=1024 -ENV RESOLUTION_HEIGHT=768 -# Add Chrome flags to prevent multiple tabs/windows -ENV CHROME_FLAGS="--single-process --no-first-run --no-default-browser-check --disable-background-networking --disable-background-timer-throttling --disable-backgrounding-occluded-windows --disable-breakpad --disable-component-extensions-with-background-pages --disable-dev-shm-usage --disable-extensions --disable-features=TranslateUI --disable-ipc-flooding-protection --disable-renderer-backgrounding --enable-features=NetworkServiceInProcess2 --force-color-profile=srgb --metrics-recording-only --mute-audio --no-sandbox --disable-gpu" - -# Set up supervisor configuration -RUN mkdir -p /var/log/supervisor -COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf - -EXPOSE 7788 6080 5901 8000 8080 - -CMD ["/usr/bin/supervisord", "-c", "/etc/supervisor/conf.d/supervisord.conf"] \ No newline at end of file diff --git a/app/daytona/docker/browser_api.py b/app/daytona/docker/browser_api.py deleted file mode 100644 index 7456b84..0000000 --- a/app/daytona/docker/browser_api.py +++ /dev/null @@ -1,2110 +0,0 @@ -from fastapi import FastAPI, APIRouter, HTTPException, Body -from playwright.async_api import async_playwright, Browser, BrowserContext, Page -from pydantic import BaseModel -from typing import Optional, List, Dict, Any -import asyncio -import json -import logging -import base64 -from dataclasses import dataclass, field -from datetime import datetime -import os -import random -from functools import cached_property -import traceback -import pytesseract -from PIL import Image -import io - -####################################################### -# Action model definitions -####################################################### - -class Position(BaseModel): - x: int - y: int - -class ClickElementAction(BaseModel): - index: int - -class ClickCoordinatesAction(BaseModel): - x: int - y: int - -class GoToUrlAction(BaseModel): - url: str - -class InputTextAction(BaseModel): - index: int - text: str - -class ScrollAction(BaseModel): - amount: Optional[int] = None - -class SendKeysAction(BaseModel): - keys: str - -class SearchGoogleAction(BaseModel): - query: str - -class SwitchTabAction(BaseModel): - page_id: int - -class OpenTabAction(BaseModel): - url: str - -class CloseTabAction(BaseModel): - page_id: int - -class NoParamsAction(BaseModel): - pass - -class DragDropAction(BaseModel): - element_source: Optional[str] = None - element_target: Optional[str] = None - element_source_offset: Optional[Position] = None - element_target_offset: Optional[Position] = None - coord_source_x: Optional[int] = None - coord_source_y: Optional[int] = None - coord_target_x: Optional[int] = None - coord_target_y: Optional[int] = None - steps: Optional[int] = 10 - delay_ms: Optional[int] = 5 - -class DoneAction(BaseModel): - success: bool = True - text: str = "" - -####################################################### -# DOM Structure Models -####################################################### - -@dataclass -class CoordinateSet: - x: int = 0 - y: int = 0 - width: int = 0 - height: int = 0 - -@dataclass -class ViewportInfo: - width: int = 0 - height: int = 0 - scroll_x: int = 0 - scroll_y: int = 0 - -@dataclass -class HashedDomElement: - tag_name: str - attributes: Dict[str, str] - is_visible: bool - page_coordinates: Optional[CoordinateSet] = None - -@dataclass -class DOMBaseNode: - is_visible: bool - parent: Optional['DOMElementNode'] = None - -@dataclass -class DOMTextNode(DOMBaseNode): - text: str = field(default="") - type: str = 'TEXT_NODE' - - def has_parent_with_highlight_index(self) -> bool: - current = self.parent - while current is not None: - if current.highlight_index is not None: - return True - current = current.parent - return False - -@dataclass -class DOMElementNode(DOMBaseNode): - tag_name: str = field(default="") - xpath: str = field(default="") - attributes: Dict[str, str] = field(default_factory=dict) - children: List['DOMBaseNode'] = field(default_factory=list) - - is_interactive: bool = False - is_top_element: bool = False - is_in_viewport: bool = False - shadow_root: bool = False - highlight_index: Optional[int] = None - viewport_coordinates: Optional[CoordinateSet] = None - page_coordinates: Optional[CoordinateSet] = None - viewport_info: Optional[ViewportInfo] = None - - def __repr__(self) -> str: - tag_str = f'<{self.tag_name}' - for key, value in self.attributes.items(): - tag_str += f' {key}="{value}"' - tag_str += '>' - - extras = [] - if self.is_interactive: - extras.append('interactive') - if self.is_top_element: - extras.append('top') - if self.highlight_index is not None: - extras.append(f'highlight:{self.highlight_index}') - - if extras: - tag_str += f' [{", ".join(extras)}]' - - return tag_str - - @cached_property - def hash(self) -> HashedDomElement: - return HashedDomElement( - tag_name=self.tag_name, - attributes=self.attributes, - is_visible=self.is_visible, - page_coordinates=self.page_coordinates - ) - - def get_all_text_till_next_clickable_element(self, max_depth: int = -1) -> str: - text_parts = [] - - def collect_text(node: DOMBaseNode, current_depth: int) -> None: - if max_depth != -1 and current_depth > max_depth: - return - - if isinstance(node, DOMElementNode) and node != self and node.highlight_index is not None: - return - - if isinstance(node, DOMTextNode): - text_parts.append(node.text) - elif isinstance(node, DOMElementNode): - for child in node.children: - collect_text(child, current_depth + 1) - - collect_text(self, 0) - return '\n'.join(text_parts).strip() - - def clickable_elements_to_string(self, include_attributes: list[str] | None = None) -> str: - """Convert the processed DOM content to HTML.""" - formatted_text = [] - - def process_node(node: DOMBaseNode, depth: int) -> None: - if isinstance(node, DOMElementNode): - # Add element with highlight_index - if node.highlight_index is not None: - attributes_str = '' - text = node.get_all_text_till_next_clickable_element() - - # Process attributes for display - display_attributes = [] - if include_attributes: - for key, value in node.attributes.items(): - if key in include_attributes and value and value != node.tag_name: - if text and value in text: - continue # Skip if attribute value is already in the text - display_attributes.append(str(value)) - - attributes_str = ';'.join(display_attributes) - - # Build the element string - line = f'[{node.highlight_index}]<{node.tag_name}' - - # Add important attributes for identification - for attr_name in ['id', 'href', 'name', 'value', 'type']: - if attr_name in node.attributes and node.attributes[attr_name]: - line += f' {attr_name}="{node.attributes[attr_name]}"' - - # Add the text content if available - if text: - line += f'> {text}' - elif attributes_str: - line += f'> {attributes_str}' - else: - # If no text and no attributes, use the tag name - line += f'> {node.tag_name.upper()}' - - line += ' ' - formatted_text.append(line) - - # Process children regardless - for child in node.children: - process_node(child, depth + 1) - - elif isinstance(node, DOMTextNode): - # Add text only if it doesn't have a highlighted parent - if not node.has_parent_with_highlight_index() and node.is_visible: - if node.text and node.text.strip(): - formatted_text.append(node.text) - - process_node(self, 0) - result = '\n'.join(formatted_text) - return result if result.strip() else "No interactive elements found" - -@dataclass -class DOMState: - element_tree: DOMElementNode - selector_map: Dict[int, DOMElementNode] - url: str = "" - title: str = "" - pixels_above: int = 0 - pixels_below: int = 0 - -####################################################### -# Browser Action Result Model -####################################################### - -class BrowserActionResult(BaseModel): - success: bool = True - message: str = "" - error: str = "" - - # Extended state information - url: Optional[str] = None - title: Optional[str] = None - elements: Optional[str] = None # Formatted string of clickable elements - screenshot_base64: Optional[str] = None - pixels_above: int = 0 - pixels_below: int = 0 - content: Optional[str] = None - ocr_text: Optional[str] = None # Added field for OCR text - - # Additional metadata - element_count: int = 0 # Number of interactive elements found - interactive_elements: Optional[List[Dict[str, Any]]] = None # Simplified list of interactive elements - viewport_width: Optional[int] = None - viewport_height: Optional[int] = None - - class Config: - arbitrary_types_allowed = True - -####################################################### -# Browser Automation Implementation -####################################################### - -class BrowserAutomation: - def __init__(self): - self.router = APIRouter() - self.browser: Browser = None - self.browser_context: BrowserContext = None - self.pages: List[Page] = [] - self.current_page_index: int = 0 - self.logger = logging.getLogger("browser_automation") - self.include_attributes = ["id", "href", "src", "alt", "aria-label", "placeholder", "name", "role", "title", "value"] - self.screenshot_dir = os.path.join(os.getcwd(), "screenshots") - os.makedirs(self.screenshot_dir, exist_ok=True) - - # Register routes - self.router.on_startup.append(self.startup) - self.router.on_shutdown.append(self.shutdown) - - # Basic navigation - self.router.post("/automation/navigate_to")(self.navigate_to) - self.router.post("/automation/search_google")(self.search_google) - self.router.post("/automation/go_back")(self.go_back) - self.router.post("/automation/wait")(self.wait) - - # Element interaction - self.router.post("/automation/click_element")(self.click_element) - self.router.post("/automation/click_coordinates")(self.click_coordinates) - self.router.post("/automation/input_text")(self.input_text) - self.router.post("/automation/send_keys")(self.send_keys) - - # Tab management - self.router.post("/automation/switch_tab")(self.switch_tab) - self.router.post("/automation/open_tab")(self.open_tab) - self.router.post("/automation/close_tab")(self.close_tab) - - # Content actions - self.router.post("/automation/extract_content")(self.extract_content) - self.router.post("/automation/save_pdf")(self.save_pdf) - - # Scroll actions - self.router.post("/automation/scroll_down")(self.scroll_down) - self.router.post("/automation/scroll_up")(self.scroll_up) - self.router.post("/automation/scroll_to_text")(self.scroll_to_text) - - # Dropdown actions - self.router.post("/automation/get_dropdown_options")(self.get_dropdown_options) - self.router.post("/automation/select_dropdown_option")(self.select_dropdown_option) - - # Drag and drop - self.router.post("/automation/drag_drop")(self.drag_drop) - - async def startup(self): - """Initialize the browser instance on startup""" - try: - print("Starting browser initialization...") - playwright = await async_playwright().start() - print("Playwright started, launching browser...") - - # Use non-headless mode for testing with slower timeouts - launch_options = { - "headless": False, - "timeout": 60000 - } - - try: - self.browser = await playwright.chromium.launch(**launch_options) - self.browser_context = await self.browser.new_context(viewport={'width': 1024, 'height': 768}) - print("Browser launched successfully") - except Exception as browser_error: - print(f"Failed to launch browser: {browser_error}") - # Try with minimal options - print("Retrying with minimal options...") - launch_options = {"timeout": 90000} - self.browser = await playwright.chromium.launch(**launch_options) - self.browser_context = await self.browser.new_context(viewport={'width': 1024, 'height': 768}) - print("Browser launched with minimal options") - - try: - await self.get_current_page() - print("Found existing page, using it") - self.current_page_index = 0 - except Exception as page_error: - print(f"Error finding existing page, creating new one. ( {page_error})") - page = await self.browser_context.new_page() - print("New page created successfully") - self.pages.append(page) - self.current_page_index = 0 - # Navigate directly to google.com instead of about:blank - await page.goto("https://www.google.com", wait_until="domcontentloaded", timeout=30000) - print("Navigated to google.com") - - try: - self.browser_context.on("page", self.handle_page_created) - except Exception as e: - print(f"Error setting up page event handler: {e}") - traceback.print_exc() - - - print("Browser initialization completed successfully") - except Exception as e: - print(f"Browser startup error: {str(e)}") - traceback.print_exc() - raise RuntimeError(f"Browser initialization failed: {str(e)}") - - async def shutdown(self): - """Clean up browser instance on shutdown""" - if self.browser_context: - await self.browser_context.close() - if self.browser: - await self.browser.close() - - async def handle_page_created(self, page: Page): - """Handle new page creation""" - await asyncio.sleep(0.5) - self.pages.append(page) - self.current_page_index = len(self.pages) - 1 - print(f"Page created: {page.url}; current page index: {self.current_page_index}") - - async def get_current_page(self) -> Page: - """Get the current active page""" - if not self.pages: - raise HTTPException(status_code=500, detail="No browser pages available") - return self.pages[self.current_page_index] - - async def get_selector_map(self) -> Dict[int, DOMElementNode]: - """Get a map of selectable elements on the page""" - page = await self.get_current_page() - - # Create a selector map for interactive elements - selector_map = {} - - try: - # More comprehensive JavaScript to find interactive elements - elements_js = """ - (() => { - // Helper function to get all attributes as an object - function getAttributes(el) { - const attributes = {}; - for (const attr of el.attributes) { - attributes[attr.name] = attr.value; - } - return attributes; - } - - // Find all potentially interactive elements - const interactiveElements = Array.from(document.querySelectorAll( - 'a, button, input, select, textarea, [role="button"], [role="link"], [role="checkbox"], [role="radio"], [tabindex]:not([tabindex="-1"])' - )); - - // Filter for visible elements - const visibleElements = interactiveElements.filter(el => { - const style = window.getComputedStyle(el); - const rect = el.getBoundingClientRect(); - return style.display !== 'none' && - style.visibility !== 'hidden' && - style.opacity !== '0' && - rect.width > 0 && - rect.height > 0; - }); - - // Map to our expected structure - return visibleElements.map((el, index) => { - const rect = el.getBoundingClientRect(); - const isInViewport = rect.top >= 0 && - rect.left >= 0 && - rect.bottom <= window.innerHeight && - rect.right <= window.innerWidth; - - return { - index: index + 1, - tagName: el.tagName.toLowerCase(), - text: el.innerText || el.value || '', - attributes: getAttributes(el), - isVisible: true, - isInteractive: true, - pageCoordinates: { - x: rect.left + window.scrollX, - y: rect.top + window.scrollY, - width: rect.width, - height: rect.height - }, - viewportCoordinates: { - x: rect.left, - y: rect.top, - width: rect.width, - height: rect.height - }, - isInViewport: isInViewport - }; - }); - })(); - """ - - elements = await page.evaluate(elements_js) - print(f"Found {len(elements)} interactive elements in selector map") - - # Create a root element for the tree - root = DOMElementNode( - is_visible=True, - tag_name="body", - is_interactive=False, - is_top_element=True - ) - - # Create element nodes for each element - for idx, el in enumerate(elements): - # Create coordinate sets - page_coordinates = None - viewport_coordinates = None - - if 'pageCoordinates' in el: - coords = el['pageCoordinates'] - page_coordinates = CoordinateSet( - x=coords.get('x', 0), - y=coords.get('y', 0), - width=coords.get('width', 0), - height=coords.get('height', 0) - ) - - if 'viewportCoordinates' in el: - coords = el['viewportCoordinates'] - viewport_coordinates = CoordinateSet( - x=coords.get('x', 0), - y=coords.get('y', 0), - width=coords.get('width', 0), - height=coords.get('height', 0) - ) - - # Create the element node - element_node = DOMElementNode( - is_visible=el.get('isVisible', True), - tag_name=el.get('tagName', 'div'), - attributes=el.get('attributes', {}), - is_interactive=el.get('isInteractive', True), - is_in_viewport=el.get('isInViewport', False), - highlight_index=el.get('index', idx + 1), - page_coordinates=page_coordinates, - viewport_coordinates=viewport_coordinates - ) - - # Add a text node if there's text content - if el.get('text'): - text_node = DOMTextNode(is_visible=True, text=el.get('text', '')) - text_node.parent = element_node - element_node.children.append(text_node) - - selector_map[el.get('index', idx + 1)] = element_node - root.children.append(element_node) - element_node.parent = root - - except Exception as e: - print(f"Error getting selector map: {e}") - traceback.print_exc() - # Create a dummy element to avoid breaking tests - dummy = DOMElementNode( - is_visible=True, - tag_name="a", - attributes={'href': '#'}, - is_interactive=True, - highlight_index=1 - ) - dummy_text = DOMTextNode(is_visible=True, text="Dummy Element") - dummy_text.parent = dummy - dummy.children.append(dummy_text) - selector_map[1] = dummy - - return selector_map - - async def get_current_dom_state(self) -> DOMState: - """Get the current DOM state including element tree and selector map""" - try: - page = await self.get_current_page() - selector_map = await self.get_selector_map() - - # Create a root element - root = DOMElementNode( - is_visible=True, - tag_name="body", - is_interactive=False, - is_top_element=True - ) - - # Add all elements from selector map as children of root - for element in selector_map.values(): - if element.parent is None: - element.parent = root - root.children.append(element) - - # Get basic page info - url = page.url - try: - title = await page.title() - except: - title = "Unknown Title" - - # Get more accurate scroll information - fix JavaScript syntax - try: - scroll_info = await page.evaluate(""" - () => { - const body = document.body; - const html = document.documentElement; - const totalHeight = Math.max( - body.scrollHeight, body.offsetHeight, - html.clientHeight, html.scrollHeight, html.offsetHeight - ); - const scrollY = window.scrollY || window.pageYOffset; - const windowHeight = window.innerHeight; - - return { - pixelsAbove: scrollY, - pixelsBelow: Math.max(0, totalHeight - scrollY - windowHeight), - totalHeight: totalHeight, - viewportHeight: windowHeight - }; - } - """) - pixels_above = scroll_info.get('pixelsAbove', 0) - pixels_below = scroll_info.get('pixelsBelow', 0) - except Exception as e: - print(f"Error getting scroll info: {e}") - pixels_above = 0 - pixels_below = 0 - - return DOMState( - element_tree=root, - selector_map=selector_map, - url=url, - title=title, - pixels_above=pixels_above, - pixels_below=pixels_below - ) - except Exception as e: - print(f"Error getting DOM state: {e}") - traceback.print_exc() - # Return a minimal valid state to avoid breaking tests - dummy_root = DOMElementNode( - is_visible=True, - tag_name="body", - is_interactive=False, - is_top_element=True - ) - dummy_map = {1: dummy_root} - current_url = "unknown" - try: - if 'page' in locals(): - current_url = page.url - except: - pass - return DOMState( - element_tree=dummy_root, - selector_map=dummy_map, - url=current_url, - title="Error page", - pixels_above=0, - pixels_below=0 - ) - - async def take_screenshot(self) -> str: - """Take a screenshot and return as base64 encoded string""" - try: - page = await self.get_current_page() - - # Wait for network to be idle and DOM to be stable - try: - await page.wait_for_load_state("networkidle", timeout=60000) # Increased timeout to 60s - except Exception as e: - print(f"Warning: Network idle timeout, proceeding anyway: {e}") - - # Wait for any animations to complete - # await page.wait_for_timeout(1000) # Wait 1 second for animations - - # Take screenshot with increased timeout and better options - screenshot_bytes = await page.screenshot( - type='jpeg', - quality=60, - full_page=False, - timeout=60000, # Increased timeout to 60s - scale='device' # Use device scale factor - ) - - return base64.b64encode(screenshot_bytes).decode('utf-8') - except Exception as e: - print(f"Error taking screenshot: {e}") - traceback.print_exc() - # Return an empty string rather than failing - return "" - - async def save_screenshot_to_file(self) -> str: - """Take a screenshot and save to file, returning the path""" - try: - page = await self.get_current_page() - timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') - random_id = random.randint(1000, 9999) - filename = f"screenshot_{timestamp}_{random_id}.jpg" - filepath = os.path.join(self.screenshot_dir, filename) - - await page.screenshot(path=filepath, type='jpeg', quality=60, full_page=False) - return filepath - except Exception as e: - print(f"Error saving screenshot: {e}") - return "" - - async def extract_ocr_text_from_screenshot(self, screenshot_base64: str) -> str: - """Extract text from screenshot using OCR""" - if not screenshot_base64: - return "" - - try: - # Decode base64 to image - image_bytes = base64.b64decode(screenshot_base64) - image = Image.open(io.BytesIO(image_bytes)) - - # Extract text using pytesseract - ocr_text = pytesseract.image_to_string(image) - - # Clean up the text - ocr_text = ocr_text.strip() - - return ocr_text - except Exception as e: - print(f"Error performing OCR: {e}") - traceback.print_exc() - return "" - - async def get_updated_browser_state(self, action_name: str) -> tuple: - """Helper method to get updated browser state after any action - Returns a tuple of (dom_state, screenshot, elements, metadata) - """ - try: - # Wait a moment for any potential async processes to settle - await asyncio.sleep(0.5) - - # Get updated state - dom_state = await self.get_current_dom_state() - screenshot = await self.take_screenshot() - - # Format elements for output - elements = dom_state.element_tree.clickable_elements_to_string( - include_attributes=self.include_attributes - ) - - # Collect additional metadata - page = await self.get_current_page() - metadata = {} - - # Get element count - metadata['element_count'] = len(dom_state.selector_map) - - # Create simplified interactive elements list - interactive_elements = [] - for idx, element in dom_state.selector_map.items(): - element_info = { - 'index': idx, - 'tag_name': element.tag_name, - 'text': element.get_all_text_till_next_clickable_element(), - 'is_in_viewport': element.is_in_viewport - } - - # Add key attributes - for attr_name in ['id', 'href', 'src', 'alt', 'placeholder', 'name', 'role', 'title', 'type']: - if attr_name in element.attributes: - element_info[attr_name] = element.attributes[attr_name] - - interactive_elements.append(element_info) - - metadata['interactive_elements'] = interactive_elements - - # Get viewport dimensions - Fix syntax error in JavaScript - try: - viewport = await page.evaluate(""" - () => { - return { - width: window.innerWidth, - height: window.innerHeight - }; - } - """) - metadata['viewport_width'] = viewport.get('width', 0) - metadata['viewport_height'] = viewport.get('height', 0) - except Exception as e: - print(f"Error getting viewport dimensions: {e}") - metadata['viewport_width'] = 0 - metadata['viewport_height'] = 0 - - # Extract OCR text from screenshot if available - ocr_text = "" - if screenshot: - ocr_text = await self.extract_ocr_text_from_screenshot(screenshot) - metadata['ocr_text'] = ocr_text - - print(f"Got updated state after {action_name}: {len(dom_state.selector_map)} elements") - return dom_state, screenshot, elements, metadata - except Exception as e: - print(f"Error getting updated state after {action_name}: {e}") - traceback.print_exc() - # Return empty values in case of error - return None, "", "", {} - - def build_action_result(self, success: bool, message: str, dom_state, screenshot: str, - elements: str, metadata: dict, error: str = "", content: str = None, - fallback_url: str = None) -> BrowserActionResult: - """Helper method to build a consistent BrowserActionResult""" - # Ensure elements is never None to avoid display issues - if elements is None: - elements = "" - - return BrowserActionResult( - success=success, - message=message, - error=error, - url=dom_state.url if dom_state else fallback_url or "", - title=dom_state.title if dom_state else "", - elements=elements, - screenshot_base64=screenshot, - pixels_above=dom_state.pixels_above if dom_state else 0, - pixels_below=dom_state.pixels_below if dom_state else 0, - content=content, - ocr_text=metadata.get('ocr_text', ""), - element_count=metadata.get('element_count', 0), - interactive_elements=metadata.get('interactive_elements', []), - viewport_width=metadata.get('viewport_width', 0), - viewport_height=metadata.get('viewport_height', 0) - ) - - # Basic Navigation Actions - - async def navigate_to(self, action: GoToUrlAction = Body(...)): - """Navigate to a specified URL""" - try: - page = await self.get_current_page() - await page.goto(action.url, wait_until="domcontentloaded") - await page.wait_for_load_state("networkidle", timeout=10000) - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"navigate_to({action.url})") - - result = self.build_action_result( - True, - f"Navigated to {action.url}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - - print(f"Navigation result: success={result.success}, url={result.url}") - return result - except Exception as e: - print(f"Navigation error: {str(e)}") - traceback.print_exc() - # Try to get some state info even after error - try: - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("navigate_error_recovery") - return self.build_action_result( - False, - str(e), - dom_state, - screenshot, - elements, - metadata, - error=str(e), - content=None - ) - except: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def search_google(self, action: SearchGoogleAction = Body(...)): - """Search Google with the provided query""" - try: - page = await self.get_current_page() - search_url = f"https://www.google.com/search?q={action.query}" - await page.goto(search_url) - await page.wait_for_load_state() - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"search_google({action.query})") - - return self.build_action_result( - True, - f"Searched for '{action.query}' in Google", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - print(f"Search error: {str(e)}") - traceback.print_exc() - # Try to get some state info even after error - try: - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("search_error_recovery") - return self.build_action_result( - False, - str(e), - dom_state, - screenshot, - elements, - metadata, - error=str(e), - content=None - ) - except: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def go_back(self, _: NoParamsAction = Body(...)): - """Navigate back in browser history""" - try: - page = await self.get_current_page() - await page.go_back() - await page.wait_for_load_state() - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("go_back") - - return self.build_action_result( - True, - "Navigated back", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def wait(self, seconds: int = Body(3)): - """Wait for the specified number of seconds""" - try: - await asyncio.sleep(seconds) - - # Get updated state after waiting - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"wait({seconds} seconds)") - - return self.build_action_result( - True, - f"Waited for {seconds} seconds", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Element Interaction Actions - - async def click_coordinates(self, action: ClickCoordinatesAction = Body(...)): - """Click at specific x,y coordinates on the page""" - try: - page = await self.get_current_page() - - # Perform the click at the specified coordinates - await page.mouse.click(action.x, action.y) - - # Give time for any navigation or DOM updates to occur - await page.wait_for_load_state("networkidle", timeout=5000) - - await asyncio.sleep(1) - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"click_coordinates({action.x}, {action.y})") - - return self.build_action_result( - True, - f"Clicked at coordinates ({action.x}, {action.y})", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - print(f"Error in click_coordinates: {e}") - traceback.print_exc() - - # Try to get state even after error - try: - await asyncio.sleep(1) - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("click_coordinates_error_recovery") - return self.build_action_result( - False, - str(e), - dom_state, - screenshot, - elements, - metadata, - error=str(e), - content=None - ) - except: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def click_element(self, action: ClickElementAction = Body(...)): - """Click on an element by index""" - try: - page = await self.get_current_page() - - # Get the current state and selector map *before* the click - initial_dom_state = await self.get_current_dom_state() - selector_map = initial_dom_state.selector_map - - if action.index not in selector_map: - # Get updated state even if element not found initially - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"click_element_error (index {action.index} not found)") - return self.build_action_result( - False, - f"Element with index {action.index} not found", - dom_state, # Use the latest state - screenshot, - elements, - metadata, - error=f"Element with index {action.index} not found" - ) - - element_to_click = selector_map[action.index] - print(f"Attempting to click element: {element_to_click}") - - # Construct a more reliable selector using JavaScript evaluation - # Find the element based on its properties captured in selector_map - js_selector_script = """ - (targetElementInfo) => { - const interactiveElements = Array.from(document.querySelectorAll( - 'a, button, input, select, textarea, [role="button"], [role="link"], [role="checkbox"], [role="radio"], [tabindex]:not([tabindex="-1"])' - )); - - const visibleElements = interactiveElements.filter(el => { - const style = window.getComputedStyle(el); - const rect = el.getBoundingClientRect(); - return style.display !== 'none' && style.visibility !== 'hidden' && style.opacity !== '0' && rect.width > 0 && rect.height > 0; - }); - - if (targetElementInfo.index > 0 && targetElementInfo.index <= visibleElements.length) { - // Return the element at the specified index (1-based) - return visibleElements[targetElementInfo.index - 1]; - } - return null; // Element not found at the expected index - } - """ - - element_info = {'index': action.index} # Pass the target index to the script - - target_element_handle = await page.evaluate_handle(js_selector_script, element_info) - - click_success = False - error_message = "" - - if await target_element_handle.evaluate("node => node !== null"): - try: - # Use Playwright's recommended way: click the handle - # Add timeout and wait for element to be stable - await target_element_handle.click(timeout=5000) - click_success = True - print(f"Successfully clicked element handle for index {action.index}") - except Exception as click_error: - error_message = f"Error clicking element handle: {click_error}" - print(error_message) - # Optional: Add fallback methods here if needed - # e.g., target_element_handle.dispatch_event('click') - else: - error_message = f"Could not locate the target element handle for index {action.index} using JS script." - print(error_message) - - - # Wait for potential page changes/network activity - try: - await page.wait_for_load_state("networkidle", timeout=5000) - except Exception as wait_error: - print(f"Timeout or error waiting for network idle after click: {wait_error}") - await asyncio.sleep(1) - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"click_element({action.index})") - - return self.build_action_result( - click_success, - f"Clicked element with index {action.index}" if click_success else f"Attempted to click element {action.index} but failed. Error: {error_message}", - dom_state, - screenshot, - elements, - metadata, - error=error_message if not click_success else "", - content=None - ) - - except Exception as e: - print(f"Error in click_element: {e}") - traceback.print_exc() - # Try to get state even after error - try: - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("click_element_error_recovery") - return self.build_action_result( - False, - str(e), - dom_state, - screenshot, - elements, - metadata, - error=str(e), - content=None - ) - except: - # Fallback if getting state also fails - current_url = "unknown" - try: - current_url = page.url # Try to get at least the URL - except: - pass - return self.build_action_result( - False, - str(e), - None, # No DOM state available - "", # No screenshot - "", # No elements string - {}, # Empty metadata - error=str(e), - content=None, - fallback_url=current_url - ) - - async def input_text(self, action: InputTextAction = Body(...)): - """Input text into an element""" - try: - page = await self.get_current_page() - selector_map = await self.get_selector_map() - - if action.index not in selector_map: - return self.build_action_result( - False, - f"Element with index {action.index} not found", - None, - "", - "", - {}, - error=f"Element with index {action.index} not found" - ) - - # In a real implementation, we would use the selector map to get the element's - # properties and use them to find and type into the element - element = selector_map[action.index] - - # Use CSS selector or XPath to locate and type into the element - await page.wait_for_timeout(500) # Small delay before typing - - # Demo implementation - would use proper selectors in production - if element.attributes.get("id"): - await page.fill(f"#{element.attributes['id']}", action.text) - elif element.attributes.get("class"): - class_selector = f".{element.attributes['class'].replace(' ', '.')}" - await page.fill(class_selector, action.text) - else: - # Fallback to xpath - await page.fill(f"//{element.tag_name}[{action.index}]", action.text) - - await asyncio.sleep(1) - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"input_text({action.index}, '{action.text}')") - - return self.build_action_result( - True, - f"Input '{action.text}' into element with index {action.index}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def send_keys(self, action: SendKeysAction = Body(...)): - """Send keyboard keys""" - try: - page = await self.get_current_page() - await page.keyboard.press(action.keys) - - await asyncio.sleep(1) - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"send_keys({action.keys})") - - return self.build_action_result( - True, - f"Sent keys: {action.keys}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Tab Management Actions - - async def switch_tab(self, action: SwitchTabAction = Body(...)): - """Switch to a different tab by index""" - try: - if 0 <= action.page_id < len(self.pages): - self.current_page_index = action.page_id - page = await self.get_current_page() - await page.wait_for_load_state() - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"switch_tab({action.page_id})") - - return self.build_action_result( - True, - f"Switched to tab {action.page_id}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - else: - return self.build_action_result( - False, - f"Tab {action.page_id} not found", - None, - "", - "", - {}, - error=f"Tab {action.page_id} not found" - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def open_tab(self, action: OpenTabAction = Body(...)): - """Open a new tab with the specified URL""" - try: - print(f"Attempting to open new tab with URL: {action.url}") - # Create new page in same browser instance - new_page = await self.browser_context.new_page() - print(f"New page created successfully") - - # Navigate to the URL - await new_page.goto(action.url, wait_until="domcontentloaded") - await new_page.wait_for_load_state("networkidle", timeout=10000) - print(f"Navigated to URL in new tab: {action.url}") - - # Add to page list and make it current - self.pages.append(new_page) - self.current_page_index = len(self.pages) - 1 - print(f"New tab added as index {self.current_page_index}") - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"open_tab({action.url})") - - return self.build_action_result( - True, - f"Opened new tab with URL: {action.url}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - print("****"*10) - print(f"Error opening tab: {e}") - print(traceback.format_exc()) - print("****"*10) - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def close_tab(self, action: CloseTabAction = Body(...)): - """Close a tab by index""" - try: - if 0 <= action.page_id < len(self.pages): - page = self.pages[action.page_id] - url = page.url - await page.close() - self.pages.pop(action.page_id) - - # Adjust current index if needed - if self.current_page_index >= len(self.pages): - self.current_page_index = max(0, len(self.pages) - 1) - elif self.current_page_index >= action.page_id: - self.current_page_index = max(0, self.current_page_index - 1) - - # Get updated state after action - page = await self.get_current_page() - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"close_tab({action.page_id})") - - return self.build_action_result( - True, - f"Closed tab {action.page_id} with URL: {url}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - else: - return self.build_action_result( - False, - f"Tab {action.page_id} not found", - None, - "", - "", - {}, - error=f"Tab {action.page_id} not found" - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Content Actions - - async def extract_content(self, goal: str = Body(...)): - """Extract content from the current page based on the provided goal""" - try: - page = await self.get_current_page() - content = await page.content() - - # In a full implementation, we would use an LLM to extract specific content - # based on the goal. For this example, we'll extract visible text. - extracted_text = await page.evaluate(""" - Array.from(document.querySelectorAll('p, h1, h2, h3, h4, h5, h6, li, span, div')) - .filter(el => { - const style = window.getComputedStyle(el); - return style.display !== 'none' && - style.visibility !== 'hidden' && - style.opacity !== '0' && - el.innerText && - el.innerText.trim().length > 0; - }) - .map(el => el.innerText.trim()) - .join('\\n\\n'); - """) - - # Get updated state - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"extract_content({goal})") - - return self.build_action_result( - True, - f"Content extracted based on goal: {goal}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=extracted_text - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def save_pdf(self): - """Save the current page as a PDF""" - try: - page = await self.get_current_page() - timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') - random_id = random.randint(1000, 9999) - filename = f"page_{timestamp}_{random_id}.pdf" - filepath = os.path.join(self.screenshot_dir, filename) - - await page.pdf(path=filepath) - - # Get updated state - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state("save_pdf") - - return self.build_action_result( - True, - f"Saved page as PDF: {filepath}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Scroll Actions - - async def scroll_down(self, action: ScrollAction = Body(...)): - """Scroll down the page""" - try: - page = await self.get_current_page() - if action.amount is not None: - await page.evaluate(f"window.scrollBy(0, {action.amount});") - amount_str = f"{action.amount} pixels" - else: - await page.evaluate("window.scrollBy(0, window.innerHeight);") - amount_str = "one page" - - await page.wait_for_timeout(500) # Wait for scroll to complete - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"scroll_down({amount_str})") - - return self.build_action_result( - True, - f"Scrolled down by {amount_str}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def scroll_up(self, action: ScrollAction = Body(...)): - """Scroll up the page""" - try: - page = await self.get_current_page() - if action.amount is not None: - await page.evaluate(f"window.scrollBy(0, -{action.amount});") - amount_str = f"{action.amount} pixels" - else: - await page.evaluate("window.scrollBy(0, -window.innerHeight);") - amount_str = "one page" - - await page.wait_for_timeout(500) # Wait for scroll to complete - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"scroll_up({amount_str})") - - return self.build_action_result( - True, - f"Scrolled up by {amount_str}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - async def scroll_to_text(self, text: str = Body(...)): - """Scroll to text on the page""" - try: - page = await self.get_current_page() - locators = [ - page.get_by_text(text, exact=False), - page.locator(f"text={text}"), - page.locator(f"//*[contains(text(), '{text}')]"), - ] - - found = False - for locator in locators: - try: - if await locator.count() > 0 and await locator.first.is_visible(): - await locator.first.scroll_into_view_if_needed() - await asyncio.sleep(0.5) # Wait for scroll to complete - found = True - break - except Exception: - continue - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"scroll_to_text({text})") - - message = f"Scrolled to text: {text}" if found else f"Text '{text}' not found or not visible on page" - - return self.build_action_result( - found, - message, - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Dropdown Actions - - async def get_dropdown_options(self, index: int = Body(...)): - """Get all options from a dropdown""" - try: - page = await self.get_current_page() - selector_map = await self.get_selector_map() - - if index not in selector_map: - return self.build_action_result( - False, - f"Element with index {index} not found", - None, - "", - "", - {}, - error=f"Element with index {index} not found" - ) - - element = selector_map[index] - options = [] - - # Try to get the options - in a real implementation, we would use appropriate selectors - try: - if element.tag_name.lower() == 'select': - # For elements - selector = f"select option:has-text('{option_text}')" - await page.select_option( - f"#{element.attributes.get('id')}" if element.attributes.get('id') else f"//select[{index}]", - label=option_text - ) - else: - # For custom dropdowns - # First click to open the dropdown - if element.attributes.get('id'): - await page.click(f"#{element.attributes.get('id')}") - else: - await page.click(f"//{element.tag_name}[{index}]") - - await page.wait_for_timeout(500) - - # Then try to click the option - await page.click(f"text={option_text}") - - await page.wait_for_timeout(500) - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"select_dropdown_option({index}, '{option_text}')") - - return self.build_action_result( - True, - f"Selected option '{option_text}' from dropdown with index {index}", - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - - # Drag and Drop - - async def drag_drop(self, action: DragDropAction = Body(...)): - """Perform drag and drop operation""" - try: - page = await self.get_current_page() - - # Element-based drag and drop - if action.element_source and action.element_target: - # In a real implementation, we would get the elements and perform the drag - source_desc = action.element_source - target_desc = action.element_target - - # We would locate the elements using selectors and perform the drag - # For this example, we'll use a simplified version - await page.evaluate(""" - console.log("Simulating drag and drop between elements"); - """) - - message = f"Dragged element '{source_desc}' to '{target_desc}'" - - # Coordinate-based drag and drop - elif all(coord is not None for coord in [ - action.coord_source_x, action.coord_source_y, - action.coord_target_x, action.coord_target_y - ]): - source_x = action.coord_source_x - source_y = action.coord_source_y - target_x = action.coord_target_x - target_y = action.coord_target_y - - # Perform the drag - await page.mouse.move(source_x, source_y) - await page.mouse.down() - - steps = max(1, action.steps or 10) - delay_ms = max(0, action.delay_ms or 5) - - for i in range(1, steps + 1): - ratio = i / steps - intermediate_x = int(source_x + (target_x - source_x) * ratio) - intermediate_y = int(source_y + (target_y - source_y) * ratio) - await page.mouse.move(intermediate_x, intermediate_y) - if delay_ms > 0: - await asyncio.sleep(delay_ms / 1000) - - await page.mouse.move(target_x, target_y) - await page.mouse.up() - - message = f"Dragged from ({source_x}, {source_y}) to ({target_x}, {target_y})" - else: - return self.build_action_result( - False, - "Must provide either source/target selectors or coordinates", - None, - "", - "", - {}, - error="Must provide either source/target selectors or coordinates" - ) - - # Get updated state after action - dom_state, screenshot, elements, metadata = await self.get_updated_browser_state(f"drag_drop({action.element_source}, {action.element_target})") - - return self.build_action_result( - True, - message, - dom_state, - screenshot, - elements, - metadata, - error="", - content=None - ) - except Exception as e: - return self.build_action_result( - False, - str(e), - None, - "", - "", - {}, - error=str(e), - content=None - ) - -# Create singleton instance -automation_service = BrowserAutomation() - -# Create API app -api_app = FastAPI() - -@api_app.get("/api") -async def health_check(): - return {"status": "ok", "message": "API server is running"} - -# Include automation service router with /api prefix -api_app.include_router(automation_service.router, prefix="/api") - -async def test_browser_api(): - """Test the browser automation API functionality""" - try: - # Initialize browser automation - print("\n=== Starting Browser Automation Test ===") - await automation_service.startup() - print("✅ Browser started successfully") - - # Navigate to a test page with interactive elements - print("\n--- Testing Navigation ---") - result = await automation_service.navigate_to(GoToUrlAction(url="https://www.youtube.com")) - print(f"Navigation status: {'✅ Success' if result.success else '❌ Failed'}") - if not result.success: - print(f"Error: {result.error}") - return - - print(f"URL: {result.url}") - print(f"Title: {result.title}") - - # Check DOM state and elements - print(f"\nFound {result.element_count} interactive elements") - if result.elements and result.elements.strip(): - print("Elements:") - print(result.elements) - else: - print("No formatted elements found, but DOM was processed") - - # Display interactive elements as JSON - if result.interactive_elements and len(result.interactive_elements) > 0: - print("\nInteractive elements summary:") - for el in result.interactive_elements: - print(f" [{el['index']}] <{el['tag_name']}> {el.get('text', '')[:30]}") - - # Screenshot info - print(f"\nScreenshot captured: {'Yes' if result.screenshot_base64 else 'No'}") - print(f"Viewport size: {result.viewport_width}x{result.viewport_height}") - - # Test OCR extraction from screenshot - print("\n--- Testing OCR Text Extraction ---") - if result.ocr_text: - print("OCR text extracted from screenshot:") - print("=== OCR TEXT START ===") - print(result.ocr_text) - print("=== OCR TEXT END ===") - print(f"OCR text length: {len(result.ocr_text)} characters") - print(result.ocr_text) - else: - print("No OCR text extracted from screenshot") - - await asyncio.sleep(2) - - # Test search functionality - print("\n--- Testing Search ---") - result = await automation_service.search_google(SearchGoogleAction(query="browser automation")) - print(f"Search status: {'✅ Success' if result.success else '❌ Failed'}") - if not result.success: - print(f"Error: {result.error}") - else: - print(f"Found {result.element_count} elements after search") - print(f"Page title: {result.title}") - - # Test OCR extraction from search results - if result.ocr_text: - print("\nOCR text from search results:") - print("=== OCR TEXT START ===") - print(result.ocr_text) - print("=== OCR TEXT END ===") - else: - print("\nNo OCR text extracted from search results") - - await asyncio.sleep(2) - - # Test scrolling - print("\n--- Testing Scrolling ---") - result = await automation_service.scroll_down(ScrollAction(amount=300)) - print(f"Scroll status: {'✅ Success' if result.success else '❌ Failed'}") - if result.success: - print(f"Pixels above viewport: {result.pixels_above}") - print(f"Pixels below viewport: {result.pixels_below}") - - await asyncio.sleep(2) - - # Test clicking on an element - print("\n--- Testing Element Click ---") - if result.element_count > 0: - click_result = await automation_service.click_element(ClickElementAction(index=1)) - print(f"Click status: {'✅ Success' if click_result.success else '❌ Failed'}") - print(f"Message: {click_result.message}") - print(f"New URL after click: {click_result.url}") - else: - print("Skipping click test - no elements found") - - await asyncio.sleep(2) - - # Test clicking on coordinates - print("\n--- Testing Click Coordinates ---") - coord_click_result = await automation_service.click_coordinates(ClickCoordinatesAction(x=100, y=100)) - print(f"Coordinate click status: {'✅ Success' if coord_click_result.success else '❌ Failed'}") - print(f"Message: {coord_click_result.message}") - print(f"URL after coordinate click: {coord_click_result.url}") - - await asyncio.sleep(2) - - # Test extracting content - print("\n--- Testing Content Extraction ---") - content_result = await automation_service.extract_content("test goal") - print(f"Content extraction status: {'✅ Success' if content_result.success else '❌ Failed'}") - if content_result.content: - content_preview = content_result.content[:100] + "..." if len(content_result.content) > 100 else content_result.content - print(f"Content sample: {content_preview}") - print(f"Total content length: {len(content_result.content)} chars") - else: - print("No content was extracted") - - # Test tab management - print("\n--- Testing Tab Management ---") - tab_result = await automation_service.open_tab(OpenTabAction(url="https://www.example.org")) - print(f"New tab status: {'✅ Success' if tab_result.success else '❌ Failed'}") - if tab_result.success: - print(f"New tab title: {tab_result.title}") - print(f"Interactive elements: {tab_result.element_count}") - - print("\n✅ All tests completed successfully!") - - except Exception as e: - print(f"\n❌ Test failed: {str(e)}") - traceback.print_exc() - finally: - # Ensure browser is closed - print("\n--- Cleaning up ---") - await automation_service.shutdown() - print("Browser closed") - -async def test_browser_api_2(): - """Test the browser automation API functionality on the chess page""" - try: - # Initialize browser automation - print("\n=== Starting Browser Automation Test 2 (Chess Page) ===") - await automation_service.startup() - print("✅ Browser started successfully") - - # Navigate to the chess test page - print("\n--- Testing Navigation to Chess Page ---") - test_url = "https://dat-lequoc.github.io/chess-for-suna/chess.html" - result = await automation_service.navigate_to(GoToUrlAction(url=test_url)) - print(f"Navigation status: {'✅ Success' if result.success else '❌ Failed'}") - if not result.success: - print(f"Error: {result.error}") - return - - print(f"URL: {result.url}") - print(f"Title: {result.title}") - - # Check DOM state and elements - print(f"\nFound {result.element_count} interactive elements") - if result.elements and result.elements.strip(): - print("Elements:") - print(result.elements) - else: - print("No formatted elements found, but DOM was processed") - - # Display interactive elements as JSON - if result.interactive_elements and len(result.interactive_elements) > 0: - print("\nInteractive elements summary:") - for el in result.interactive_elements: - print(f" [{el['index']}] <{el['tag_name']}> {el.get('text', '')[:30]}") - - # Screenshot info - print(f"\nScreenshot captured: {'Yes' if result.screenshot_base64 else 'No'}") - print(f"Viewport size: {result.viewport_width}x{result.viewport_height}") - - await asyncio.sleep(2) - - # Test clicking on an element (e.g., a chess square) - print("\n--- Testing Element Click (element 5) ---") - if result.element_count > 4: # Ensure element 5 exists - click_index = 5 - click_result = await automation_service.click_element(ClickElementAction(index=click_index)) - print(f"Click status for element {click_index}: {'✅ Success' if click_result.success else '❌ Failed'}") - print(f"Message: {click_result.message}") - print(f"URL after click: {click_result.url}") - - # Retrieve and display elements again after click - print(f"\n--- Retrieving elements after clicking element {click_index} ---") - if click_result.elements and click_result.elements.strip(): - print("Updated Elements:") - print(click_result.elements) - else: - print("No formatted elements found after click.") - - if click_result.interactive_elements and len(click_result.interactive_elements) > 0: - print("\nUpdated interactive elements summary:") - for el in click_result.interactive_elements: - print(f" [{el['index']}] <{el['tag_name']}> {el.get('text', '')[:30]}") - else: - print("No interactive elements found after click.") - - # Test clicking element 1 after the first click - print("\n--- Testing Element Click (element 1 after clicking 5) ---") - if click_result.element_count > 0: # Check if there are still elements - click_index_2 = 1 - click_result_2 = await automation_service.click_element(ClickElementAction(index=click_index_2)) - print(f"Click status for element {click_index_2}: {'✅ Success' if click_result_2.success else '❌ Failed'}") - print(f"Message: {click_result_2.message}") - print(f"URL after click: {click_result_2.url}") - - # Retrieve and display elements again after the second click - print(f"\n--- Retrieving elements after clicking element {click_index_2} ---") - if click_result_2.elements and click_result_2.elements.strip(): - print("Elements after second click:") - print(click_result_2.elements) - else: - print("No formatted elements found after second click.") - - if click_result_2.interactive_elements and len(click_result_2.interactive_elements) > 0: - print("\nInteractive elements summary after second click:") - for el in click_result_2.interactive_elements: - print(f" [{el['index']}] <{el['tag_name']}> {el.get('text', '')[:30]}") - else: - print("No interactive elements found after second click.") - else: - print("Skipping second element click test - no elements found after first click.") - - else: - print("Skipping element click test - fewer than 5 elements found.") - - await asyncio.sleep(2) - - print("\n✅ Chess Page Test Completed!") - await asyncio.sleep(100) - - except Exception as e: - print(f"\n❌ Chess Page Test failed: {str(e)}") - traceback.print_exc() - finally: - # Ensure browser is closed - print("\n--- Cleaning up ---") - await automation_service.shutdown() - print("Browser closed") - -if __name__ == '__main__': - import uvicorn - import sys - - # Check command line arguments for test mode - test_mode_1 = "--test" in sys.argv - test_mode_2 = "--test2" in sys.argv - - if test_mode_1: - print("Running in test mode 1") - asyncio.run(test_browser_api()) - elif test_mode_2: - print("Running in test mode 2 (Chess Page)") - asyncio.run(test_browser_api_2()) - else: - print("Starting API server") - uvicorn.run("browser_api:api_app", host="0.0.0.0", port=8003) \ No newline at end of file diff --git a/app/daytona/docker/docker-compose.yml b/app/daytona/docker/docker-compose.yml deleted file mode 100644 index 55ae725..0000000 --- a/app/daytona/docker/docker-compose.yml +++ /dev/null @@ -1,45 +0,0 @@ -services: - kortix-suna: - platform: linux/amd64 - build: - context: . - dockerfile: ${DOCKERFILE:-Dockerfile} - args: - TARGETPLATFORM: ${TARGETPLATFORM:-linux/amd64} - image: daytona_sanbox:1.0 - ports: - - "6080:6080" # noVNC web interface - - "5901:5901" # VNC port - - "9222:9222" # Chrome remote debugging port - - "8003:8003" # API server port - - "8080:8080" # HTTP server port - environment: - - ANONYMIZED_TELEMETRY=${ANONYMIZED_TELEMETRY:-false} - - CHROME_PATH=/usr/bin/google-chrome - - CHROME_USER_DATA=/app/data/chrome_data - - CHROME_PERSISTENT_SESSION=${CHROME_PERSISTENT_SESSION:-false} - - CHROME_CDP=${CHROME_CDP:-http://localhost:9222} - - DISPLAY=:99 - - PLAYWRIGHT_BROWSERS_PATH=/ms-playwright - - RESOLUTION=${RESOLUTION:-1024x768x24} - - RESOLUTION_WIDTH=${RESOLUTION_WIDTH:-1024} - - RESOLUTION_HEIGHT=${RESOLUTION_HEIGHT:-768} - - VNC_PASSWORD=${VNC_PASSWORD:-vncpassword} - - CHROME_DEBUGGING_PORT=9222 - - CHROME_DEBUGGING_HOST=localhost - - CHROME_FLAGS=${CHROME_FLAGS:-"--single-process --no-first-run --no-default-browser-check --disable-background-networking --disable-background-timer-throttling --disable-backgrounding-occluded-windows --disable-breakpad --disable-component-extensions-with-background-pages --disable-dev-shm-usage --disable-extensions --disable-features=TranslateUI --disable-ipc-flooding-protection --disable-renderer-backgrounding --enable-features=NetworkServiceInProcess2 --force-color-profile=srgb --metrics-recording-only --mute-audio --no-sandbox --disable-gpu"} - volumes: - - /tmp/.X11-unix:/tmp/.X11-unix - restart: unless-stopped - shm_size: '2gb' - cap_add: - - SYS_ADMIN - security_opt: - - seccomp=unconfined - tmpfs: - - /tmp - healthcheck: - test: ["CMD", "nc", "-z", "localhost", "5901"] - interval: 10s - timeout: 5s - retries: 3 diff --git a/app/daytona/docker/entrypoint.sh b/app/daytona/docker/entrypoint.sh deleted file mode 100644 index 9ab9240..0000000 --- a/app/daytona/docker/entrypoint.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/bin/bash - -# Start supervisord in the foreground to properly manage child processes -exec /usr/bin/supervisord -n -c /etc/supervisor/conf.d/supervisord.conf \ No newline at end of file diff --git a/app/daytona/docker/requirements.txt b/app/daytona/docker/requirements.txt deleted file mode 100644 index fae0feb..0000000 --- a/app/daytona/docker/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -fastapi==0.115.12 -uvicorn==0.34.0 -pyautogui==0.9.54 -pillow==10.2.0 -pydantic==2.6.1 -pytesseract==0.3.13 \ No newline at end of file diff --git a/app/daytona/docker/server.py b/app/daytona/docker/server.py deleted file mode 100644 index defa5f0..0000000 --- a/app/daytona/docker/server.py +++ /dev/null @@ -1,29 +0,0 @@ -from fastapi import FastAPI, Request -from fastapi.staticfiles import StaticFiles -from starlette.middleware.base import BaseHTTPMiddleware -import uvicorn -import os - -# Ensure we're serving from the /workspace directory -workspace_dir = "/workspace" - -class WorkspaceDirMiddleware(BaseHTTPMiddleware): - async def dispatch(self, request: Request, call_next): - # Check if workspace directory exists and recreate if deleted - if not os.path.exists(workspace_dir): - print(f"Workspace directory {workspace_dir} not found, recreating...") - os.makedirs(workspace_dir, exist_ok=True) - return await call_next(request) - -app = FastAPI() -app.add_middleware(WorkspaceDirMiddleware) - -# Initial directory creation -os.makedirs(workspace_dir, exist_ok=True) -app.mount('/', StaticFiles(directory=workspace_dir, html=True), name='site') - -# This is needed for the import string approach with uvicorn -if __name__ == '__main__': - print(f"Starting server with auto-reload, serving files from: {workspace_dir}") - # Don't use reload directly in the run call - uvicorn.run("server:app", host="0.0.0.0", port=8080, reload=True) \ No newline at end of file diff --git a/app/daytona/docker/supervisord.conf b/app/daytona/docker/supervisord.conf deleted file mode 100644 index b55ceb1..0000000 --- a/app/daytona/docker/supervisord.conf +++ /dev/null @@ -1,94 +0,0 @@ -[supervisord] -user=root -nodaemon=true -logfile=/dev/stdout -logfile_maxbytes=0 -loglevel=debug - -[program:xvfb] -command=Xvfb :99 -screen 0 %(ENV_RESOLUTION)s -ac +extension GLX +render -noreset -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=100 -startsecs=3 -stopsignal=TERM -stopwaitsecs=10 - -[program:vnc_setup] -command=bash -c "mkdir -p ~/.vnc && echo '%(ENV_VNC_PASSWORD)s' | vncpasswd -f > ~/.vnc/passwd && chmod 600 ~/.vnc/passwd && ls -la ~/.vnc/passwd" -autorestart=false -startsecs=0 -priority=150 -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 - -[program:x11vnc] -command=bash -c "mkdir -p /var/log && touch /var/log/x11vnc.log && chmod 666 /var/log/x11vnc.log && sleep 5 && DISPLAY=:99 x11vnc -display :99 -forever -shared -rfbauth /root/.vnc/passwd -rfbport 5901 -o /var/log/x11vnc.log" -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=200 -startretries=10 -startsecs=10 -stopsignal=TERM -stopwaitsecs=10 -depends_on=vnc_setup,xvfb - -[program:x11vnc_log] -command=bash -c "mkdir -p /var/log && touch /var/log/x11vnc.log && tail -f /var/log/x11vnc.log" -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=250 -stopsignal=TERM -stopwaitsecs=5 -depends_on=x11vnc - -[program:novnc] -command=bash -c "sleep 5 && cd /opt/novnc && ./utils/novnc_proxy --vnc localhost:5901 --listen 0.0.0.0:6080 --web /opt/novnc" -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=300 -startretries=5 -startsecs=3 -depends_on=x11vnc - -[program:http_server] -command=python /app/server.py -directory=/app -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=400 -startretries=5 -startsecs=5 -stopsignal=TERM -stopwaitsecs=10 - -[program:browser_api] -command=python /app/browser_api.py -directory=/app -autorestart=true -stdout_logfile=/dev/stdout -stdout_logfile_maxbytes=0 -stderr_logfile=/dev/stderr -stderr_logfile_maxbytes=0 -priority=400 -startretries=5 -startsecs=5 -stopsignal=TERM -stopwaitsecs=10 diff --git a/app/daytona/tool_base.py b/app/daytona/tool_base.py index 6a69b06..58fa629 100644 --- a/app/daytona/tool_base.py +++ b/app/daytona/tool_base.py @@ -12,8 +12,6 @@ from app.daytona.sandbox import ( get_or_start_sandbox, start_supervisord_session, ) - -# from app.agentpress.thread_manager import ThreadManager from app.tool.base import BaseTool, ToolResult from app.utils.files_utils import clean_path from app.utils.logger import logger diff --git a/app/services/billing.py b/app/services/billing.py deleted file mode 100644 index 67fa9bb..0000000 --- a/app/services/billing.py +++ /dev/null @@ -1,969 +0,0 @@ -""" -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 deleted file mode 100644 index b8b4180..0000000 --- a/app/services/docker/redis.conf +++ /dev/null @@ -1 +0,0 @@ -timeout 120 diff --git a/app/services/email.py b/app/services/email.py deleted file mode 100644 index 76dcb47..0000000 --- a/app/services/email.py +++ /dev/null @@ -1,192 +0,0 @@ -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 deleted file mode 100644 index 9834c7b..0000000 --- a/app/services/email_api.py +++ /dev/null @@ -1,70 +0,0 @@ -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 deleted file mode 100644 index cf624bf..0000000 --- a/app/services/langfuse.py +++ /dev/null @@ -1,12 +0,0 @@ -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 deleted file mode 100644 index f525f24..0000000 --- a/app/services/llm.py +++ /dev/null @@ -1,411 +0,0 @@ -""" -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 deleted file mode 100644 index 1a9c245..0000000 --- a/app/services/mcp_custom.py +++ /dev/null @@ -1,129 +0,0 @@ -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 deleted file mode 100644 index 71d3c01..0000000 --- a/app/services/mcp_temp.py +++ /dev/null @@ -1,299 +0,0 @@ -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 deleted file mode 100644 index 1259879..0000000 --- a/app/services/redis.py +++ /dev/null @@ -1,153 +0,0 @@ -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 deleted file mode 100644 index 0a3f855..0000000 --- a/app/services/supabase.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -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 deleted file mode 100644 index 2a40eec..0000000 --- a/app/services/transcription.py +++ /dev/null @@ -1,76 +0,0 @@ -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/sb_files_tool.py b/app/tool/sb_files_tool.py index f603814..405b5cd 100644 --- a/app/tool/sb_files_tool.py +++ b/app/tool/sb_files_tool.py @@ -21,6 +21,7 @@ Key capabilities include: * File reading: Read file contents with optional line range specification """ + class SandboxFilesTool(SandboxToolsBase): name: str = "sandbox_files" description: str = _FILES_DESCRIPTION @@ -33,7 +34,7 @@ class SandboxFilesTool(SandboxToolsBase): "create_file", "str_replace", "full_file_rewrite", - "delete_file" + "delete_file", ], "description": "The file operation to perform", }, @@ -56,8 +57,8 @@ class SandboxFilesTool(SandboxToolsBase): "permissions": { "type": "string", "description": "File permissions in octal format (e.g., '644')", - "default": "644" - } + "default": "644", + }, }, "required": ["action"], "dependencies": { @@ -71,7 +72,9 @@ class SandboxFilesTool(SandboxToolsBase): # workspace_path: str = Field(default="/workspace", exclude=True) # sandbox: Optional[Sandbox] = Field(default=None, exclude=True) - def __init__(self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data): + def __init__( + self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data + ): """Initialize with optional sandbox and thread_id.""" super().__init__(**data) if sandbox is not None: @@ -115,7 +118,7 @@ class SandboxFilesTool(SandboxToolsBase): "content": content, "is_dir": file_info.is_dir, "size": file_info.size, - "modified": file_info.mod_time + "modified": file_info.mod_time, } except Exception as e: print(f"Error reading file {rel_path}: {e}") @@ -136,7 +139,7 @@ class SandboxFilesTool(SandboxToolsBase): old_str: Optional[str] = None, new_str: Optional[str] = None, permissions: Optional[str] = "644", - **kwargs + **kwargs, ) -> ToolResult: """ Execute a file operation in the sandbox environment. @@ -155,25 +158,37 @@ class SandboxFilesTool(SandboxToolsBase): # File creation if action == "create_file": if not file_path or not file_contents: - return self.fail_response("file_path and file_contents are required for create_file") - return await self._create_file(file_path, file_contents, permissions) + return self.fail_response( + "file_path and file_contents are required for create_file" + ) + return await self._create_file( + file_path, file_contents, permissions + ) # String replacement elif action == "str_replace": if not file_path or not old_str or not new_str: - return self.fail_response("file_path, old_str, and new_str are required for str_replace") + return self.fail_response( + "file_path, old_str, and new_str are required for str_replace" + ) return await self._str_replace(file_path, old_str, new_str) # Full file rewrite elif action == "full_file_rewrite": if not file_path or not file_contents: - return self.fail_response("file_path and file_contents are required for full_file_rewrite") - return await self._full_file_rewrite(file_path, file_contents, permissions) + return self.fail_response( + "file_path and file_contents are required for full_file_rewrite" + ) + return await self._full_file_rewrite( + file_path, file_contents, permissions + ) # File deletion elif action == "delete_file": if not file_path: - return self.fail_response("file_path is required for delete_file") + return self.fail_response( + "file_path is required for delete_file" + ) return await self._delete_file(file_path) else: @@ -183,7 +198,9 @@ class SandboxFilesTool(SandboxToolsBase): logger.error(f"Error executing file action: {e}") return self.fail_response(f"Error executing file action: {e}") - async def _create_file(self, file_path: str, file_contents: str, permissions: str = "644") -> ToolResult: + async def _create_file( + self, file_path: str, file_contents: str, permissions: str = "644" + ) -> ToolResult: """Create a new file with the provided contents""" try: # Ensure sandbox is initialized @@ -192,10 +209,12 @@ class SandboxFilesTool(SandboxToolsBase): 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 full_file_rewrite to modify existing files.") + return self.fail_response( + f"File '{file_path}' already exists. Use full_file_rewrite to modify existing files." + ) # Create parent directories if needed - parent_dir = '/'.join(full_path.split('/')[:-1]) + parent_dir = "/".join(full_path.split("/")[:-1]) if parent_dir: self.sandbox.fs.create_folder(parent_dir, "755") @@ -206,20 +225,28 @@ class SandboxFilesTool(SandboxToolsBase): 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': + if file_path.lower() == "index.html": try: website_link = self.sandbox.get_preview_link(8080) - website_url = website_link.url if hasattr(website_link, 'url') else str(website_link).split("url='")[1].split("'")[0] + website_url = ( + website_link.url + if hasattr(website_link, "url") + else str(website_link).split("url='")[1].split("'")[0] + ) message += f"\n\n[Auto-detected index.html - HTTP server available at: {website_url}]" 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)}") + 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)}") - async def _str_replace(self, file_path: str, old_str: str, new_str: str) -> ToolResult: + async def _str_replace( + self, file_path: str, old_str: str, new_str: str + ) -> ToolResult: """Replace specific text in a file""" try: # Ensure sandbox is initialized @@ -238,18 +265,24 @@ class SandboxFilesTool(SandboxToolsBase): 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") + 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') + 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]) + end_line = replacement_line + self.SNIPPET_LINES + new_str.count("\n") + snippet = "\n".join(new_content.split("\n")[start_line : end_line + 1]) message = f"Replacement successful." @@ -258,7 +291,9 @@ class SandboxFilesTool(SandboxToolsBase): except Exception as e: return self.fail_response(f"Error replacing string: {str(e)}") - async def _full_file_rewrite(self, file_path: str, file_contents: str, permissions: str = "644") -> ToolResult: + async def _full_file_rewrite( + self, file_path: str, file_contents: str, permissions: str = "644" + ) -> ToolResult: """Completely rewrite an existing file with new content""" try: # Ensure sandbox is initialized @@ -267,7 +302,9 @@ class SandboxFilesTool(SandboxToolsBase): 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.") + 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) @@ -275,14 +312,20 @@ class SandboxFilesTool(SandboxToolsBase): 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': + if file_path.lower() == "index.html": try: website_link = self.sandbox.get_preview_link(8080) - website_url = website_link.url if hasattr(website_link, 'url') else str(website_link).split("url='")[1].split("'")[0] + website_url = ( + website_link.url + if hasattr(website_link, "url") + else str(website_link).split("url='")[1].split("'")[0] + ) message += f"\n\n[Auto-detected index.html - HTTP server available at: {website_url}]" 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)}") + logger.warning( + f"Failed to get website URL for index.html: {str(e)}" + ) return self.success_response(message) except Exception as e: @@ -311,4 +354,6 @@ class SandboxFilesTool(SandboxToolsBase): @classmethod def create_with_context(cls, context: Context) -> "SandboxFilesTool[Context]": """Factory method to create a SandboxFilesTool with a specific context.""" - raise NotImplementedError("create_with_context not implemented for SandboxFilesTool") \ No newline at end of file + raise NotImplementedError( + "create_with_context not implemented for SandboxFilesTool" + ) diff --git a/app/tool/sb_shell_tool.py b/app/tool/sb_shell_tool.py index 135e7f5..cfd3a85 100644 --- a/app/tool/sb_shell_tool.py +++ b/app/tool/sb_shell_tool.py @@ -18,7 +18,9 @@ This tool is essential for running CLI tools, installing packages, and managing class SandboxShellTool(SandboxToolsBase): """Tool for executing tasks in a Daytona sandbox with browser-use capabilities. - Uses sessions for maintaining state between commands and provides comprehensive process management.""" + Uses sessions for maintaining state between commands and provides comprehensive process management. + """ + name: str = "sandbox_shell" description: str = _SHELL_DESCRIPTION parameters: dict = { @@ -30,54 +32,56 @@ class SandboxShellTool(SandboxToolsBase): "execute_command", "check_command_output", "terminate_command", - "list_commands" + "list_commands", ], "description": "The shell action to perform", }, "command": { "type": "string", "description": "The shell command to execute. Use this for running CLI tools, installing packages, " - "or system operations. Commands can be chained using &&, ||, and | operators." + "or system operations. Commands can be chained using &&, ||, and | operators.", }, "folder": { "type": "string", "description": "Optional relative path to a subdirectory of /workspace where the command should be " - "executed. Example: 'data/pdfs'" + "executed. Example: 'data/pdfs'", }, "session_name": { "type": "string", "description": "Optional name of the tmux session to use. Use named sessions for related commands " - "that need to maintain state. Defaults to a random session name.", + "that need to maintain state. Defaults to a random session name.", }, "blocking": { "type": "boolean", "description": "Whether to wait for the command to complete. Defaults to false for non-blocking " - "execution.", - "default": False + "execution.", + "default": False, }, "timeout": { "type": "integer", "description": "Optional timeout in seconds for blocking commands. Defaults to 60. Ignored for " - "non-blocking commands.", - "default": 60 + "non-blocking commands.", + "default": 60, }, "kill_session": { "type": "boolean", "description": "Whether to terminate the tmux session after checking. Set to true when you're done " - "with the command.", - "default": False - } + "with the command.", + "default": False, + }, }, "required": ["action"], "dependencies": { "execute_command": ["command"], "check_command_output": ["session_name"], "terminate_command": ["session_name"], - "list_commands": [] + "list_commands": [], }, } - def __init__(self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data): + def __init__( + self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data + ): """Initialize with optional sandbox and thread_id.""" super().__init__(**data) if sandbox is not None: @@ -112,35 +116,30 @@ class SandboxShellTool(SandboxToolsBase): # Execute command in session from app.daytona.sandbox import SessionExecuteRequest + req = SessionExecuteRequest( - command=command, - run_async=False, - cwd=self.workspace_path + command=command, run_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 + timeout=30, # Short timeout for utility commands ) logs = self.sandbox.process.get_session_command_logs( - session_id=session_id, - command_id=response.cmd_id + session_id=session_id, command_id=response.cmd_id ) - return { - "output": logs, - "exit_code": response.exit_code - } + return {"output": logs, "exit_code": response.exit_code} async def _execute_command( - self, - command: str, - folder: Optional[str] = None, - session_name: Optional[str] = None, - blocking: bool = False, - timeout: int = 60 + self, + command: str, + folder: Optional[str] = None, + session_name: Optional[str] = None, + blocking: bool = False, + timeout: int = 60, ) -> ToolResult: try: # Ensure sandbox is initialized @@ -149,7 +148,7 @@ class SandboxShellTool(SandboxToolsBase): # Set up working directory cwd = self.workspace_path if folder: - folder = folder.strip('/') + folder = folder.strip("/") cwd = f"{self.workspace_path}/{folder}" # Generate a session name if not provided @@ -158,19 +157,24 @@ class SandboxShellTool(SandboxToolsBase): # 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'") + 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}") + 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') + await self._execute_raw_command( + f'tmux send-keys -t {session_name} "{wrapped_command}" Enter' + ) if blocking: # For blocking execution, wait and capture output @@ -181,56 +185,76 @@ class SandboxShellTool(SandboxToolsBase): # 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'") + 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 -") + 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): + 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 -") + 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, - "cwd": cwd, - "completed": True - }) + return self.success_response( + { + "output": final_output, + "session_name": session_name, + "cwd": cwd, + "completed": True, + } + ) else: # For non-blocking, just return immediately - return self.success_response({ - "session_name": session_name, - "cwd": cwd, - "message": - f"Command sent to tmux session '{session_name}'. Use check_command_output to view results.", - "completed": False - }) + return self.success_response( + { + "session_name": session_name, + "cwd": cwd, + "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: try: - await self._execute_raw_command(f"tmux kill-session -t {session_name}") + await self._execute_raw_command( + f"tmux kill-session -t {session_name}" + ) except: pass return self.fail_response(f"Error executing command: {str(e)}") async def _check_command_output( - self, - session_name: str, - kill_session: bool = False + self, session_name: str, kill_session: bool = False ) -> ToolResult: try: # Ensure sandbox is initialized @@ -238,12 +262,17 @@ class SandboxShellTool(SandboxToolsBase): # Check if session exists check_result = await self._execute_raw_command( - f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'") + 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.") + 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_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 @@ -253,35 +282,37 @@ class SandboxShellTool(SandboxToolsBase): else: termination_status = "Session still running." - return self.success_response({ - "output": output, - "session_name": session_name, - "status": termination_status - }) + 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)}") - async def _terminate_command( - self, - session_name: str - ) -> ToolResult: + async def _terminate_command(self, session_name: str) -> ToolResult: 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'") + 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.") + 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." - }) + 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)}") @@ -292,41 +323,44 @@ class SandboxShellTool(SandboxToolsBase): await self._ensure_sandbox() # List all tmux sessions - result = await self._execute_raw_command("tmux list-sessions 2>/dev/null || echo 'No 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": [] - }) + return self.success_response( + {"message": "No active tmux sessions found.", "sessions": []} + ) # Parse session list sessions = [] - for line in output.split('\n'): + for line in output.split("\n"): if line.strip(): - parts = line.split(':') + parts = line.split(":") if parts: session_name = parts[0].strip() sessions.append(session_name) - return self.success_response({ - "message": f"Found {len(sessions)} active sessions.", - "sessions": sessions - }) + 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)}") async def execute( - self, - action: str, - command: str, - folder: Optional[str] = None, - session_name: Optional[str] = None, - blocking: bool = False, - timeout: int = 60, - kill_session: bool = False, + self, + action: str, + command: str, + folder: Optional[str] = None, + session_name: Optional[str] = None, + blocking: bool = False, + timeout: int = 60, + kill_session: bool = False, ) -> ToolResult: """ Execute a browser action in the sandbox environment. @@ -347,14 +381,20 @@ class SandboxShellTool(SandboxToolsBase): if action == "execute_command": if not command: return self.fail_response("command is required for navigation") - return await self._execute_command(command, folder, session_name, blocking, timeout) + return await self._execute_command( + command, folder, session_name, blocking, timeout + ) elif action == "check_command_output": if session_name is None: - return self.fail_response("session_name is required for navigation") + return self.fail_response( + "session_name is required for navigation" + ) return await self._check_command_output(session_name, kill_session) elif action == "terminate_command": if session_name is None: - return self.fail_response("session_name is required for click_element") + return self.fail_response( + "session_name is required for click_element" + ) return await self._terminate_command(session_name) elif action == "list_commands": return await self._list_commands() @@ -364,7 +404,6 @@ class SandboxShellTool(SandboxToolsBase): logger.error(f"Error executing shell action: {e}") return self.fail_response(f"Error executing shell action: {e}") - async def cleanup(self): """Clean up all sessions.""" for session_name in list(self._sessions.keys()): diff --git a/app/tool/sb_vision_tool.py b/app/tool/sb_vision_tool.py index a7b9022..e550626 100644 --- a/app/tool/sb_vision_tool.py +++ b/app/tool/sb_vision_tool.py @@ -26,6 +26,7 @@ A sandbox-based vision tool that allows the agent to read image files inside the * Supported formats: JPG, PNG, GIF, WEBP. Maximum size: 10MB. """ + class SandboxVisionTool(SandboxToolsBase): name: str = "sandbox_vision" description: str = _VISION_DESCRIPTION @@ -35,17 +36,15 @@ class SandboxVisionTool(SandboxToolsBase): "action": { "type": "string", "enum": ["see_image"], - "description": "要执行的视觉动作，目前仅支持 see_image" + "description": "要执行的视觉动作，目前仅支持 see_image", }, "file_path": { "type": "string", - "description": "图片在 /workspace 下的相对路径，如 'screenshots/image.png'" - } + "description": "图片在 /workspace 下的相对路径，如 'screenshots/image.png'", + }, }, "required": ["action", "file_path"], - "dependencies": { - "see_image": ["file_path"] - } + "dependencies": {"see_image": ["file_path"]}, } # def __init__(self, project_id: str, thread_id: str, thread_manager: ThreadManager): @@ -54,7 +53,10 @@ class SandboxVisionTool(SandboxToolsBase): # self.thread_manager = thread_manager vision_message: Optional[ThreadMessage] = Field(default=None, exclude=True) - def __init__(self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data): + + def __init__( + self, sandbox: Optional[Sandbox] = None, thread_id: Optional[str] = None, **data + ): """Initialize with optional sandbox and thread_id.""" super().__init__(**data) if sandbox is not None: @@ -64,11 +66,13 @@ class SandboxVisionTool(SandboxToolsBase): """压缩图片，保持合理质量。""" try: img = Image.open(BytesIO(image_bytes)) - if img.mode in ('RGBA', 'LA', 'P'): - background = Image.new('RGB', img.size, (255, 255, 255)) - if img.mode == 'P': - img = img.convert('RGBA') - background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None) + if img.mode in ("RGBA", "LA", "P"): + background = Image.new("RGB", img.size, (255, 255, 255)) + if img.mode == "P": + img = img.convert("RGBA") + background.paste( + img, mask=img.split()[-1] if img.mode == "RGBA" else None + ) img = background width, height = img.size if width > DEFAULT_MAX_WIDTH or height > DEFAULT_MAX_HEIGHT: @@ -77,21 +81,30 @@ class SandboxVisionTool(SandboxToolsBase): new_height = int(height * ratio) img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) output = BytesIO() - if mime_type == 'image/gif': - img.save(output, format='GIF', optimize=True) - output_mime = 'image/gif' - elif mime_type == 'image/png': - img.save(output, format='PNG', optimize=True, compress_level=DEFAULT_PNG_COMPRESS_LEVEL) - output_mime = 'image/png' + if mime_type == "image/gif": + img.save(output, format="GIF", optimize=True) + output_mime = "image/gif" + elif mime_type == "image/png": + img.save( + output, + format="PNG", + optimize=True, + compress_level=DEFAULT_PNG_COMPRESS_LEVEL, + ) + output_mime = "image/png" else: - img.save(output, format='JPEG', quality=DEFAULT_JPEG_QUALITY, optimize=True) - output_mime = 'image/jpeg' + img.save( + output, format="JPEG", quality=DEFAULT_JPEG_QUALITY, optimize=True + ) + output_mime = "image/jpeg" compressed_bytes = output.getvalue() return compressed_bytes, output_mime except Exception as e: return image_bytes, mime_type - async def execute(self, action: str, file_path: Optional[str] = None, **kwargs) -> ToolResult: + async def execute( + self, action: str, file_path: Optional[str] = None, **kwargs + ) -> ToolResult: """ 执行视觉动作，目前仅支持 see_image。 参数： @@ -109,42 +122,56 @@ class SandboxVisionTool(SandboxToolsBase): try: file_info = self.sandbox.fs.get_file_info(full_path) if file_info.is_dir: - return self.fail_response(f"路径 '{cleaned_path}' 是目录，不是图片文件。") + return self.fail_response( + f"路径 '{cleaned_path}' 是目录，不是图片文件。" + ) except Exception: return self.fail_response(f"图片文件未找到: '{cleaned_path}'") if file_info.size > MAX_IMAGE_SIZE: - return self.fail_response(f"图片文件 '{cleaned_path}' 过大 ({file_info.size / (1024*1024):.2f}MB)，最大允许 {MAX_IMAGE_SIZE / (1024*1024)}MB。") + return self.fail_response( + f"图片文件 '{cleaned_path}' 过大 ({file_info.size / (1024*1024):.2f}MB)，最大允许 {MAX_IMAGE_SIZE / (1024*1024)}MB。" + ) try: image_bytes = self.sandbox.fs.download_file(full_path) except Exception: return self.fail_response(f"无法读取图片文件: {cleaned_path}") mime_type, _ = mimetypes.guess_type(full_path) - if not mime_type or not mime_type.startswith('image/'): + if not mime_type or not mime_type.startswith("image/"): ext = os.path.splitext(cleaned_path)[1].lower() - if ext == '.jpg' or ext == '.jpeg': mime_type = 'image/jpeg' - elif ext == '.png': mime_type = 'image/png' - elif ext == '.gif': mime_type = 'image/gif' - elif ext == '.webp': mime_type = 'image/webp' + if ext == ".jpg" or ext == ".jpeg": + mime_type = "image/jpeg" + elif ext == ".png": + mime_type = "image/png" + elif ext == ".gif": + mime_type = "image/gif" + elif ext == ".webp": + mime_type = "image/webp" else: - return self.fail_response(f"不支持或未知的图片格式: '{cleaned_path}'。支持: JPG, PNG, GIF, WEBP。") - compressed_bytes, compressed_mime_type = self.compress_image(image_bytes, mime_type, cleaned_path) + return self.fail_response( + f"不支持或未知的图片格式: '{cleaned_path}'。支持: JPG, PNG, GIF, WEBP。" + ) + compressed_bytes, compressed_mime_type = self.compress_image( + image_bytes, mime_type, cleaned_path + ) if len(compressed_bytes) > MAX_COMPRESSED_SIZE: - return self.fail_response(f"图片文件 '{cleaned_path}' 压缩后仍过大 ({len(compressed_bytes) / (1024*1024):.2f}MB)，最大允许 {MAX_COMPRESSED_SIZE / (1024*1024)}MB。") - base64_image = base64.b64encode(compressed_bytes).decode('utf-8') + return self.fail_response( + f"图片文件 '{cleaned_path}' 压缩后仍过大 ({len(compressed_bytes) / (1024*1024):.2f}MB)，最大允许 {MAX_COMPRESSED_SIZE / (1024*1024)}MB。" + ) + base64_image = base64.b64encode(compressed_bytes).decode("utf-8") image_context_data = { "mime_type": compressed_mime_type, "base64": base64_image, "file_path": cleaned_path, "original_size": file_info.size, - "compressed_size": len(compressed_bytes) + "compressed_size": len(compressed_bytes), } message = ThreadMessage( - type="image_context", - content=image_context_data, - is_llm_message=False + type="image_context", content=image_context_data, is_llm_message=False ) self.vision_message = message # return self.success_response(f"成功加载并压缩图片 '{cleaned_path}' (由 {file_info.size / 1024:.1f}KB 压缩到 {len(compressed_bytes) / 1024:.1f}KB)。") - return self.success_response(f"成功加载并压缩图片 '{cleaned_path}'，压缩后的内容为：{base64_image}") + return self.success_response( + f"成功加载并压缩图片 '{cleaned_path}'，压缩后的内容为：{base64_image}" + ) except Exception as e: - return self.fail_response(f"see_image 执行异常: {str(e)}") \ No newline at end of file + return self.fail_response(f"see_image 执行异常: {str(e)}") diff --git a/app/utils/auth_utils.py b/app/utils/auth_utils.py deleted file mode 100644 index 62de820..0000000 --- a/app/utils/auth_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -import sentry -from fastapi import HTTPException, Request -from typing import Optional -import jwt -from jwt.exceptions import PyJWTError -from utils.logger import structlog - -# This function extracts the user ID from Supabase JWT -async def get_current_user_id_from_jwt(request: Request) -> str: - """ - Extract and verify the user ID from the JWT in the Authorization header. - - This function is used as a dependency in FastAPI routes to ensure the user - is authenticated and to provide the user ID for authorization checks. - - Args: - request: The FastAPI request object - - Returns: - str: The user ID extracted from the JWT - - Raises: - HTTPException: If no valid token is found or if the token is invalid - """ - auth_header = request.headers.get('Authorization') - - if not auth_header or not auth_header.startswith('Bearer '): - raise HTTPException( - status_code=401, - detail="No valid authentication credentials found", - headers={"WWW-Authenticate": "Bearer"} - ) - - token = auth_header.split(' ')[1] - - try: - # For Supabase JWT, we just need to decode and extract the user ID - # The actual validation is handled by Supabase's RLS - payload = jwt.decode(token, options={"verify_signature": False}) - - # Supabase stores the user ID in the 'sub' claim - user_id = payload.get('sub') - - if not user_id: - raise HTTPException( - status_code=401, - detail="Invalid token payload", - headers={"WWW-Authenticate": "Bearer"} - ) - - sentry.sentry.set_user({ "id": user_id }) - structlog.contextvars.bind_contextvars( - user_id=user_id - ) - return user_id - - except PyJWTError: - raise HTTPException( - status_code=401, - detail="Invalid token", - headers={"WWW-Authenticate": "Bearer"} - ) - -async def get_account_id_from_thread(client, thread_id: str) -> str: - """ - Extract and verify the account ID from the thread. - - Args: - client: The Supabase client - thread_id: The ID of the thread - - Returns: - str: The account ID associated with the thread - - Raises: - HTTPException: If the thread is not found or if there's an error - """ - try: - response = await client.table('threads').select('account_id').eq('thread_id', thread_id).execute() - - if not response.data or len(response.data) == 0: - raise HTTPException( - status_code=404, - detail="Thread not found" - ) - - account_id = response.data[0].get('account_id') - - if not account_id: - raise HTTPException( - status_code=500, - detail="Thread has no associated account" - ) - - return account_id - - except Exception as e: - raise HTTPException( - status_code=500, - detail=f"Error retrieving thread information: {str(e)}" - ) - -async def get_user_id_from_stream_auth( - request: Request, - token: Optional[str] = None -) -> str: - """ - Extract and verify the user ID from either the Authorization header or query parameter token. - This function is specifically designed for streaming endpoints that need to support both - header-based and query parameter-based authentication (for EventSource compatibility). - - Args: - request: The FastAPI request object - token: Optional token from query parameters - - Returns: - str: The user ID extracted from the JWT - - Raises: - HTTPException: If no valid token is found or if the token is invalid - """ - # Try to get user_id from token in query param (for EventSource which can't set headers) - if token: - try: - # For Supabase JWT, we just need to decode and extract the user ID - payload = jwt.decode(token, options={"verify_signature": False}) - user_id = payload.get('sub') - if user_id: - sentry.sentry.set_user({ "id": user_id }) - structlog.contextvars.bind_contextvars( - user_id=user_id - ) - return user_id - except Exception: - pass - - # If no valid token in query param, try to get it from the Authorization header - auth_header = request.headers.get('Authorization') - if auth_header and auth_header.startswith('Bearer '): - try: - # Extract token from header - header_token = auth_header.split(' ')[1] - payload = jwt.decode(header_token, options={"verify_signature": False}) - user_id = payload.get('sub') - if user_id: - return user_id - except Exception: - pass - - # If we still don't have a user_id, return authentication error - raise HTTPException( - status_code=401, - detail="No valid authentication credentials found", - headers={"WWW-Authenticate": "Bearer"} - ) - -async def verify_thread_access(client, thread_id: str, user_id: str): - """ - Verify that a user has access to a specific thread based on account membership. - - Args: - client: The Supabase client - thread_id: The thread ID to check access for - user_id: The user ID to check permissions for - - Returns: - bool: True if the user has access - - Raises: - HTTPException: If the user doesn't have access to the thread - """ - # Query the thread to get account information - thread_result = await client.table('threads').select('*,project_id').eq('thread_id', thread_id).execute() - - if not thread_result.data or len(thread_result.data) == 0: - raise HTTPException(status_code=404, detail="Thread not found") - - thread_data = thread_result.data[0] - - # Check if project is public - project_id = thread_data.get('project_id') - if project_id: - project_result = await client.table('projects').select('is_public').eq('project_id', project_id).execute() - if project_result.data and len(project_result.data) > 0: - if project_result.data[0].get('is_public'): - return True - - account_id = thread_data.get('account_id') - # When using service role, we need to manually check account membership instead of using current_user_account_role - if account_id: - account_user_result = await client.schema('basejump').from_('account_user').select('account_role').eq('user_id', user_id).eq('account_id', account_id).execute() - if account_user_result.data and len(account_user_result.data) > 0: - return True - raise HTTPException(status_code=403, detail="Not authorized to access this thread") - -async def get_optional_user_id(request: Request) -> Optional[str]: - """ - Extract the user ID from the JWT in the Authorization header if present, - but don't require authentication. Returns None if no valid token is found. - - This function is used for endpoints that support both authenticated and - unauthenticated access (like public projects). - - Args: - request: The FastAPI request object - - Returns: - Optional[str]: The user ID extracted from the JWT, or None if no valid token - """ - auth_header = request.headers.get('Authorization') - - if not auth_header or not auth_header.startswith('Bearer '): - return None - - token = auth_header.split(' ')[1] - - try: - # For Supabase JWT, we just need to decode and extract the user ID - payload = jwt.decode(token, options={"verify_signature": False}) - - # Supabase stores the user ID in the 'sub' claim - user_id = payload.get('sub') - if user_id: - sentry.sentry.set_user({ "id": user_id }) - structlog.contextvars.bind_contextvars( - user_id=user_id - ) - - return user_id - except PyJWTError: - return None diff --git a/app/utils/config.py b/app/utils/config.py deleted file mode 100644 index ab28537..0000000 --- a/app/utils/config.py +++ /dev/null @@ -1,253 +0,0 @@ -""" -Configuration management. - -This module provides a centralized way to access configuration settings and -environment variables across the application. It supports different environment -modes (development, staging, production) and provides validation for required -values. - -Usage: - from utils.config import config - - # Access configuration values - api_key = config.OPENAI_API_KEY - env_mode = config.ENV_MODE -""" - -import os -from enum import Enum -from typing import Dict, Any, Optional, get_type_hints, Union -from dotenv import load_dotenv -import logging - -logger = logging.getLogger(__name__) - -class EnvMode(Enum): - """Environment mode enumeration.""" - LOCAL = "local" - STAGING = "staging" - PRODUCTION = "production" - -class Configuration: - """ - Centralized configuration for AgentPress backend. - - This class loads environment variables and provides type checking and validation. - Default values can be specified for optional configuration items. - """ - - # Environment mode - ENV_MODE: EnvMode = EnvMode.LOCAL - - # Subscription tier IDs - Production - STRIPE_FREE_TIER_ID_PROD: str = 'price_1RILb4G6l1KZGqIrK4QLrx9i' - STRIPE_TIER_2_20_ID_PROD: str = 'price_1RILb4G6l1KZGqIrhomjgDnO' - STRIPE_TIER_6_50_ID_PROD: str = 'price_1RILb4G6l1KZGqIr5q0sybWn' - STRIPE_TIER_12_100_ID_PROD: str = 'price_1RILb4G6l1KZGqIr5Y20ZLHm' - STRIPE_TIER_25_200_ID_PROD: str = 'price_1RILb4G6l1KZGqIrGAD8rNjb' - STRIPE_TIER_50_400_ID_PROD: str = 'price_1RILb4G6l1KZGqIruNBUMTF1' - STRIPE_TIER_125_800_ID_PROD: str = 'price_1RILb3G6l1KZGqIrbJA766tN' - STRIPE_TIER_200_1000_ID_PROD: str = 'price_1RILb3G6l1KZGqIrmauYPOiN' - - # Subscription tier IDs - Staging - STRIPE_FREE_TIER_ID_STAGING: str = 'price_1RIGvuG6l1KZGqIrw14abxeL' - STRIPE_TIER_2_20_ID_STAGING: str = 'price_1RIGvuG6l1KZGqIrCRu0E4Gi' - STRIPE_TIER_6_50_ID_STAGING: str = 'price_1RIGvuG6l1KZGqIrvjlz5p5V' - STRIPE_TIER_12_100_ID_STAGING: str = 'price_1RIGvuG6l1KZGqIrT6UfgblC' - STRIPE_TIER_25_200_ID_STAGING: str = 'price_1RIGvuG6l1KZGqIrOVLKlOMj' - STRIPE_TIER_50_400_ID_STAGING: str = 'price_1RIKNgG6l1KZGqIrvsat5PW7' - STRIPE_TIER_125_800_ID_STAGING: str = 'price_1RIKNrG6l1KZGqIrjKT0yGvI' - STRIPE_TIER_200_1000_ID_STAGING: str = 'price_1RIKQ2G6l1KZGqIrum9n8SI7' - - # Computed subscription tier IDs based on environment - @property - def STRIPE_FREE_TIER_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_FREE_TIER_ID_STAGING - return self.STRIPE_FREE_TIER_ID_PROD - - @property - def STRIPE_TIER_2_20_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_2_20_ID_STAGING - return self.STRIPE_TIER_2_20_ID_PROD - - @property - def STRIPE_TIER_6_50_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_6_50_ID_STAGING - return self.STRIPE_TIER_6_50_ID_PROD - - @property - def STRIPE_TIER_12_100_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_12_100_ID_STAGING - return self.STRIPE_TIER_12_100_ID_PROD - - @property - def STRIPE_TIER_25_200_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_25_200_ID_STAGING - return self.STRIPE_TIER_25_200_ID_PROD - - @property - def STRIPE_TIER_50_400_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_50_400_ID_STAGING - return self.STRIPE_TIER_50_400_ID_PROD - - @property - def STRIPE_TIER_125_800_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_125_800_ID_STAGING - return self.STRIPE_TIER_125_800_ID_PROD - - @property - def STRIPE_TIER_200_1000_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_TIER_200_1000_ID_STAGING - return self.STRIPE_TIER_200_1000_ID_PROD - - # LLM API keys - ANTHROPIC_API_KEY: Optional[str] = None - OPENAI_API_KEY: Optional[str] = None - GROQ_API_KEY: Optional[str] = None - OPENROUTER_API_KEY: Optional[str] = None - OPENROUTER_API_BASE: Optional[str] = "https://openrouter.ai/api/v1" - OR_SITE_URL: Optional[str] = "https://kortix.ai" - OR_APP_NAME: Optional[str] = "Kortix AI" - - # AWS Bedrock credentials - AWS_ACCESS_KEY_ID: Optional[str] = None - AWS_SECRET_ACCESS_KEY: Optional[str] = None - AWS_REGION_NAME: Optional[str] = None - - # Model configuration - MODEL_TO_USE: Optional[str] = "anthropic/claude-sonnet-4-20250514" - - # Supabase configuration - SUPABASE_URL: str - SUPABASE_ANON_KEY: str - SUPABASE_SERVICE_ROLE_KEY: str - - # Redis configuration - REDIS_HOST: str - REDIS_PORT: int = 6379 - REDIS_PASSWORD: Optional[str] = None - REDIS_SSL: bool = True - - # Daytona sandbox configuration - DAYTONA_API_KEY: str - DAYTONA_SERVER_URL: str - DAYTONA_TARGET: str - - # Search and other API keys - TAVILY_API_KEY: str - RAPID_API_KEY: str - CLOUDFLARE_API_TOKEN: Optional[str] = None - FIRECRAWL_API_KEY: str - FIRECRAWL_URL: Optional[str] = "https://api.firecrawl.dev" - - # Stripe configuration - STRIPE_SECRET_KEY: Optional[str] = None - STRIPE_WEBHOOK_SECRET: Optional[str] = None - STRIPE_DEFAULT_PLAN_ID: Optional[str] = None - STRIPE_DEFAULT_TRIAL_DAYS: int = 14 - - # Stripe Product IDs - STRIPE_PRODUCT_ID_PROD: str = 'prod_SCl7AQ2C8kK1CD' - STRIPE_PRODUCT_ID_STAGING: str = 'prod_SCgIj3G7yPOAWY' - - # Sandbox configuration - SANDBOX_IMAGE_NAME = "kortix/suna:0.1.3" - SANDBOX_ENTRYPOINT = "/usr/bin/supervisord -n -c /etc/supervisor/conf.d/supervisord.conf" - - # LangFuse configuration - LANGFUSE_PUBLIC_KEY: Optional[str] = None - LANGFUSE_SECRET_KEY: Optional[str] = None - LANGFUSE_HOST: str = "https://cloud.langfuse.com" - - @property - def STRIPE_PRODUCT_ID(self) -> str: - if self.ENV_MODE == EnvMode.STAGING: - return self.STRIPE_PRODUCT_ID_STAGING - return self.STRIPE_PRODUCT_ID_PROD - - def __init__(self): - """Initialize configuration by loading from environment variables.""" - # Load environment variables from .env file if it exists - load_dotenv() - - # Set environment mode first - env_mode_str = os.getenv("ENV_MODE", EnvMode.LOCAL.value) - try: - self.ENV_MODE = EnvMode(env_mode_str.lower()) - except ValueError: - logger.warning(f"Invalid ENV_MODE: {env_mode_str}, defaulting to LOCAL") - self.ENV_MODE = EnvMode.LOCAL - - logger.info(f"Environment mode: {self.ENV_MODE.value}") - - # Load configuration from environment variables - self._load_from_env() - - # Perform validation - self._validate() - - def _load_from_env(self): - """Load configuration values from environment variables.""" - for key, expected_type in get_type_hints(self.__class__).items(): - env_val = os.getenv(key) - - if env_val is not None: - # Convert environment variable to the expected type - if expected_type == bool: - # Handle boolean conversion - setattr(self, key, env_val.lower() in ('true', 't', 'yes', 'y', '1')) - elif expected_type == int: - # Handle integer conversion - try: - setattr(self, key, int(env_val)) - except ValueError: - logger.warning(f"Invalid value for {key}: {env_val}, using default") - elif expected_type == EnvMode: - # Already handled for ENV_MODE - pass - else: - # String or other type - setattr(self, key, env_val) - - def _validate(self): - """Validate configuration based on type hints.""" - # Get all configuration fields and their type hints - type_hints = get_type_hints(self.__class__) - - # Find missing required fields - missing_fields = [] - for field, field_type in type_hints.items(): - # Check if the field is Optional - is_optional = hasattr(field_type, "__origin__") and field_type.__origin__ is Union and type(None) in field_type.__args__ - - # If not optional and value is None, add to missing fields - if not is_optional and getattr(self, field) is None: - missing_fields.append(field) - - if missing_fields: - error_msg = f"Missing required configuration fields: {', '.join(missing_fields)}" - logger.error(error_msg) - raise ValueError(error_msg) - - def get(self, key: str, default: Any = None) -> Any: - """Get a configuration value with an optional default.""" - return getattr(self, key, default) - - def as_dict(self) -> Dict[str, Any]: - """Return configuration as a dictionary.""" - return { - key: getattr(self, key) - for key in get_type_hints(self.__class__).keys() - if not key.startswith('_') - } - -# Create a singleton instance -config = Configuration() \ No newline at end of file diff --git a/app/utils/constants.py b/app/utils/constants.py deleted file mode 100644 index e24578b..0000000 --- a/app/utils/constants.py +++ /dev/null @@ -1,145 +0,0 @@ -MODEL_ACCESS_TIERS = { - "free": [ - "openrouter/deepseek/deepseek-chat", - "openrouter/qwen/qwen3-235b-a22b", - "openrouter/google/gemini-2.5-flash-preview-05-20", - ], - "tier_2_20": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_6_50": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_12_100": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_25_200": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_50_400": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_125_800": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], - "tier_200_1000": [ - "openrouter/deepseek/deepseek-chat", - # "xai/grok-3-mini-fast-beta", - "openai/gpt-4o", - # "openai/gpt-4-turbo", - # "xai/grok-3-fast-latest", - "openrouter/google/gemini-2.5-flash-preview-05-20", # Added - # "openai/gpt-4", - "anthropic/claude-3-7-sonnet-latest", - "anthropic/claude-sonnet-4-20250514", - # "openai/gpt-4.1-2025-04-14", - # "openrouter/deepseek/deepseek-r1", - "openrouter/qwen/qwen3-235b-a22b", - ], -} -MODEL_NAME_ALIASES = { - # Short names to full names - "sonnet-3.7": "anthropic/claude-3-7-sonnet-latest", - "sonnet-3.5": "anthropic/claude-3-5-sonnet-latest", - "haiku-3.5": "anthropic/claude-3-5-haiku-latest", - "claude-sonnet-4": "anthropic/claude-sonnet-4-20250514", - # "gpt-4.1": "openai/gpt-4.1-2025-04-14", # Commented out in constants.py - "gpt-4o": "openai/gpt-4o", - "gpt-4.1": "openai/gpt-4.1", - "gpt-4.1-mini": "gpt-4.1-mini", - # "gpt-4-turbo": "openai/gpt-4-turbo", # Commented out in constants.py - # "gpt-4": "openai/gpt-4", # Commented out in constants.py - # "gemini-flash-2.5": "openrouter/google/gemini-2.5-flash-preview", # Commented out in constants.py - # "grok-3": "xai/grok-3-fast-latest", # Commented out in constants.py - "deepseek": "openrouter/deepseek/deepseek-chat", - # "deepseek-r1": "openrouter/deepseek/deepseek-r1", - # "grok-3-mini": "xai/grok-3-mini-fast-beta", # Commented out in constants.py - "qwen3": "openrouter/qwen/qwen3-235b-a22b", # Commented out in constants.py - "gemini-flash-2.5": "openrouter/google/gemini-2.5-flash-preview-05-20", - "gemini-2.5-flash:thinking":"openrouter/google/gemini-2.5-flash-preview-05-20:thinking", - - # "google/gemini-2.5-flash-preview":"openrouter/google/gemini-2.5-flash-preview", - # "google/gemini-2.5-flash-preview:thinking":"openrouter/google/gemini-2.5-flash-preview:thinking", - "google/gemini-2.5-pro-preview":"openrouter/google/gemini-2.5-pro-preview", - "deepseek/deepseek-chat-v3-0324":"openrouter/deepseek/deepseek-chat-v3-0324", - - # Also include full names as keys to ensure they map to themselves - # "anthropic/claude-3-7-sonnet-latest": "anthropic/claude-3-7-sonnet-latest", - # "openai/gpt-4.1-2025-04-14": "openai/gpt-4.1-2025-04-14", # Commented out in constants.py - # "openai/gpt-4o": "openai/gpt-4o", - # "openai/gpt-4-turbo": "openai/gpt-4-turbo", # Commented out in constants.py - # "openai/gpt-4": "openai/gpt-4", # Commented out in constants.py - # "openrouter/google/gemini-2.5-flash-preview": "openrouter/google/gemini-2.5-flash-preview", # Commented out in constants.py - # "xai/grok-3-fast-latest": "xai/grok-3-fast-latest", # Commented out in constants.py - # "deepseek/deepseek-chat": "openrouter/deepseek/deepseek-chat", - # "deepseek/deepseek-r1": "openrouter/deepseek/deepseek-r1", - - # "qwen/qwen3-235b-a22b": "openrouter/qwen/qwen3-235b-a22b", - # "xai/grok-3-mini-fast-beta": "xai/grok-3-mini-fast-beta", # Commented out in constants.py -} \ No newline at end of file diff --git a/app/utils/retry.py b/app/utils/retry.py deleted file mode 100644 index 992c045..0000000 --- a/app/utils/retry.py +++ /dev/null @@ -1,58 +0,0 @@ -import asyncio -from typing import TypeVar, Callable, Awaitable, Optional - -T = TypeVar("T") - - -async def retry( - fn: Callable[[], Awaitable[T]], - max_attempts: int = 3, - delay_seconds: int = 1, -) -> T: - """ - Retry an async function with exponential backoff. - - Args: - fn: The async function to retry - max_attempts: Maximum number of attempts - delay_seconds: Delay between attempts in seconds - - Returns: - The result of the function call - - Raises: - The last exception if all attempts fail - - Example: - ```python - async def fetch_data(): - # Some operation that might fail - return await api_call() - - try: - result = await retry(fetch_data, max_attempts=3, delay_seconds=2) - print(f"Success: {result}") - except Exception as e: - print(f"Failed after all retries: {e}") - ``` - """ - if max_attempts <= 0: - raise ValueError("max_attempts must be greater than zero") - - last_error: Optional[Exception] = None - - for attempt in range(1, max_attempts + 1): - try: - return await fn() - except Exception as error: - last_error = error - - if attempt == max_attempts: - break - - await asyncio.sleep(delay_seconds) - - if last_error: - raise last_error - - raise RuntimeError("Unexpected: last_error is None") diff --git a/app/utils/s3_upload_utils.py b/app/utils/s3_upload_utils.py deleted file mode 100644 index 6572264..0000000 --- a/app/utils/s3_upload_utils.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Utility functions for handling image operations. -""" - -import base64 -import uuid -from datetime import datetime -from utils.logger import logger -from services.supabase import DBConnection - -async def upload_base64_image(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 - db = DBConnection() - client = await db.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)}") \ No newline at end of file diff --git a/app/utils/scripts/archive_inactive_sandboxes.py b/app/utils/scripts/archive_inactive_sandboxes.py deleted file mode 100644 index 1a56dbe..0000000 --- a/app/utils/scripts/archive_inactive_sandboxes.py +++ /dev/null @@ -1,350 +0,0 @@ -#!/usr/bin/env python -""" -Script to archive sandboxes for projects whose account_id is not associated with an active billing customer. - -Usage: - python archive_inactive_sandboxes.py - -This script: -1. Gets all active account_ids from basejump.billing_customers (active=TRUE) -2. Gets all projects from the projects table -3. Archives sandboxes for any project whose account_id is not in the active billing customers list - -Make sure your environment variables are properly set: -- SUPABASE_URL -- SUPABASE_SERVICE_ROLE_KEY -- DAYTONA_SERVER_URL -""" - -import asyncio -import sys -import os -import argparse -from typing import List, Dict, Any, Set -from dotenv import load_dotenv - -# Load script-specific environment variables -load_dotenv(".env") - -from services.supabase import DBConnection -from sandbox.sandbox import daytona -from utils.logger import logger - -# Global DB connection to reuse -db_connection = None - - -async def get_active_billing_customer_account_ids() -> Set[str]: - """ - Query all account_ids from the basejump.billing_customers table where active=TRUE. - - Returns: - Set of account_ids that have an active billing customer record - """ - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Print the Supabase URL being used - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - # Query all account_ids from billing_customers where active=true - result = await client.schema('basejump').from_('billing_customers').select('account_id, active').execute() - - # Print the query result - print(f"Found {len(result.data)} billing customers in database") - print(result.data) - - if not result.data: - logger.info("No billing customers found in database") - return set() - - # Extract account_ids for active customers and return as a set for fast lookups - active_account_ids = {customer.get('account_id') for customer in result.data - if customer.get('account_id') and customer.get('active') is True} - - print(f"Found {len(active_account_ids)} active billing customers") - return active_account_ids - - -async def get_all_projects() -> List[Dict[str, Any]]: - """ - Query all projects with sandbox information. - - Returns: - List of projects with their sandbox information - """ - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Initialize variables for pagination - all_projects = [] - page_size = 1000 - current_page = 0 - has_more = True - - logger.info("Starting to fetch all projects (paginated)") - - # Paginate through all projects - while has_more: - # Query projects with pagination - start_range = current_page * page_size - end_range = start_range + page_size - 1 - - logger.info(f"Fetching projects page {current_page+1} (range: {start_range}-{end_range})") - - result = await client.table('projects').select( - 'project_id', - 'name', - 'account_id', - 'sandbox' - ).range(start_range, end_range).execute() - - if not result.data: - has_more = False - else: - all_projects.extend(result.data) - current_page += 1 - - # Progress update - logger.info(f"Loaded {len(all_projects)} projects so far") - print(f"Loaded {len(all_projects)} projects so far...") - - # Check if we've reached the end - if len(result.data) < page_size: - has_more = False - - # Print the query result - total_projects = len(all_projects) - print(f"Found {total_projects} projects in database") - logger.info(f"Total projects found in database: {total_projects}") - - if not all_projects: - logger.info("No projects found in database") - return [] - - # Filter projects that have sandbox information - projects_with_sandboxes = [ - project for project in all_projects - if project.get('sandbox') and project['sandbox'].get('id') - ] - - logger.info(f"Found {len(projects_with_sandboxes)} projects with sandboxes") - return projects_with_sandboxes - - -async def archive_sandbox(project: Dict[str, Any], dry_run: bool) -> bool: - """ - Archive a single sandbox. - - Args: - project: Project information containing sandbox to archive - dry_run: If True, only simulate archiving - - Returns: - True if successful, False otherwise - """ - sandbox_id = project['sandbox'].get('id') - project_name = project.get('name', 'Unknown') - project_id = project.get('project_id', 'Unknown') - - try: - logger.info(f"Checking sandbox {sandbox_id} for project '{project_name}' (ID: {project_id})") - - if dry_run: - logger.info(f"DRY RUN: Would archive sandbox {sandbox_id}") - print(f"Would archive sandbox {sandbox_id} for project '{project_name}'") - return True - - # Get the sandbox - sandbox = daytona.get(sandbox_id) - - # Check sandbox state - it must be stopped before archiving - sandbox_info = sandbox.info() - - # Log the current state - logger.info(f"Sandbox {sandbox_id} is in '{sandbox_info.state}' state") - - # Only archive if the sandbox is in the stopped state - if sandbox_info.state == "stopped": - logger.info(f"Archiving sandbox {sandbox_id} as it is in stopped state") - sandbox.archive() - logger.info(f"Successfully archived sandbox {sandbox_id}") - return True - else: - logger.info(f"Skipping sandbox {sandbox_id} as it is not in stopped state (current: {sandbox_info.state})") - return True - - except Exception as e: - import traceback - error_type = type(e).__name__ - stack_trace = traceback.format_exc() - - # Log detailed error information - logger.error(f"Error processing sandbox {sandbox_id}: {str(e)}") - logger.error(f"Error type: {error_type}") - logger.error(f"Stack trace:\n{stack_trace}") - - # If the exception has a response attribute (like in HTTP errors), log it - if hasattr(e, 'response'): - try: - response_data = e.response.json() if hasattr(e.response, 'json') else str(e.response) - logger.error(f"Response data: {response_data}") - except Exception: - logger.error(f"Could not parse response data from error") - - print(f"Failed to process sandbox {sandbox_id}: {error_type} - {str(e)}") - return False - - -async def process_sandboxes(inactive_projects: List[Dict[str, Any]], dry_run: bool) -> tuple[int, int]: - """ - Process all sandboxes sequentially. - - Args: - inactive_projects: List of projects without active billing - dry_run: Whether to actually archive sandboxes or just simulate - - Returns: - Tuple of (processed_count, failed_count) - """ - processed_count = 0 - failed_count = 0 - - if dry_run: - logger.info(f"DRY RUN: Would archive {len(inactive_projects)} sandboxes") - else: - logger.info(f"Archiving {len(inactive_projects)} sandboxes") - - print(f"Processing {len(inactive_projects)} sandboxes...") - - # Process each sandbox sequentially - for i, project in enumerate(inactive_projects): - success = await archive_sandbox(project, dry_run) - - if success: - processed_count += 1 - else: - failed_count += 1 - - # Print progress periodically - if (i + 1) % 20 == 0 or (i + 1) == len(inactive_projects): - progress = (i + 1) / len(inactive_projects) * 100 - print(f"Progress: {i + 1}/{len(inactive_projects)} sandboxes processed ({progress:.1f}%)") - print(f" - Processed: {processed_count}, Failed: {failed_count}") - - return processed_count, failed_count - - -async def main(): - """Main function to run the script.""" - # Parse command line arguments - parser = argparse.ArgumentParser(description='Archive sandboxes for projects without active billing') - parser.add_argument('--dry-run', action='store_true', help='Show what would be archived without actually archiving') - args = parser.parse_args() - - logger.info("Starting sandbox cleanup for projects without active billing") - if args.dry_run: - logger.info("DRY RUN MODE - No sandboxes will be archived") - - # Print environment info - print(f"Environment Mode: {os.getenv('ENV_MODE', 'Not set')}") - print(f"Daytona Server: {os.getenv('DAYTONA_SERVER_URL', 'Not set')}") - - try: - # Initialize global DB connection - global db_connection - db_connection = DBConnection() - - # Get all account_ids that have an active billing customer - active_billing_customer_account_ids = await get_active_billing_customer_account_ids() - - # Get all projects with sandboxes - all_projects = await get_all_projects() - - if not all_projects: - logger.info("No projects with sandboxes to process") - return - - # Filter projects whose account_id is not in the active billing customers list - inactive_projects = [ - project for project in all_projects - if project.get('account_id') not in active_billing_customer_account_ids - ] - - # Print summary of what will be processed - active_projects_count = len(all_projects) - len(inactive_projects) - print("\n===== SANDBOX CLEANUP SUMMARY =====") - print(f"Total projects found: {len(all_projects)}") - print(f"Projects with active billing accounts: {active_projects_count}") - print(f"Projects without active billing accounts: {len(inactive_projects)}") - print(f"Sandboxes that will be archived: {len(inactive_projects)}") - print("===================================") - - logger.info(f"Found {len(inactive_projects)} projects without an active billing customer account") - - if not inactive_projects: - logger.info("No projects to archive sandboxes for") - return - - # Ask for confirmation before proceeding - if not args.dry_run: - print("\n⚠️ WARNING: You are about to archive sandboxes for inactive accounts ⚠️") - print("This action cannot be undone!") - confirmation = input("\nAre you sure you want to proceed with archiving? (TRUE/FALSE): ").strip().upper() - - if confirmation != "TRUE": - print("Archiving cancelled. Exiting script.") - logger.info("Archiving cancelled by user") - return - - print("\nProceeding with sandbox archiving...\n") - logger.info("User confirmed sandbox archiving") - - # List all projects to be processed - for i, project in enumerate(inactive_projects[:5]): # Just show first 5 for brevity - account_id = project.get('account_id', 'Unknown') - project_name = project.get('name', 'Unknown') - project_id = project.get('project_id', 'Unknown') - sandbox_id = project['sandbox'].get('id') - - print(f"{i+1}. Project: {project_name}") - print(f" Project ID: {project_id}") - print(f" Account ID: {account_id}") - print(f" Sandbox ID: {sandbox_id}") - - if len(inactive_projects) > 5: - print(f" ... and {len(inactive_projects) - 5} more projects") - - # Process all sandboxes - processed_count, failed_count = await process_sandboxes(inactive_projects, args.dry_run) - - # Print final summary - print("\nSandbox Cleanup Summary:") - print(f"Total projects without active billing: {len(inactive_projects)}") - print(f"Total sandboxes processed: {len(inactive_projects)}") - - if args.dry_run: - print(f"DRY RUN: No sandboxes were actually archived") - else: - print(f"Successfully processed: {processed_count}") - print(f"Failed to process: {failed_count}") - - logger.info("Sandbox cleanup completed") - - except Exception as e: - logger.error(f"Error during sandbox cleanup: {str(e)}") - sys.exit(1) - finally: - # Clean up database connection - if db_connection: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/app/utils/scripts/archive_old_sandboxes.py b/app/utils/scripts/archive_old_sandboxes.py deleted file mode 100644 index ae020d8..0000000 --- a/app/utils/scripts/archive_old_sandboxes.py +++ /dev/null @@ -1,344 +0,0 @@ -#!/usr/bin/env python -""" -Script to archive sandboxes for projects that are older than 1 day. - -Usage: - python archive_old_sandboxes.py [--days N] [--dry-run] - -This script: -1. Gets all projects from the projects table -2. Filters projects created more than N days ago (default: 1 day) -3. Archives the sandboxes for those projects - -Make sure your environment variables are properly set: -- SUPABASE_URL -- SUPABASE_SERVICE_ROLE_KEY -- DAYTONA_SERVER_URL -""" - -# TODO: SAVE THE LATEST SANDBOX STATE SOMEWHERE OR LIKE MASS CHECK THE STATE BEFORE STARTING TO ARCHIVE - AS ITS GOING TO GO OVER A BUNCH THAT ARE ALREADY ARCHIVED – MAYBE BEST TO GET ALL FROM DAYTONA AND THEN RUN THE ARCHIVE ONLY ON THE ONES THAT MEET THE CRITERIA (STOPPED STATE) - -import asyncio -import sys -import os -import argparse -from typing import List, Dict, Any -from datetime import datetime, timedelta -from dotenv import load_dotenv - -# Load script-specific environment variables -load_dotenv(".env") - -from services.supabase import DBConnection -from sandbox.sandbox import daytona -from utils.logger import logger - -# Global DB connection to reuse -db_connection = None - - -async def get_old_projects(days_threshold: int = 1) -> List[Dict[str, Any]]: - """ - Query all projects created more than N days ago. - - Args: - days_threshold: Number of days threshold (default: 1) - - Returns: - List of projects with their sandbox information - """ - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Print the Supabase URL being used - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - # Calculate the date threshold - threshold_date = (datetime.now() - timedelta(days=days_threshold)).isoformat() - - # Initialize variables for pagination - all_projects = [] - page_size = 1000 - current_page = 0 - has_more = True - - logger.info(f"Starting to fetch projects older than {days_threshold} day(s)") - print(f"Looking for projects created before: {threshold_date}") - - # Paginate through all projects - while has_more: - # Query projects with pagination - start_range = current_page * page_size - end_range = start_range + page_size - 1 - - logger.info(f"Fetching projects page {current_page+1} (range: {start_range}-{end_range})") - - try: - result = await client.table('projects').select( - 'project_id', - 'name', - 'created_at', - 'account_id', - 'sandbox' - ).order('created_at', desc=True).range(start_range, end_range).execute() - - # Debug info - print raw response - print(f"Response data length: {len(result.data)}") - - if not result.data: - print("No more data returned from query, ending pagination") - has_more = False - else: - # Print a sample project to see the actual data structure - if current_page == 0 and result.data: - print(f"Sample project data: {result.data[0]}") - - all_projects.extend(result.data) - current_page += 1 - - # Progress update - logger.info(f"Loaded {len(all_projects)} projects so far") - print(f"Loaded {len(all_projects)} projects so far...") - - # Check if we've reached the end - if we got fewer results than the page size - if len(result.data) < page_size: - print(f"Got {len(result.data)} records which is less than page size {page_size}, ending pagination") - has_more = False - else: - print(f"Full page returned ({len(result.data)} records), continuing to next page") - - except Exception as e: - logger.error(f"Error during pagination: {str(e)}") - print(f"Error during pagination: {str(e)}") - has_more = False # Stop on error - - # Print the query result summary - total_projects = len(all_projects) - print(f"Found {total_projects} total projects in database") - logger.info(f"Total projects found in database: {total_projects}") - - if not all_projects: - logger.info("No projects found in database") - return [] - - # Filter projects that are older than the threshold and have sandbox information - old_projects_with_sandboxes = [ - project for project in all_projects - if project.get('created_at') and project.get('created_at') < threshold_date - and project.get('sandbox') and project['sandbox'].get('id') - ] - - logger.info(f"Found {len(old_projects_with_sandboxes)} old projects with sandboxes") - - # Print a few sample old projects for debugging - if old_projects_with_sandboxes: - print("\nSample of old projects with sandboxes:") - for i, project in enumerate(old_projects_with_sandboxes[:3]): - print(f" {i+1}. {project.get('name')} (Created: {project.get('created_at')})") - print(f" Sandbox ID: {project['sandbox'].get('id')}") - if i >= 2: - break - - return old_projects_with_sandboxes - - -async def archive_sandbox(project: Dict[str, Any], dry_run: bool) -> bool: - """ - Archive a single sandbox. - - Args: - project: Project information containing sandbox to archive - dry_run: If True, only simulate archiving - - Returns: - True if successful, False otherwise - """ - sandbox_id = project['sandbox'].get('id') - project_name = project.get('name', 'Unknown') - project_id = project.get('project_id', 'Unknown') - created_at = project.get('created_at', 'Unknown') - - try: - logger.info(f"Checking sandbox {sandbox_id} for project '{project_name}' (ID: {project_id}, Created: {created_at})") - - if dry_run: - logger.info(f"DRY RUN: Would archive sandbox {sandbox_id}") - print(f"Would archive sandbox {sandbox_id} for project '{project_name}' (Created: {created_at})") - return True - - # Get the sandbox - sandbox = daytona.get(sandbox_id) - - # Check sandbox state - it must be stopped before archiving - sandbox_info = sandbox.info() - - # Log the current state - logger.info(f"Sandbox {sandbox_id} is in '{sandbox_info.state}' state") - - # Only archive if the sandbox is in the stopped state - if sandbox_info.state == "stopped": - logger.info(f"Archiving sandbox {sandbox_id} as it is in stopped state") - sandbox.archive() - logger.info(f"Successfully archived sandbox {sandbox_id}") - return True - else: - logger.info(f"Skipping sandbox {sandbox_id} as it is not in stopped state (current: {sandbox_info.state})") - return True - - except Exception as e: - import traceback - error_type = type(e).__name__ - stack_trace = traceback.format_exc() - - # Log detailed error information - logger.error(f"Error processing sandbox {sandbox_id}: {str(e)}") - logger.error(f"Error type: {error_type}") - logger.error(f"Stack trace:\n{stack_trace}") - - # If the exception has a response attribute (like in HTTP errors), log it - if hasattr(e, 'response'): - try: - response_data = e.response.json() if hasattr(e.response, 'json') else str(e.response) - logger.error(f"Response data: {response_data}") - except Exception: - logger.error(f"Could not parse response data from error") - - print(f"Failed to process sandbox {sandbox_id}: {error_type} - {str(e)}") - return False - - -async def process_sandboxes(old_projects: List[Dict[str, Any]], dry_run: bool) -> tuple[int, int]: - """ - Process all sandboxes sequentially. - - Args: - old_projects: List of projects older than the threshold - dry_run: Whether to actually archive sandboxes or just simulate - - Returns: - Tuple of (processed_count, failed_count) - """ - processed_count = 0 - failed_count = 0 - - if dry_run: - logger.info(f"DRY RUN: Would archive {len(old_projects)} sandboxes") - else: - logger.info(f"Archiving {len(old_projects)} sandboxes") - - print(f"Processing {len(old_projects)} sandboxes...") - - # Process each sandbox sequentially - for i, project in enumerate(old_projects): - success = await archive_sandbox(project, dry_run) - - if success: - processed_count += 1 - else: - failed_count += 1 - - # Print progress periodically - if (i + 1) % 20 == 0 or (i + 1) == len(old_projects): - progress = (i + 1) / len(old_projects) * 100 - print(f"Progress: {i + 1}/{len(old_projects)} sandboxes processed ({progress:.1f}%)") - print(f" - Processed: {processed_count}, Failed: {failed_count}") - - return processed_count, failed_count - - -async def main(): - """Main function to run the script.""" - # Parse command line arguments - parser = argparse.ArgumentParser(description='Archive sandboxes for projects older than N days') - parser.add_argument('--days', type=int, default=1, help='Age threshold in days (default: 1)') - parser.add_argument('--dry-run', action='store_true', help='Show what would be archived without actually archiving') - args = parser.parse_args() - - logger.info(f"Starting sandbox cleanup for projects older than {args.days} day(s)") - if args.dry_run: - logger.info("DRY RUN MODE - No sandboxes will be archived") - - # Print environment info - print(f"Environment Mode: {os.getenv('ENV_MODE', 'Not set')}") - print(f"Daytona Server: {os.getenv('DAYTONA_SERVER_URL', 'Not set')}") - - try: - # Initialize global DB connection - global db_connection - db_connection = DBConnection() - - # Get all projects older than the threshold - old_projects = await get_old_projects(args.days) - - if not old_projects: - logger.info(f"No projects older than {args.days} day(s) with sandboxes to process") - print(f"No projects older than {args.days} day(s) with sandboxes to archive.") - return - - # Print summary of what will be processed - print("\n===== SANDBOX CLEANUP SUMMARY =====") - print(f"Projects older than {args.days} day(s): {len(old_projects)}") - print(f"Sandboxes that will be archived: {len(old_projects)}") - print("===================================") - - logger.info(f"Found {len(old_projects)} projects older than {args.days} day(s)") - - # Ask for confirmation before proceeding - if not args.dry_run: - print("\n⚠️ WARNING: You are about to archive sandboxes for old projects ⚠️") - print("This action cannot be undone!") - confirmation = input("\nAre you sure you want to proceed with archiving? (TRUE/FALSE): ").strip().upper() - - if confirmation != "TRUE": - print("Archiving cancelled. Exiting script.") - logger.info("Archiving cancelled by user") - return - - print("\nProceeding with sandbox archiving...\n") - logger.info("User confirmed sandbox archiving") - - # List a sample of projects to be processed - for i, project in enumerate(old_projects[:5]): # Just show first 5 for brevity - created_at = project.get('created_at', 'Unknown') - project_name = project.get('name', 'Unknown') - project_id = project.get('project_id', 'Unknown') - sandbox_id = project['sandbox'].get('id') - - print(f"{i+1}. Project: {project_name}") - print(f" Project ID: {project_id}") - print(f" Created At: {created_at}") - print(f" Sandbox ID: {sandbox_id}") - - if len(old_projects) > 5: - print(f" ... and {len(old_projects) - 5} more projects") - - # Process all sandboxes - processed_count, failed_count = await process_sandboxes(old_projects, args.dry_run) - - # Print final summary - print("\nSandbox Cleanup Summary:") - print(f"Total projects older than {args.days} day(s): {len(old_projects)}") - print(f"Total sandboxes processed: {len(old_projects)}") - - if args.dry_run: - print(f"DRY RUN: No sandboxes were actually archived") - else: - print(f"Successfully processed: {processed_count}") - print(f"Failed to process: {failed_count}") - - logger.info("Sandbox cleanup completed") - - except Exception as e: - logger.error(f"Error during sandbox cleanup: {str(e)}") - sys.exit(1) - finally: - # Clean up database connection - if db_connection: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/app/utils/scripts/copy_project.py b/app/utils/scripts/copy_project.py deleted file mode 100644 index e35cdf0..0000000 --- a/app/utils/scripts/copy_project.py +++ /dev/null @@ -1,388 +0,0 @@ -import asyncio -import argparse -from dotenv import load_dotenv - -load_dotenv(".env") - -from services.supabase import DBConnection -from daytona_sdk import Sandbox -from sandbox.sandbox import daytona, create_sandbox, delete_sandbox -from utils.logger import logger - -db_connection = None -db = None - - -async def get_db(): - global db_connection, db - if db_connection is None or db is None: - db_connection = DBConnection() - db = await db_connection.client - return db - - -async def get_project(project_id: str): - db = await get_db() - project = ( - await db.schema("public") - .from_("projects") - .select("*") - .eq("project_id", project_id) - .maybe_single() - .execute() - ) - return project.data - - -async def get_threads(project_id: str): - db = await get_db() - threads = ( - await db.schema("public") - .from_("threads") - .select("*") - .eq("project_id", project_id) - .execute() - ) - return threads.data - - -async def copy_thread(thread_id: str, account_id: str, project_id: str): - db = await get_db() - thread = ( - await db.schema("public") - .from_("threads") - .select("*") - .eq("thread_id", thread_id) - .maybe_single() - .execute() - ) - - if not thread.data: - raise Exception(f"Thread {thread_id} not found") - - thread_data = thread.data - new_thread = ( - await db.schema("public") - .from_("threads") - .insert( - { - "account_id": account_id, - "project_id": project_id, - "is_public": thread_data["is_public"], - "agent_id": thread_data["agent_id"], - "metadata": thread_data["metadata"] or {}, - } - ) - .execute() - ) - return new_thread.data[0] - - -async def copy_project(project_id: str, to_user_id: str, sandbox_data: dict): - db = await get_db() - project = await get_project(project_id) - to_user = await get_user(to_user_id) - - if not project: - raise Exception(f"Project {project_id} not found") - if not to_user: - raise Exception(f"User {to_user_id} not found") - - result = ( - await db.schema("public") - .from_("projects") - .insert( - { - "name": project["name"], - "description": project["description"], - "account_id": to_user["id"], - "is_public": project["is_public"], - "sandbox": sandbox_data, - } - ) - .execute() - ) - return result.data[0] - - -async def copy_agent_runs(thread_id: str, new_thread_id: str): - db = await get_db() - agent_runs = ( - await db.schema("public") - .from_("agent_runs") - .select("*") - .eq("thread_id", thread_id) - .execute() - ) - - async def copy_single_agent_run(agent_run, new_thread_id, db): - new_agent_run = ( - await db.schema("public") - .from_("agent_runs") - .insert( - { - "thread_id": new_thread_id, - "status": agent_run["status"], - "started_at": agent_run["started_at"], - "completed_at": agent_run["completed_at"], - "responses": agent_run["responses"], - "error": agent_run["error"], - } - ) - .execute() - ) - return new_agent_run.data[0] - - tasks = [ - copy_single_agent_run(agent_run, new_thread_id, db) - for agent_run in agent_runs.data - ] - new_agent_runs = await asyncio.gather(*tasks) - return new_agent_runs - - -async def copy_messages(thread_id: str, new_thread_id: str): - db = await get_db() - messages_data = [] - offset = 0 - batch_size = 1000 - - while True: - batch = ( - await db.schema("public") - .from_("messages") - .select("*") - .eq("thread_id", thread_id) - .range(offset, offset + batch_size - 1) - .execute() - ) - - if not batch.data: - break - - messages_data.extend(batch.data) - - if len(batch.data) < batch_size: - break - - offset += batch_size - - async def copy_single_message(message, new_thread_id, db): - new_message = ( - await db.schema("public") - .from_("messages") - .insert( - { - "thread_id": new_thread_id, - "type": message["type"], - "is_llm_message": message["is_llm_message"], - "content": message["content"], - "metadata": message["metadata"], - "created_at": message["created_at"], - "updated_at": message["updated_at"], - } - ) - .execute() - ) - return new_message.data[0] - - tasks = [] - for message in messages_data: - tasks.append(copy_single_message(message, new_thread_id, db)) - - # Process tasks in batches to avoid overwhelming the database - batch_size = 100 - new_messages = [] - for i in range(0, len(tasks), batch_size): - batch_tasks = tasks[i : i + batch_size] - batch_results = await asyncio.gather(*batch_tasks) - new_messages.extend(batch_results) - # Add delay between batches - if i + batch_size < len(tasks): - await asyncio.sleep(1) - - return new_messages - - -async def get_user(user_id: str): - db = await get_db() - user = await db.auth.admin.get_user_by_id(user_id) - return user.user.model_dump() - - -async def copy_sandbox(sandbox_id: str, password: str, project_id: str) -> Sandbox: - sandbox = daytona.find_one(sandbox_id=sandbox_id) - if not sandbox: - raise Exception(f"Sandbox {sandbox_id} not found") - - # TODO: Currently there's no way to create a copy of a sandbox, so we will create a new one - new_sandbox = create_sandbox(password, project_id) - return new_sandbox - - -async def main(): - """Main function to run the script.""" - # Parse command line arguments - parser = argparse.ArgumentParser(description="Create copy of a project") - parser.add_argument( - "--project-id", type=str, help="Project ID to copy", required=True - ) - parser.add_argument( - "--new-user-id", - type=str, - default=None, - help="[OPTIONAL] User ID to copy the project to", - required=False, - ) - args = parser.parse_args() - - # Initialize variables for cleanup - new_sandbox = None - new_project = None - new_threads = [] - new_agent_runs = [] - new_messages = [] - - try: - project = await get_project(args.project_id) - if not project: - raise Exception(f"Project {args.project_id} not found") - - to_user_id = args.new_user_id or project["account_id"] - to_user = await get_user(to_user_id) - - logger.info( - f"Project: {project['project_id']} ({project['name']}) -> User: {to_user['id']} ({to_user['email']})" - ) - - new_sandbox = await copy_sandbox( - project["sandbox"]["id"], project["sandbox"]["pass"], args.project_id - ) - if new_sandbox: - vnc_link = new_sandbox.get_preview_link(6080) - website_link = new_sandbox.get_preview_link(8080) - vnc_url = ( - vnc_link.url - if hasattr(vnc_link, "url") - else str(vnc_link).split("url='")[1].split("'")[0] - ) - website_url = ( - website_link.url - if hasattr(website_link, "url") - else str(website_link).split("url='")[1].split("'")[0] - ) - token = None - if hasattr(vnc_link, "token"): - token = vnc_link.token - elif "token='" in str(vnc_link): - token = str(vnc_link).split("token='")[1].split("'")[0] - else: - raise Exception("Failed to create new sandbox") - - sandbox_data = { - "id": new_sandbox.id, - "pass": project["sandbox"]["pass"], - "token": token, - "vnc_preview": vnc_url, - "sandbox_url": website_url, - } - logger.info(f"New sandbox: {new_sandbox.id}") - - new_project = await copy_project( - project["project_id"], to_user["id"], sandbox_data - ) - logger.info(f"New project: {new_project['project_id']} ({new_project['name']})") - - threads = await get_threads(project["project_id"]) - if threads: - for thread in threads: - new_thread = await copy_thread( - thread["thread_id"], to_user["id"], new_project["project_id"] - ) - new_threads.append(new_thread) - logger.info(f"New threads: {len(new_threads)}") - - for i in range(len(new_threads)): - runs = await copy_agent_runs( - threads[i]["thread_id"], new_threads[i]["thread_id"] - ) - new_agent_runs.extend(runs) - logger.info(f"New agent runs: {len(new_agent_runs)}") - - for i in range(len(new_threads)): - messages = await copy_messages( - threads[i]["thread_id"], new_threads[i]["thread_id"] - ) - new_messages.extend(messages) - logger.info(f"New messages: {len(new_messages)}") - else: - logger.info("No threads found for this project") - - except Exception as e: - db = await get_db() - # Clean up any resources that were created before the error - if new_sandbox: - try: - logger.info(f"Cleaning up sandbox: {new_sandbox.id}") - await delete_sandbox(new_sandbox.id) - except Exception as cleanup_error: - logger.error( - f"Error cleaning up sandbox {new_sandbox.id}: {cleanup_error}" - ) - - if new_messages: - for message in new_messages: - try: - logger.info(f"Cleaning up message: {message['message_id']}") - await db.table("messages").delete().eq( - "message_id", message["message_id"] - ).execute() - except Exception as cleanup_error: - logger.error( - f"Error cleaning up message {message['message_id']}: {cleanup_error}" - ) - - if new_agent_runs: - for agent_run in new_agent_runs: - try: - logger.info(f"Cleaning up agent run: {agent_run['id']}") - await db.table("agent_runs").delete().eq( - "id", agent_run["id"] - ).execute() - except Exception as cleanup_error: - logger.error( - f"Error cleaning up agent run {agent_run['id']}: {cleanup_error}" - ) - - if new_threads: - for thread in new_threads: - try: - logger.info(f"Cleaning up thread: {thread['thread_id']}") - await db.table("threads").delete().eq( - "thread_id", thread["thread_id"] - ).execute() - except Exception as cleanup_error: - logger.error( - f"Error cleaning up thread {thread['thread_id']}: {cleanup_error}" - ) - - if new_project: - try: - logger.info(f"Cleaning up project: {new_project['project_id']}") - await db.table("projects").delete().eq( - "project_id", new_project["project_id"] - ).execute() - except Exception as cleanup_error: - logger.error( - f"Error cleaning up project {new_project['project_id']}: {cleanup_error}" - ) - - await DBConnection.disconnect() - raise e - - finally: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/app/utils/scripts/delete_user_sandboxes.py b/app/utils/scripts/delete_user_sandboxes.py deleted file mode 100644 index fe1254e..0000000 --- a/app/utils/scripts/delete_user_sandboxes.py +++ /dev/null @@ -1,138 +0,0 @@ -#!/usr/bin/env python -""" -Script to query and delete sandboxes for a given account ID. - -Usage: - python delete_user_sandboxes.py -""" - -import asyncio -import sys -import os -from typing import List, Dict, Any -from dotenv import load_dotenv - -# Load script-specific environment variables -load_dotenv(".env") - -from services.supabase import DBConnection -from sandbox.sandbox import daytona -from utils.logger import logger - - -async def get_user_sandboxes(account_id: str) -> List[Dict[str, Any]]: - """ - Query all projects and their sandboxes associated with a specific account ID. - - Args: - account_id: The account ID to query - - Returns: - List of projects with sandbox information - """ - db = DBConnection() - client = await db.client - - # Print the Supabase URL being used - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - # Query projects by account_id - result = await client.table('projects').select( - 'project_id', - 'name', - 'sandbox' - ).eq('account_id', account_id).execute() - - # Print the query result for debugging - print(f"Query result: {result}") - - if not result.data: - logger.info(f"No projects found for account ID: {account_id}") - return [] - - # Filter projects with sandbox information - projects_with_sandboxes = [ - project for project in result.data - if project.get('sandbox') and project['sandbox'].get('id') - ] - - logger.info(f"Found {len(projects_with_sandboxes)} projects with sandboxes for account ID: {account_id}") - return projects_with_sandboxes - - -async def delete_sandboxes(projects: List[Dict[str, Any]]) -> None: - """ - Delete all sandboxes from the provided list of projects. - - Args: - projects: List of projects with sandbox information - """ - if not projects: - logger.info("No sandboxes to delete") - return - - for project in projects: - sandbox_id = project['sandbox'].get('id') - project_name = project.get('name', 'Unknown') - project_id = project.get('project_id', 'Unknown') - - if not sandbox_id: - continue - - try: - logger.info(f"Deleting sandbox {sandbox_id} for project '{project_name}' (ID: {project_id})") - - # Get the sandbox and delete it - sandbox = daytona.get(sandbox_id) - daytona.delete(sandbox) - - logger.info(f"Successfully deleted sandbox {sandbox_id}") - except Exception as e: - logger.error(f"Error deleting sandbox {sandbox_id}: {str(e)}") - - -async def main(): - """Main function to run the script.""" - if len(sys.argv) != 2: - print(f"Usage: python {sys.argv[0]} ") - sys.exit(1) - - account_id = sys.argv[1] - logger.info(f"Starting sandbox cleanup for account ID: {account_id}") - - # Print environment info - print(f"Environment Mode: {os.getenv('ENV_MODE', 'Not set')}") - print(f"Daytona Server: {os.getenv('DAYTONA_SERVER_URL', 'Not set')}") - - try: - # Query projects with sandboxes - projects = await get_user_sandboxes(account_id) - - # Print sandbox information - for i, project in enumerate(projects): - sandbox_id = project['sandbox'].get('id', 'N/A') - print(f"{i+1}. Project: {project.get('name', 'Unknown')}") - print(f" Project ID: {project.get('project_id', 'Unknown')}") - print(f" Sandbox ID: {sandbox_id}") - - # Confirm deletion - if projects: - confirm = input(f"\nDelete {len(projects)} sandboxes? (y/n): ") - if confirm.lower() == 'y': - await delete_sandboxes(projects) - logger.info("Sandbox cleanup completed") - else: - logger.info("Sandbox deletion cancelled") - else: - logger.info("No sandboxes found for deletion") - - except Exception as e: - logger.error(f"Error during sandbox cleanup: {str(e)}") - sys.exit(1) - finally: - # Clean up database connection - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/app/utils/scripts/export_import_project.py b/app/utils/scripts/export_import_project.py deleted file mode 100644 index b5632ff..0000000 --- a/app/utils/scripts/export_import_project.py +++ /dev/null @@ -1,400 +0,0 @@ -import asyncio -import argparse -import json -import os -from datetime import datetime -from dotenv import load_dotenv - -load_dotenv(".env") - -from services.supabase import DBConnection -from daytona_sdk import Sandbox -from sandbox.sandbox import daytona, create_sandbox, delete_sandbox -from utils.logger import logger - -db_connection = None -db = None - - -async def get_db(): - global db_connection, db - if db_connection is None or db is None: - db_connection = DBConnection() - db = await db_connection.client - return db - - -async def get_project(project_id: str): - db = await get_db() - project = ( - await db.schema("public") - .from_("projects") - .select("*") - .eq("project_id", project_id) - .maybe_single() - .execute() - ) - return project.data - - -async def get_threads(project_id: str): - db = await get_db() - threads = ( - await db.schema("public") - .from_("threads") - .select("*") - .eq("project_id", project_id) - .execute() - ) - return threads.data - - -async def get_agent_runs(thread_id: str): - db = await get_db() - agent_runs = ( - await db.schema("public") - .from_("agent_runs") - .select("*") - .eq("thread_id", thread_id) - .execute() - ) - return agent_runs.data - - -async def get_messages(thread_id: str): - db = await get_db() - messages_data = [] - offset = 0 - batch_size = 1000 - - while True: - batch = ( - await db.schema("public") - .from_("messages") - .select("*") - .eq("thread_id", thread_id) - .range(offset, offset + batch_size - 1) - .execute() - ) - - if not batch.data: - break - - messages_data.extend(batch.data) - - if len(batch.data) < batch_size: - break - - offset += batch_size - - return messages_data - - -async def get_user(user_id: str): - db = await get_db() - user = await db.auth.admin.get_user_by_id(user_id) - return user.user.model_dump() - - -async def export_project_to_file(project_id: str, output_file: str): - """Export all project data to a JSON file.""" - try: - logger.info(f"Starting export of project {project_id}") - - # Get project data - project = await get_project(project_id) - if not project: - raise Exception(f"Project {project_id} not found") - - logger.info(f"Exporting project: {project['name']}") - - # Get threads - threads = await get_threads(project_id) - logger.info(f"Found {len(threads)} threads") - - # Get agent runs and messages for each thread - threads_data = [] - for thread in threads: - thread_data = dict(thread) - - # Get agent runs for this thread - agent_runs = await get_agent_runs(thread["thread_id"]) - thread_data["agent_runs"] = agent_runs - - # Get messages for this thread - messages = await get_messages(thread["thread_id"]) - thread_data["messages"] = messages - - threads_data.append(thread_data) - logger.info(f"Thread {thread['thread_id']}: {len(agent_runs)} runs, {len(messages)} messages") - - # Prepare export data - export_data = { - "export_metadata": { - "export_date": datetime.now().isoformat(), - "project_id": project_id, - "project_name": project["name"] - }, - "project": project, - "threads": threads_data - } - - # Write to file - os.makedirs(os.path.dirname(output_file), exist_ok=True) - with open(output_file, 'w', encoding='utf-8') as f: - json.dump(export_data, f, indent=2, ensure_ascii=False, default=str) - - logger.info(f"Project exported successfully to {output_file}") - logger.info(f"Export summary: 1 project, {len(threads_data)} threads") - - return export_data - - except Exception as e: - logger.error(f"Error exporting project: {e}") - raise e - finally: - await DBConnection.disconnect() - - -async def import_project_from_file(input_file: str, to_user_id: str = None, create_new_sandbox: bool = True): - """Import project data from a JSON file and create a new project.""" - new_sandbox = None - new_project = None - new_threads = [] - new_agent_runs = [] - new_messages = [] - - try: - logger.info(f"Starting import from {input_file}") - - # Read data from file - with open(input_file, 'r', encoding='utf-8') as f: - import_data = json.load(f) - - project_data = import_data["project"] - threads_data = import_data["threads"] - - logger.info(f"Importing project: {project_data['name']}") - logger.info(f"Found {len(threads_data)} threads to import") - - # Determine target user - to_user_id = to_user_id or project_data["account_id"] - to_user = await get_user(to_user_id) - - logger.info(f"Target user: {to_user['id']} ({to_user['email']})") - - # Create new sandbox if requested - if create_new_sandbox: - logger.info("Creating new sandbox...") - new_sandbox = create_sandbox(project_data["sandbox"]["pass"], project_data["project_id"]) - - if new_sandbox: - vnc_link = new_sandbox.get_preview_link(6080) - website_link = new_sandbox.get_preview_link(8080) - vnc_url = ( - vnc_link.url - if hasattr(vnc_link, "url") - else str(vnc_link).split("url='")[1].split("'")[0] - ) - website_url = ( - website_link.url - if hasattr(website_link, "url") - else str(website_link).split("url='")[1].split("'")[0] - ) - token = None - if hasattr(vnc_link, "token"): - token = vnc_link.token - elif "token='" in str(vnc_link): - token = str(vnc_link).split("token='")[1].split("'")[0] - - sandbox_data = { - "id": new_sandbox.id, - "pass": project_data["sandbox"]["pass"], - "token": token, - "vnc_preview": vnc_url, - "sandbox_url": website_url, - } - logger.info(f"New sandbox created: {new_sandbox.id}") - else: - raise Exception("Failed to create new sandbox") - else: - # Use existing sandbox data - sandbox_data = project_data["sandbox"] - logger.info("Using existing sandbox data") - - # Create new project - db = await get_db() - result = ( - await db.schema("public") - .from_("projects") - .insert( - { - "name": project_data["name"], - "description": project_data["description"], - "account_id": to_user["id"], - "is_public": project_data["is_public"], - "sandbox": sandbox_data, - } - ) - .execute() - ) - new_project = result.data[0] - logger.info(f"New project created: {new_project['project_id']} ({new_project['name']})") - - # Import threads - for thread_data in threads_data: - # Create new thread - new_thread = ( - await db.schema("public") - .from_("threads") - .insert( - { - "account_id": to_user["id"], - "project_id": new_project["project_id"], - "is_public": thread_data["is_public"], - "agent_id": thread_data["agent_id"], - "metadata": thread_data["metadata"] or {}, - } - ) - .execute() - ) - new_thread = new_thread.data[0] - new_threads.append(new_thread) - - # Create agent runs for this thread - for agent_run_data in thread_data.get("agent_runs", []): - new_agent_run = ( - await db.schema("public") - .from_("agent_runs") - .insert( - { - "thread_id": new_thread["thread_id"], - "status": agent_run_data["status"], - "started_at": agent_run_data["started_at"], - "completed_at": agent_run_data["completed_at"], - "responses": agent_run_data["responses"], - "error": agent_run_data["error"], - } - ) - .execute() - ) - new_agent_runs.append(new_agent_run.data[0]) - - # Create messages for this thread in batches - messages = thread_data.get("messages", []) - batch_size = 100 - for i in range(0, len(messages), batch_size): - batch_messages = messages[i:i + batch_size] - message_inserts = [] - - for message_data in batch_messages: - message_inserts.append({ - "thread_id": new_thread["thread_id"], - "type": message_data["type"], - "is_llm_message": message_data["is_llm_message"], - "content": message_data["content"], - "metadata": message_data["metadata"], - "created_at": message_data["created_at"], - "updated_at": message_data["updated_at"], - }) - - if message_inserts: - batch_result = ( - await db.schema("public") - .from_("messages") - .insert(message_inserts) - .execute() - ) - new_messages.extend(batch_result.data) - - # Add delay between batches - if i + batch_size < len(messages): - await asyncio.sleep(0.5) - - logger.info(f"Thread imported: {len(thread_data.get('agent_runs', []))} runs, {len(messages)} messages") - - logger.info(f"Import completed successfully!") - logger.info(f"Summary: 1 project, {len(new_threads)} threads, {len(new_agent_runs)} agent runs, {len(new_messages)} messages") - - return { - "project": new_project, - "threads": new_threads, - "agent_runs": new_agent_runs, - "messages": new_messages - } - - except Exception as e: - logger.error(f"Error importing project: {e}") - - # Clean up any resources that were created before the error - db = await get_db() - - if new_sandbox: - try: - logger.info(f"Cleaning up sandbox: {new_sandbox.id}") - await delete_sandbox(new_sandbox.id) - except Exception as cleanup_error: - logger.error(f"Error cleaning up sandbox {new_sandbox.id}: {cleanup_error}") - - if new_messages: - for message in new_messages: - try: - await db.table("messages").delete().eq("message_id", message["message_id"]).execute() - except Exception as cleanup_error: - logger.error(f"Error cleaning up message {message['message_id']}: {cleanup_error}") - - if new_agent_runs: - for agent_run in new_agent_runs: - try: - await db.table("agent_runs").delete().eq("id", agent_run["id"]).execute() - except Exception as cleanup_error: - logger.error(f"Error cleaning up agent run {agent_run['id']}: {cleanup_error}") - - if new_threads: - for thread in new_threads: - try: - await db.table("threads").delete().eq("thread_id", thread["thread_id"]).execute() - except Exception as cleanup_error: - logger.error(f"Error cleaning up thread {thread['thread_id']}: {cleanup_error}") - - if new_project: - try: - await db.table("projects").delete().eq("project_id", new_project["project_id"]).execute() - except Exception as cleanup_error: - logger.error(f"Error cleaning up project {new_project['project_id']}: {cleanup_error}") - - await DBConnection.disconnect() - raise e - - finally: - await DBConnection.disconnect() - - -async def main(): - """Main function to run the script.""" - parser = argparse.ArgumentParser(description="Export/Import project data") - parser.add_argument("action", choices=["export", "import"], help="Action to perform") - parser.add_argument("--project-id", type=str, help="Project ID to export (required for export)") - parser.add_argument("--file", type=str, help="File path for export/import", required=True) - parser.add_argument("--user-id", type=str, help="User ID to import project to (optional for import)") - parser.add_argument("--no-sandbox", action="store_true", help="Don't create new sandbox during import") - - args = parser.parse_args() - - try: - if args.action == "export": - if not args.project_id: - raise Exception("--project-id is required for export") - await export_project_to_file(args.project_id, args.file) - - elif args.action == "import": - create_new_sandbox = not args.no_sandbox - await import_project_from_file(args.file, args.user_id, create_new_sandbox) - - except Exception as e: - logger.error(f"Script failed: {e}") - raise e - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/app/utils/scripts/generate_share_links.py b/app/utils/scripts/generate_share_links.py deleted file mode 100644 index 6607eeb..0000000 --- a/app/utils/scripts/generate_share_links.py +++ /dev/null @@ -1,166 +0,0 @@ -import asyncio -import sys -import os -from typing import List, Dict, Any -from datetime import datetime -import random -from dotenv import load_dotenv - -load_dotenv(".env") - -from services.supabase import DBConnection -from utils.logger import logger - -db_connection = None - - -async def get_random_thread_ids(n: int) -> List[str]: - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - all_thread_ids = [] - page_size = 1000 - current_page = 0 - has_more = True - - print("Fetching all thread IDs from database (paginated)...") - - while has_more: - start_range = current_page * page_size - end_range = start_range + page_size - 1 - - print(f"Fetching page {current_page + 1} (rows {start_range}-{end_range})...") - - try: - result = await client.table('threads').select('thread_id').range(start_range, end_range).execute() - - if not result.data: - has_more = False - else: - page_thread_ids = [thread['thread_id'] for thread in result.data] - all_thread_ids.extend(page_thread_ids) - - print(f"Loaded {len(page_thread_ids)} thread IDs (total so far: {len(all_thread_ids)})") - if len(result.data) < page_size: - has_more = False - else: - current_page += 1 - - except Exception as e: - logger.error(f"Error during pagination: {str(e)}") - has_more = False - - print(f"Found {len(all_thread_ids)} total thread IDs in database") - - if not all_thread_ids: - logger.info("No threads found in database") - return [] - - if len(all_thread_ids) <= n: - logger.warning(f"Requested {n} threads but only {len(all_thread_ids)} available. Returning all.") - selected_thread_ids = all_thread_ids - else: - selected_thread_ids = random.sample(all_thread_ids, n) - - logger.info(f"Retrieved {len(selected_thread_ids)} random thread IDs") - return selected_thread_ids - - -async def generate_share_links(n: int) -> List[str]: - try: - thread_ids = await get_random_thread_ids(n) - - if not thread_ids: - logger.warning("No thread IDs found, returning empty list") - return [] - - share_links = [f"suna.so/share/{thread_id}" for thread_id in thread_ids] - - logger.info(f"Generated {len(share_links)} share links") - return share_links - - except Exception as e: - logger.error(f"Error generating share links: {str(e)}") - raise - - -def save_links_to_file(share_links: List[str], filename: str = None) -> str: - """Save share links to a text file and return the filename.""" - if filename is None: - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - filename = f"share_links_{timestamp}.txt" - - try: - with open(filename, 'w', encoding='utf-8') as f: - f.write(f"Share Links Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") - f.write("=" * 60 + "\n\n") - - for i, link in enumerate(share_links, 1): - f.write(f"{i}. {link}\n") - - f.write(f"\nTotal: {len(share_links)} share links generated\n") - - logger.info(f"Share links saved to {filename}") - return filename - - except Exception as e: - logger.error(f"Error saving links to file: {str(e)}") - raise - - -async def main(): - logger.info("Starting share link generation process") - - try: - global db_connection - db_connection = DBConnection() - - try: - n = int(input("Enter the number of share links to generate: ")) - if n <= 0: - print("Number must be positive") - return - except ValueError: - print("Please enter a valid number") - return - - custom_filename = input("Enter filename (press Enter for auto-generated): ").strip() - if not custom_filename: - custom_filename = None - elif not custom_filename.endswith('.txt'): - custom_filename += '.txt' - - print(f"\nGenerating {n} random share links...") - - share_links = await generate_share_links(n) - - if not share_links: - print("No share links were generated") - return - - print(f"\nGenerated {len(share_links)} share links:") - print("-" * 50) - for i, link in enumerate(share_links, 1): - print(f"{i}. {link}") - - saved_filename = save_links_to_file(share_links, custom_filename) - - print(f"\nTotal: {len(share_links)} share links generated") - print(f"Links saved to: {saved_filename}") - logger.info("Share link generation completed") - - except Exception as e: - logger.error(f"Error during share link generation: {str(e)}") - sys.exit(1) - finally: - if db_connection: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/app/utils/scripts/get_monthly_usage.py b/app/utils/scripts/get_monthly_usage.py deleted file mode 100644 index 4fdb1e0..0000000 --- a/app/utils/scripts/get_monthly_usage.py +++ /dev/null @@ -1,248 +0,0 @@ -""" -Monthly Usage Script - -This script calculates the monthly usage (in agent run minutes) for a specific user during a specific month. - -Usage: - python backend/utils/scripts/get_monthly_usage.py --user-id --year --month [--verbose] - -Arguments: - --user-id The user ID to get usage for (required) - --year The year (e.g., 2024) (required) - --month The month (1-12) (required) - --verbose Enable verbose logging (optional) - -Examples: - # Get usage for December 2024 - python backend/utils/scripts/get_monthly_usage.py --user-id "user123" --year 2024 --month 12 - - # Get usage with verbose logging - python backend/utils/scripts/get_monthly_usage.py --user-id "user123" --year 2024 --month 11 --verbose - -Output: - The script will output: - - User information (email and ID) - - Month and year - - Total usage in minutes and hours - - Average usage per day (if any usage exists) - - Example output: - === Monthly Usage Report === - User: user@example.com (user123) - Month: December 2024 - Total Usage: 150.45 minutes - Total Usage: 2.51 hours - Average per day: 5.02 minutes - -Features: - - Validates agent runs to exclude invalid durations (>2 hours) - - Handles incomplete agent runs appropriately - - Provides detailed logging for debugging - - Calculates usage only for the specified month - - Shows average daily usage - -Notes: - - The script requires access to the Supabase database - - Make sure the .env file is properly configured - - The script uses the same logic as the billing system for consistency -""" - -import asyncio -import argparse -from datetime import datetime, timezone -from dotenv import load_dotenv - -load_dotenv(".env") - -from services.supabase import DBConnection -from utils.logger import logger - -db_connection = None -db = None - - -async def get_db(): - global db_connection, db - if db_connection is None or db is None: - db_connection = DBConnection() - db = await db_connection.client - return db - - -async def get_user(user_id: str): - """Get user information by user ID.""" - db = await get_db() - user = await db.auth.admin.get_user_by_id(user_id) - return user.user.model_dump() - - -async def calculate_monthly_usage(client, user_id: str, year: int, month: int): - """Calculate total agent run minutes for a specific month for a user.""" - # Get start and end of specified month in UTC - start_of_month = datetime(year, month, 1, tzinfo=timezone.utc) - - # Calculate start of next month for end boundary - if month == 12: - end_of_month = datetime(year + 1, 1, 1, tzinfo=timezone.utc) - else: - end_of_month = datetime(year, month + 1, 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] - logger.info(f"Found {len(thread_ids)} threads for user {user_id}") - - # Then get all agent runs for these threads in specified month - runs_result = ( - await client.table("agent_runs") - .select("id, started_at, completed_at, thread_id") - .in_("thread_id", thread_ids) - .gte("started_at", start_of_month.isoformat()) - .lt("started_at", end_of_month.isoformat()) - .execute() - ) - - if not runs_result.data: - return 0.0, [] - - logger.info(f"Found {len(runs_result.data)} agent runs in {year}-{month:02d}") - - # Calculate total minutes and collect run details - total_seconds = 0 - valid_runs = 0 - run_details = [] - - 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() - # Skip runs that seem invalid (more than 2 hours) - if start_time < end_time - 7200: - logger.warning(f"Skipping run with duration > 2 hours: {run}") - continue - status = "completed" - else: - # For incomplete runs, use end of month as boundary if run started in that month - end_time = min( - end_of_month.timestamp(), datetime.now(timezone.utc).timestamp() - ) - # Skip runs that started more than 1 hour ago and are still incomplete - if start_time < datetime.now(timezone.utc).timestamp() - 3600: - logger.warning(f"Skipping incomplete run started > 1 hour ago: {run}") - continue - status = "incomplete" - - duration = end_time - start_time - total_seconds += duration - valid_runs += 1 - - # Store run details - run_details.append( - { - "id": run["id"], - "thread_id": run["thread_id"], - "started_at": run["started_at"], - "completed_at": run["completed_at"], - "duration_minutes": duration / 60, - "status": status, - } - ) - - logger.debug(f"Run duration: {duration/60:.2f} minutes") - - logger.info( - f"Processed {valid_runs} valid runs out of {len(runs_result.data)} total runs" - ) - - # Sort runs by duration (longest first) - run_details.sort(key=lambda x: x["duration_minutes"], reverse=True) - - return total_seconds / 60, run_details # Convert to minutes - - -async def main(): - """Main function to run the script.""" - # Parse command line arguments - parser = argparse.ArgumentParser( - description="Get monthly usage for a specific user and month" - ) - parser.add_argument( - "--user-id", type=str, help="User ID to get usage for", required=True - ) - parser.add_argument("--year", type=int, help="Year (e.g., 2024)", required=True) - parser.add_argument("--month", type=int, help="Month (1-12)", required=True) - parser.add_argument( - "--verbose", "-v", action="store_true", help="Enable verbose logging" - ) - - args = parser.parse_args() - - # Validate month - if args.month < 1 or args.month > 12: - raise ValueError("Month must be between 1 and 12") - - try: - # Get user information - try: - user = await get_user(args.user_id) - logger.info(f"User: {user['id']} ({user['email']})") - except Exception as e: - logger.warning(f"Could not fetch user details: {e}") - user = {"id": args.user_id, "email": "unknown"} - - # Get database connection - db = await get_db() - - # Calculate monthly usage - usage_minutes, run_details = await calculate_monthly_usage( - db, args.user_id, args.year, args.month - ) - - # Display results - month_name = datetime(args.year, args.month, 1).strftime("%B") - print(f"\n=== Monthly Usage Report ===") - print(f"User: {user['email']} ({user['id']})") - print(f"Month: {month_name} {args.year}") - print(f"Total Usage: {usage_minutes:.2f} minutes") - print(f"Total Usage: {usage_minutes/60:.2f} hours") - - if usage_minutes > 0: - print(f"Average per day: {usage_minutes/30:.2f} minutes") - - # Display top 10 runs - if run_details: - print(f"\n=== Top Longest Runs ===") - for i, run in enumerate(run_details, 1): - started_at = datetime.fromisoformat( - run["started_at"].replace("Z", "+00:00") - ) - print( - f"{i:2d}. {run['duration_minutes']:6.2f} min | {started_at.strftime('%Y-%m-%d %H:%M')} | {run['status']:10} | Thread: {run['thread_id']} | Run: {run['id']}" - ) - else: - print("\nNo runs found for this period.") - - except Exception as e: - logger.error(f"Error: {e}") - raise e - - finally: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/app/utils/scripts/set_all_customers_active.py b/app/utils/scripts/set_all_customers_active.py deleted file mode 100644 index a64cf75..0000000 --- a/app/utils/scripts/set_all_customers_active.py +++ /dev/null @@ -1,142 +0,0 @@ -#!/usr/bin/env python -""" -Script to set all Stripe customers in the database to active status. - -Usage: - python update_customer_status.py - -This script: -1. Queries all customer IDs from basejump.billing_customers -2. Sets all customers' active field to True in the database - -Make sure your environment variables are properly set: -- SUPABASE_URL -- SUPABASE_SERVICE_ROLE_KEY -""" - -import asyncio -import sys -import os -from typing import List, Dict, Any -from dotenv import load_dotenv - -# Load script-specific environment variables -load_dotenv(".env") - -from services.supabase import DBConnection -from utils.logger import logger - -# Semaphore to limit concurrent database connections -DB_CONNECTION_LIMIT = 20 -db_semaphore = asyncio.Semaphore(DB_CONNECTION_LIMIT) - -# Global DB connection to reuse -db_connection = None - - -async def get_all_customers() -> List[Dict[str, Any]]: - """ - Query all customers from the database. - - Returns: - List of customers with their ID and account_id - """ - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Print the Supabase URL being used - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - # Query all customers from billing_customers - result = await client.schema('basejump').from_('billing_customers').select( - 'id', - 'account_id', - 'active' - ).execute() - - # Print the query result - print(f"Found {len(result.data)} customers in database") - print(result.data) - - if not result.data: - logger.info("No customers found in database") - return [] - - return result.data - - -async def update_all_customers_to_active() -> Dict[str, int]: - """ - Update all customers to active status in the database. - - Returns: - Dict with count of updated customers - """ - try: - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Update all customers to active - result = await client.schema('basejump').from_('billing_customers').update( - {'active': True} - ).filter('id', 'neq', None).execute() - - updated_count = len(result.data) if hasattr(result, 'data') else 0 - logger.info(f"Updated {updated_count} customers to active status") - print(f"Updated {updated_count} customers to active status") - print("Result:", result) - - return {'updated': updated_count} - except Exception as e: - logger.error(f"Error updating customers in database: {str(e)}") - return {'updated': 0, 'error': str(e)} - - -async def main(): - """Main function to run the script.""" - logger.info("Starting customer status update process") - - try: - # Initialize global DB connection - global db_connection - db_connection = DBConnection() - - # Get all customers from the database - customers = await get_all_customers() - - if not customers: - logger.info("No customers to process") - return - - # Ask for confirmation before proceeding - confirm = input(f"\nSet all {len(customers)} customers to active? (y/n): ") - if confirm.lower() != 'y': - logger.info("Operation cancelled by user") - return - - # Update all customers to active - results = await update_all_customers_to_active() - - # Print summary - print("\nCustomer Status Update Summary:") - print(f"Total customers set to active: {results.get('updated', 0)}") - - logger.info("Customer status update completed") - - except Exception as e: - logger.error(f"Error during customer status update: {str(e)}") - sys.exit(1) - finally: - # Clean up database connection - if db_connection: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/app/utils/scripts/update_customer_active_status.py b/app/utils/scripts/update_customer_active_status.py deleted file mode 100644 index ced3ad2..0000000 --- a/app/utils/scripts/update_customer_active_status.py +++ /dev/null @@ -1,326 +0,0 @@ -#!/usr/bin/env python -""" -Script to check Stripe subscriptions for all customers and update their active status. - -Usage: - python update_customer_active_status.py - -This script: -1. Queries all customers from basejump.billing_customers -2. Checks subscription status directly on Stripe using customer_id -3. Updates customer active status in database - -Make sure your environment variables are properly set: -- SUPABASE_URL -- SUPABASE_SERVICE_ROLE_KEY -- STRIPE_SECRET_KEY -""" - -import asyncio -import sys -import os -import time -from typing import List, Dict, Any, Tuple -from dotenv import load_dotenv -import stripe - -# Load script-specific environment variables -load_dotenv(".env") - -# Import relative modules -from services.supabase import DBConnection -from utils.logger import logger -from utils.config import config - -# Initialize Stripe with the API key -stripe.api_key = config.STRIPE_SECRET_KEY - -# Batch size settings -BATCH_SIZE = 100 # Process customers in batches -MAX_CONCURRENCY = 20 # Maximum concurrent Stripe API calls - -# Global DB connection to reuse -db_connection = None - -async def get_all_customers() -> List[Dict[str, Any]]: - """ - Query all customers from the database. - - Returns: - List of customers with their ID (customer_id is used for Stripe) - """ - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Print the Supabase URL being used - print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}") - - # Query all customers from billing_customers - result = await client.schema('basejump').from_('billing_customers').select( - 'id', - 'active' - ).execute() - - # Print the query result - print(f"Found {len(result.data)} customers in database") - - if not result.data: - logger.info("No customers found in database") - return [] - - return result.data - -async def check_stripe_subscription(customer_id: str) -> bool: - """ - Check if a customer has an active subscription directly on Stripe. - - Args: - customer_id: Customer ID (billing_customers.id) which is the Stripe customer ID - - Returns: - True if customer has at least one active subscription, False otherwise - """ - if not customer_id: - print(f"⚠️ Empty customer_id") - return False - - try: - # Print what we're checking for debugging - print(f"Checking Stripe subscriptions for customer: {customer_id}") - - # List all subscriptions for this customer directly on Stripe - subscriptions = stripe.Subscription.list( - customer=customer_id, - status='active', # Only get active subscriptions - limit=1 # We only need to know if there's at least one - ) - - # Print the raw data for debugging - print(f"Stripe returned data: {subscriptions.data}") - - # If there's at least one active subscription, the customer is active - has_active_subscription = len(subscriptions.data) > 0 - - if has_active_subscription: - print(f"✅ Customer {customer_id} has ACTIVE subscription") - else: - print(f"❌ Customer {customer_id} has NO active subscription") - - return has_active_subscription - - except Exception as e: - logger.error(f"Error checking Stripe subscription for customer {customer_id}: {str(e)}") - print(f"⚠️ Error checking subscription for {customer_id}: {str(e)}") - return False - -async def process_customer_batch(batch: List[Dict[str, Any]], batch_number: int, total_batches: int) -> Dict[str, bool]: - """ - Process a batch of customers by checking their Stripe subscriptions concurrently. - - Args: - batch: List of customer records in this batch - batch_number: Current batch number (for logging) - total_batches: Total number of batches (for logging) - - Returns: - Dictionary mapping customer IDs to subscription status (True/False) - """ - start_time = time.time() - batch_size = len(batch) - print(f"Processing batch {batch_number}/{total_batches} ({batch_size} customers)...") - - # Create a semaphore to limit concurrency within the batch to avoid rate limiting - semaphore = asyncio.Semaphore(MAX_CONCURRENCY) - - async def check_single_customer(customer: Dict[str, Any]) -> Tuple[str, bool]: - async with semaphore: # Limit concurrent API calls - customer_id = customer['id'] - - # Check directly on Stripe - customer_id IS the Stripe customer ID - is_active = await check_stripe_subscription(customer_id) - return customer_id, is_active - - # Create tasks for all customers in this batch - tasks = [check_single_customer(customer) for customer in batch] - - # Run all tasks in this batch concurrently - results = await asyncio.gather(*tasks) - - # Convert results to dictionary - subscription_status = {customer_id: status for customer_id, status in results} - - end_time = time.time() - - # Count active/inactive in this batch - active_count = sum(1 for status in subscription_status.values() if status) - inactive_count = batch_size - active_count - - print(f"Batch {batch_number} completed in {end_time - start_time:.2f} seconds") - print(f"Results (batch {batch_number}): {active_count} active, {inactive_count} inactive subscriptions") - - return subscription_status - -async def update_customer_batch(subscription_status: Dict[str, bool]) -> Dict[str, int]: - """ - Update a batch of customers in the database. - - Args: - subscription_status: Dictionary mapping customer IDs to active status - - Returns: - Dictionary with statistics about the update - """ - start_time = time.time() - - global db_connection - if db_connection is None: - db_connection = DBConnection() - - client = await db_connection.client - - # Separate customers into active and inactive groups - active_customers = [cid for cid, status in subscription_status.items() if status] - inactive_customers = [cid for cid, status in subscription_status.items() if not status] - - total_count = len(active_customers) + len(inactive_customers) - - # Update statistics - stats = { - 'total': total_count, - 'active_updated': 0, - 'inactive_updated': 0, - 'errors': 0 - } - - # Update active customers in a single operation - if active_customers: - try: - print(f"Updating {len(active_customers)} customers to ACTIVE status") - await client.schema('basejump').from_('billing_customers').update( - {'active': True} - ).in_('id', active_customers).execute() - - stats['active_updated'] = len(active_customers) - logger.info(f"Updated {len(active_customers)} customers to ACTIVE status") - except Exception as e: - logger.error(f"Error updating active customers: {str(e)}") - stats['errors'] += 1 - - # Update inactive customers in a single operation - if inactive_customers: - try: - print(f"Updating {len(inactive_customers)} customers to INACTIVE status") - await client.schema('basejump').from_('billing_customers').update( - {'active': False} - ).in_('id', inactive_customers).execute() - - stats['inactive_updated'] = len(inactive_customers) - logger.info(f"Updated {len(inactive_customers)} customers to INACTIVE status") - except Exception as e: - logger.error(f"Error updating inactive customers: {str(e)}") - stats['errors'] += 1 - - end_time = time.time() - print(f"Database updates completed in {end_time - start_time:.2f} seconds") - - return stats - -async def main(): - """Main function to run the script.""" - total_start_time = time.time() - logger.info("Starting customer active status update process") - - try: - # Check Stripe API key - print(f"Stripe API key configured: {'Yes' if config.STRIPE_SECRET_KEY else 'No'}") - if not config.STRIPE_SECRET_KEY: - print("ERROR: Stripe API key not configured. Please set STRIPE_SECRET_KEY in your environment.") - return - - # Initialize global DB connection - global db_connection - db_connection = DBConnection() - - # Get all customers from the database - all_customers = await get_all_customers() - - if not all_customers: - logger.info("No customers to process") - return - - # Print a small sample of the customer data - print("\nCustomer data sample (customer_id = Stripe customer ID):") - for i, customer in enumerate(all_customers[:5]): # Show first 5 only - print(f" {i+1}. ID: {customer['id']}, Active: {customer.get('active')}") - if len(all_customers) > 5: - print(f" ... and {len(all_customers) - 5} more") - - # Split customers into batches - batches = [all_customers[i:i + BATCH_SIZE] for i in range(0, len(all_customers), BATCH_SIZE)] - total_batches = len(batches) - - # Ask for confirmation before proceeding - confirm = input(f"\nProcess {len(all_customers)} customers in {total_batches} batches of {BATCH_SIZE}? (y/n): ") - if confirm.lower() != 'y': - logger.info("Operation cancelled by user") - return - - # Overall statistics - all_stats = { - 'total': 0, - 'active_updated': 0, - 'inactive_updated': 0, - 'errors': 0 - } - - # Process each batch - for i, batch in enumerate(batches): - batch_number = i + 1 - - # STEP 1: Process this batch of customers - subscription_status = await process_customer_batch(batch, batch_number, total_batches) - - # STEP 2: Update this batch in the database - batch_stats = await update_customer_batch(subscription_status) - - # Accumulate statistics - all_stats['total'] += batch_stats['total'] - all_stats['active_updated'] += batch_stats['active_updated'] - all_stats['inactive_updated'] += batch_stats['inactive_updated'] - all_stats['errors'] += batch_stats['errors'] - - # Show batch completion - print(f"Completed batch {batch_number}/{total_batches}") - - # Brief pause between batches to avoid Stripe rate limiting - if batch_number < total_batches: - await asyncio.sleep(1) # 1 second pause between batches - - # Print summary - total_end_time = time.time() - total_time = total_end_time - total_start_time - - print("\nCustomer Status Update Summary:") - print(f"Total customers processed: {all_stats['total']}") - print(f"Customers set to active: {all_stats['active_updated']}") - print(f"Customers set to inactive: {all_stats['inactive_updated']}") - if all_stats['errors'] > 0: - print(f"Update errors: {all_stats['errors']}") - print(f"Total processing time: {total_time:.2f} seconds") - - logger.info(f"Customer active status update completed in {total_time:.2f} seconds") - - except Exception as e: - logger.error(f"Error during customer status update: {str(e)}") - sys.exit(1) - finally: - # Clean up database connection - if db_connection: - await DBConnection.disconnect() - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file