mirrors/foundationagents--openmanus
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

feat:support 4 sandbox tools by using daytona

Browse Source
This commit is contained in:
GhostC
2025-07-27 21:59:18 +08:00
parent f99dfac5bb
commit 66f4ae046e
45 changed files with 475 additions and 11129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/agent/sandbox_agent.py
-6
View File
@@ -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
app/agentpress/__init__.py
-1
View File
@@ -1 +0,0 @@
# Utility functions and constants for agent tools
app/agentpress/response_processor.py
-1890
View File
File diff suppressed because it is too large Load Diff
app/agentpress/tool.py
-240
View File
@@ -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='''
<str-replace file_path="path/to/file">
<old_str>text to replace</old_str>
<new_str>replacement text</new_str>
</str-replace>
'''
)
"""
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
app/agentpress/tool_registry.py
-152
View File
@@ -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
app/agentpress/utils/__init__.py
-1
View File
@@ -1 +0,0 @@
# Utils module for AgentPress
app/agentpress/utils/json_helpers.py
-174
View File
@@ -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
app/agentpress/xml_tool_parser.py
-300
View File
@@ -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:
<function_calls>
<invoke name="function_name">
<parameter name="param_name">param_value</parameter>
...
</invoke>
</function_calls>
"""
# Regex patterns for extracting XML blocks
FUNCTION_CALLS_PATTERN = re.compile(
r'<function_calls>(.*?)</function_calls>',
re.DOTALL | re.IGNORECASE
)
INVOKE_PATTERN = re.compile(
r'<invoke\s+name=["\']([^"\']+)["\']>(.*?)</invoke>',
re.DOTALL | re.IGNORECASE
)
PARAMETER_PATTERN = re.compile(
r'<parameter\s+name=["\']([^"\']+)["\']>(.*?)</parameter>',
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'<invoke\s+name=["\']{re.escape(function_name)}["\']>.*?</invoke>',
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"<invoke name=\"{function_name}\">...</invoke>"
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 <tool_name>...</tool_name> or
<tool_name param="value">...</tool_name>
"""
tool_calls = []
# Pattern for finding XML-like tags
tag_pattern = re.compile(r'<([a-zA-Z][\w\-]*)((?:\s+[\w\-]+=["\'][^"\']*["\'])*)\s*>(.*?)</\1>', 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 = ['<function_calls>', '<invoke name="{}">'.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('<parameter name="{}">{}</parameter>'.format(
param_name, value_str
))
lines.extend(['</invoke>', '</function_calls>'])
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)
app/daytona/api.py
+199 -112
View File
@@ -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))
app/daytona/docker/Dockerfile
-133
View File
@@ -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"]
app/daytona/docker/browser_api.py
-2110
View File
File diff suppressed because it is too large Load Diff
app/daytona/docker/docker-compose.yml
-45
View File
@@ -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
app/daytona/docker/entrypoint.sh
-4
View File
@@ -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
app/daytona/docker/requirements.txt
-6
View File
@@ -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
app/daytona/docker/server.py
-29
View File
@@ -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)
app/daytona/docker/supervisord.conf
-94
View File
@@ -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
app/daytona/tool_base.py
-2
View File
@@ -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
app/services/billing.py
-969
View File
@@ -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)}")
app/services/docker/redis.conf
-1
View File
@@ -1 +0,0 @@
timeout 120
app/services/email.py
-192
View File
@@ -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"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Welcome to Kortix Suna</title>
<style>
body {{
font-family: Arial, sans-serif;
background-color: #ffffff;
color: #000000;
margin: 0;
padding: 0;
line-height: 1.6;
}}
.container {{
max-width: 600px;
margin: 40px auto;
padding: 30px;
background-color: #ffffff;
}}
.logo-container {{
text-align: center;
margin-bottom: 30px;
padding: 10px 0;
}}
.logo {{
max-width: 100%;
height: auto;
max-height: 60px;
display: inline-block;
}}
h1 {{
font-size: 24px;
color: #000000;
margin-bottom: 20px;
}}
p {{
margin-bottom: 16px;
}}
a {{
color: #3366cc;
text-decoration: none;
}}
a:hover {{
text-decoration: underline;
}}
.button {{
display: inline-block;
margin-top: 30px;
background-color: #3B82F6;
color: white !important;
padding: 14px 24px;
text-align: center;
text-decoration: none;
font-weight: bold;
border-radius: 6px;
border: none;
}}
.button:hover {{
background-color: #2563EB;
text-decoration: none;
}}
.emoji {{
font-size: 20px;
}}
</style>
</head>
<body>
<div class="container">
<div class="logo-container">
<img src="https://i.postimg.cc/WdNtRx5Z/kortix-suna-logo.png" alt="Kortix Suna Logo" class="logo">
</div>
<h1>Welcome to Kortix Suna!</h1>
<p>Hi {user_name},</p>
<p><em><strong>Welcome to Kortix Suna — we're excited to have you on board!</strong></em></p>
<p>To get started, we'd like to get to know you better: fill out this short <a href="https://docs.google.com/forms/d/e/1FAIpQLSef1EHuqmIh_iQz-kwhjnzSC3Ml-V_5wIySDpMoMU9W_j24JQ/viewform">form</a>!</p>
<p>To celebrate your arrival, here's a <strong>15% discount</strong> to try out the best version of Suna (1 month):</p>
<p>🎁 Use code <strong>WELCOME15</strong> at checkout.</p>
<p>Let us know if you need help getting started or have questions — we're always here, and join our <a href="https://discord.com/invite/FjD644cfcs">Discord community</a>.</p>
<p><strong>For your business:</strong> if you want to automate manual and ordinary tasks for your company, book a call with us <a href="https://cal.com/team/kortix/enterprise-demo">here</a></p>
<p>Thanks again, and welcome to the Suna community <span class="emoji">🌞</span></p>
<p>— The Suna Team</p>
<a href="https://www.suna.so/" class="button">Go to the platform</a>
</div>
</body>
</html>"""
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()
app/services/email_api.py
-70
View File
@@ -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"
)
app/services/langfuse.py
-12
View File
@@ -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)
app/services/llm.py
-411
View File
@@ -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 <timestamp>.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!")
app/services/mcp_custom.py
-129
View File
@@ -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
app/services/mcp_temp.py
-299
View File
@@ -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
app/services/redis.py
-153
View File
@@ -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)
app/services/supabase.py
-113
View File
@@ -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)}")
app/services/transcription.py
-76
View File
@@ -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)}")
app/tool/sb_files_tool.py
+75 -30
View File
@@ -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")
raise NotImplementedError(
"create_with_context not implemented for SandboxFilesTool"
)
app/tool/sb_shell_tool.py
+136 -97
View File
@@ -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()):
app/tool/sb_vision_tool.py
+65 -38
View File
@@ -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)}")
return self.fail_response(f"see_image 执行异常: {str(e)}")
app/utils/auth_utils.py
-231
View File
@@ -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
app/utils/config.py
-253
View File
@@ -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()
app/utils/constants.py
-145
View File
@@ -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
}
app/utils/retry.py
-58
View File
@@ -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")
app/utils/s3_upload_utils.py
-51
View File
@@ -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)}")
app/utils/scripts/archive_inactive_sandboxes.py
-350
View File
@@ -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())
app/utils/scripts/archive_old_sandboxes.py
-344
View File
@@ -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())
app/utils/scripts/copy_project.py
-388
View File
@@ -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())
app/utils/scripts/delete_user_sandboxes.py
-138
View File
@@ -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 <account_id>
"""
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]} <account_id>")
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())
app/utils/scripts/export_import_project.py
-400
View File
@@ -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())
app/utils/scripts/generate_share_links.py
-166
View File
@@ -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())
app/utils/scripts/get_monthly_usage.py
-248
View File
@@ -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 <USER_ID> --year <YEAR> --month <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())
app/utils/scripts/set_all_customers_active.py
-142
View File
@@ -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())
app/utils/scripts/update_customer_active_status.py
-326
View File
@@ -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())