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

modify BaseTool

Browse Source
This commit is contained in:
GhostC
2025-06-30 23:12:44 +08:00
parent 527fd6483b
commit 9513686155
34 changed files with 6572 additions and 2370 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/agentpress/__init__.py
+1
View File
@@ -0,0 +1 @@
# Utility functions and constants for agent tools
app/agentpress/context_manager.py
+298
View File
@@ -0,0 +1,298 @@
"""
Context Management for AgentPress Threads.
This module handles token counting and thread summarization to prevent
reaching the context window limitations of LLM models.
"""
import json
from typing import List, Dict, Any, Optional
from litellm import token_counter, completion_cost
from services.supabase import DBConnection
from services.llm import make_llm_api_call
from utils.logger import logger
# Constants for token management
DEFAULT_TOKEN_THRESHOLD = 120000 # 80k tokens threshold for summarization
SUMMARY_TARGET_TOKENS = 10000 # Target ~10k tokens for the summary message
RESERVE_TOKENS = 5000 # Reserve tokens for new messages
class ContextManager:
"""Manages thread context including token counting and summarization."""
def __init__(self, token_threshold: int = DEFAULT_TOKEN_THRESHOLD):
"""Initialize the ContextManager.
Args:
token_threshold: Token count threshold to trigger summarization
"""
self.db = DBConnection()
self.token_threshold = token_threshold
async def get_thread_token_count(self, thread_id: str) -> int:
"""Get the current token count for a thread using LiteLLM.
Args:
thread_id: ID of the thread to analyze
Returns:
The total token count for relevant messages in the thread
"""
logger.debug(f"Getting token count for thread {thread_id}")
try:
# Get messages for the thread
messages = await self.get_messages_for_summarization(thread_id)
if not messages:
logger.debug(f"No messages found for thread {thread_id}")
return 0
# Use litellm's token_counter for accurate model-specific counting
# This is much more accurate than the SQL-based estimation
token_count = token_counter(model="gpt-4", messages=messages)
logger.info(f"Thread {thread_id} has {token_count} tokens (calculated with litellm)")
return token_count
except Exception as e:
logger.error(f"Error getting token count: {str(e)}")
return 0
async def get_messages_for_summarization(self, thread_id: str) -> List[Dict[str, Any]]:
"""Get all LLM messages from the thread that need to be summarized.
This gets messages after the most recent summary or all messages if
no summary exists. Unlike get_llm_messages, this includes ALL messages
since the last summary, even if we're generating a new summary.
Args:
thread_id: ID of the thread to get messages from
Returns:
List of message objects to summarize
"""
logger.debug(f"Getting messages for summarization for thread {thread_id}")
client = await self.db.client
try:
# Find the most recent summary message
summary_result = await client.table('messages').select('created_at') \
.eq('thread_id', thread_id) \
.eq('type', 'summary') \
.eq('is_llm_message', True) \
.order('created_at', desc=True) \
.limit(1) \
.execute()
# Get messages after the most recent summary or all messages if no summary
if summary_result.data and len(summary_result.data) > 0:
last_summary_time = summary_result.data[0]['created_at']
logger.debug(f"Found last summary at {last_summary_time}")
# Get all messages after the summary, but NOT including the summary itself
messages_result = await client.table('messages').select('*') \
.eq('thread_id', thread_id) \
.eq('is_llm_message', True) \
.gt('created_at', last_summary_time) \
.order('created_at') \
.execute()
else:
logger.debug("No previous summary found, getting all messages")
# Get all messages
messages_result = await client.table('messages').select('*') \
.eq('thread_id', thread_id) \
.eq('is_llm_message', True) \
.order('created_at') \
.execute()
# Parse the message content if needed
messages = []
for msg in messages_result.data:
# Skip existing summary messages - we don't want to summarize summaries
if msg.get('type') == 'summary':
logger.debug(f"Skipping summary message from {msg.get('created_at')}")
continue
# Parse content if it's a string
content = msg['content']
if isinstance(content, str):
try:
content = json.loads(content)
except json.JSONDecodeError:
pass # Keep as string if not valid JSON
# Ensure we have the proper format for the LLM
if 'role' not in content and 'type' in msg:
# Convert message type to role if needed
role = msg['type']
if role == 'assistant' or role == 'user' or role == 'system' or role == 'tool':
content = {'role': role, 'content': content}
messages.append(content)
logger.info(f"Got {len(messages)} messages to summarize for thread {thread_id}")
return messages
except Exception as e:
logger.error(f"Error getting messages for summarization: {str(e)}", exc_info=True)
return []
async def create_summary(
self,
thread_id: str,
messages: List[Dict[str, Any]],
model: str = "gpt-4o-mini"
) -> Optional[Dict[str, Any]]:
"""Generate a summary of conversation messages.
Args:
thread_id: ID of the thread to summarize
messages: Messages to summarize
model: LLM model to use for summarization
Returns:
Summary message object or None if summarization failed
"""
if not messages:
logger.warning("No messages to summarize")
return None
logger.info(f"Creating summary for thread {thread_id} with {len(messages)} messages")
# Create system message with summarization instructions
system_message = {
"role": "system",
"content": f"""You are a specialized summarization assistant. Your task is to create a concise but comprehensive summary of the conversation history.
The summary should:
1. Preserve all key information including decisions, conclusions, and important context
2. Include any tools that were used and their results
3. Maintain chronological order of events
4. Be presented as a narrated list of key points with section headers
5. Include only factual information from the conversation (no new information)
6. Be concise but detailed enough that the conversation can continue with this summary as context
VERY IMPORTANT: This summary will replace older parts of the conversation in the LLM's context window, so ensure it contains ALL key information and LATEST STATE OF THE CONVERSATION - SO WE WILL KNOW HOW TO PICK UP WHERE WE LEFT OFF.
THE CONVERSATION HISTORY TO SUMMARIZE IS AS FOLLOWS:
===============================================================
==================== CONVERSATION HISTORY ====================
{messages}
==================== END OF CONVERSATION HISTORY ====================
===============================================================
"""
}
try:
# Call LLM to generate summary
response = await make_llm_api_call(
model_name=model,
messages=[system_message, {"role": "user", "content": "PLEASE PROVIDE THE SUMMARY NOW."}],
temperature=0,
max_tokens=SUMMARY_TARGET_TOKENS,
stream=False
)
if response and hasattr(response, 'choices') and response.choices:
summary_content = response.choices[0].message.content
# Track token usage
try:
token_count = token_counter(model=model, messages=[{"role": "user", "content": summary_content}])
cost = completion_cost(model=model, prompt="", completion=summary_content)
logger.info(f"Summary generated with {token_count} tokens at cost ${cost:.6f}")
except Exception as e:
logger.error(f"Error calculating token usage: {str(e)}")
# Format the summary message with clear beginning and end markers
formatted_summary = f"""
======== CONVERSATION HISTORY SUMMARY ========
{summary_content}
======== END OF SUMMARY ========
The above is a summary of the conversation history. The conversation continues below.
"""
# Format the summary message
summary_message = {
"role": "user",
"content": formatted_summary
}
return summary_message
else:
logger.error("Failed to generate summary: Invalid response")
return None
except Exception as e:
logger.error(f"Error creating summary: {str(e)}", exc_info=True)
return None
async def check_and_summarize_if_needed(
self,
thread_id: str,
add_message_callback,
model: str = "gpt-4o-mini",
force: bool = False
) -> bool:
"""Check if thread needs summarization and summarize if so.
Args:
thread_id: ID of the thread to check
add_message_callback: Callback to add the summary message to the thread
model: LLM model to use for summarization
force: Whether to force summarization regardless of token count
Returns:
True if summarization was performed, False otherwise
"""
try:
# Get token count using LiteLLM (accurate model-specific counting)
token_count = await self.get_thread_token_count(thread_id)
# If token count is below threshold and not forcing, no summarization needed
if token_count < self.token_threshold and not force:
logger.debug(f"Thread {thread_id} has {token_count} tokens, below threshold {self.token_threshold}")
return False
# Log reason for summarization
if force:
logger.info(f"Forced summarization of thread {thread_id} with {token_count} tokens")
else:
logger.info(f"Thread {thread_id} exceeds token threshold ({token_count} >= {self.token_threshold}), summarizing...")
# Get messages to summarize
messages = await self.get_messages_for_summarization(thread_id)
# If there are too few messages, don't summarize
if len(messages) < 3:
logger.info(f"Thread {thread_id} has too few messages ({len(messages)}) to summarize")
return False
# Create summary
summary = await self.create_summary(thread_id, messages, model)
if summary:
# Add summary message to thread
await add_message_callback(
thread_id=thread_id,
type="summary",
content=summary,
is_llm_message=True,
metadata={"token_count": token_count}
)
logger.info(f"Successfully added summary to thread {thread_id}")
return True
else:
logger.error(f"Failed to create summary for thread {thread_id}")
return False
except Exception as e:
logger.error(f"Error in check_and_summarize_if_needed: {str(e)}", exc_info=True)
return False
app/agentpress/response_processor.py
+1890
View File
File diff suppressed because it is too large Load Diff
app/agentpress/thread_manager.py
+787
View File
@@ -0,0 +1,787 @@
"""
Conversation thread management system for AgentPress.
This module provides comprehensive conversation management, including:
- Thread creation and persistence
- Message handling with support for text and images
- Tool registration and execution
- LLM interaction with streaming support
- Error handling and cleanup
- Context summarization to manage token limits
"""
import json
from typing import List, Dict, Any, Optional, Type, Union, AsyncGenerator, Literal
from services.llm import make_llm_api_call
from agentpress.tool import Tool
from agentpress.tool_registry import ToolRegistry
from agentpress.context_manager import ContextManager
from agentpress.response_processor import (
ResponseProcessor,
ProcessorConfig
)
from services.supabase import DBConnection
from utils.logger import logger
from langfuse.client import StatefulGenerationClient, StatefulTraceClient
from services.langfuse import langfuse
import datetime
from litellm import token_counter
# Type alias for tool choice
ToolChoice = Literal["auto", "required", "none"]
class ThreadManager:
"""Manages conversation threads with LLM models and tool execution.
Provides comprehensive conversation management, handling message threading,
tool registration, and LLM interactions with support for both standard and
XML-based tool execution patterns.
"""
def __init__(self, trace: Optional[StatefulTraceClient] = None, is_agent_builder: bool = False, target_agent_id: Optional[str] = None, agent_config: Optional[dict] = None):
"""Initialize ThreadManager.
Args:
trace: Optional trace client for logging
is_agent_builder: Whether this is an agent builder session
target_agent_id: ID of the agent being built (if in agent builder mode)
agent_config: Optional agent configuration with version information
"""
self.db = DBConnection()
self.tool_registry = ToolRegistry()
self.trace = trace
self.is_agent_builder = is_agent_builder
self.target_agent_id = target_agent_id
self.agent_config = agent_config
if not self.trace:
self.trace = langfuse.trace(name="anonymous:thread_manager")
self.response_processor = ResponseProcessor(
tool_registry=self.tool_registry,
add_message_callback=self.add_message,
trace=self.trace,
is_agent_builder=self.is_agent_builder,
target_agent_id=self.target_agent_id,
agent_config=self.agent_config
)
self.context_manager = ContextManager()
def _is_tool_result_message(self, msg: Dict[str, Any]) -> bool:
if not ("content" in msg and msg['content']):
return False
content = msg['content']
if isinstance(content, str) and "ToolResult" in content: return True
if isinstance(content, dict) and "tool_execution" in content: return True
if isinstance(content, dict) and "interactive_elements" in content: return True
if isinstance(content, str):
try:
parsed_content = json.loads(content)
if isinstance(parsed_content, dict) and "tool_execution" in parsed_content: return True
if isinstance(parsed_content, dict) and "interactive_elements" in content: return True
except (json.JSONDecodeError, TypeError):
pass
return False
def _compress_message(self, msg_content: Union[str, dict], message_id: Optional[str] = None, max_length: int = 3000) -> Union[str, dict]:
"""Compress the message content."""
# print("max_length", max_length)
if isinstance(msg_content, str):
if len(msg_content) > max_length:
return msg_content[:max_length] + "... (truncated)" + f"\n\nmessage_id \"{message_id}\"\nUse expand-message tool to see contents"
else:
return msg_content
elif isinstance(msg_content, dict):
if len(json.dumps(msg_content)) > max_length:
return json.dumps(msg_content)[:max_length] + "... (truncated)" + f"\n\nmessage_id \"{message_id}\"\nUse expand-message tool to see contents"
else:
return msg_content
def _safe_truncate(self, msg_content: Union[str, dict], max_length: int = 100000) -> Union[str, dict]:
"""Truncate the message content safely by removing the middle portion."""
max_length = min(max_length, 100000)
if isinstance(msg_content, str):
if len(msg_content) > max_length:
# Calculate how much to keep from start and end
keep_length = max_length - 150 # Reserve space for truncation message
start_length = keep_length // 2
end_length = keep_length - start_length
start_part = msg_content[:start_length]
end_part = msg_content[-end_length:] if end_length > 0 else ""
return start_part + f"\n\n... (middle truncated) ...\n\n" + end_part + f"\n\nThis message is too long, repeat relevant information in your response to remember it"
else:
return msg_content
elif isinstance(msg_content, dict):
json_str = json.dumps(msg_content)
if len(json_str) > max_length:
# Calculate how much to keep from start and end
keep_length = max_length - 150 # Reserve space for truncation message
start_length = keep_length // 2
end_length = keep_length - start_length
start_part = json_str[:start_length]
end_part = json_str[-end_length:] if end_length > 0 else ""
return start_part + f"\n\n... (middle truncated) ...\n\n" + end_part + f"\n\nThis message is too long, repeat relevant information in your response to remember it"
else:
return msg_content
def _compress_tool_result_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]:
"""Compress the tool result messages except the most recent one."""
uncompressed_total_token_count = token_counter(model=llm_model, messages=messages)
if uncompressed_total_token_count > (max_tokens or (100 * 1000)):
_i = 0 # Count the number of ToolResult messages
for msg in reversed(messages): # Start from the end and work backwards
if self._is_tool_result_message(msg): # Only compress ToolResult messages
_i += 1 # Count the number of ToolResult messages
msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message
if msg_token_count > token_threshold: # If the message is too long
if _i > 1: # If this is not the most recent ToolResult message
message_id = msg.get('message_id') # Get the message_id
if message_id:
msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3)
else:
logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}")
else:
msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2))
return messages
def _compress_user_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]:
"""Compress the user messages except the most recent one."""
uncompressed_total_token_count = token_counter(model=llm_model, messages=messages)
if uncompressed_total_token_count > (max_tokens or (100 * 1000)):
_i = 0 # Count the number of User messages
for msg in reversed(messages): # Start from the end and work backwards
if msg.get('role') == 'user': # Only compress User messages
_i += 1 # Count the number of User messages
msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message
if msg_token_count > token_threshold: # If the message is too long
if _i > 1: # If this is not the most recent User message
message_id = msg.get('message_id') # Get the message_id
if message_id:
msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3)
else:
logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}")
else:
msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2))
return messages
def _compress_assistant_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int], token_threshold: Optional[int] = 1000) -> List[Dict[str, Any]]:
"""Compress the assistant messages except the most recent one."""
uncompressed_total_token_count = token_counter(model=llm_model, messages=messages)
if uncompressed_total_token_count > (max_tokens or (100 * 1000)):
_i = 0 # Count the number of Assistant messages
for msg in reversed(messages): # Start from the end and work backwards
if msg.get('role') == 'assistant': # Only compress Assistant messages
_i += 1 # Count the number of Assistant messages
msg_token_count = token_counter(messages=[msg]) # Count the number of tokens in the message
if msg_token_count > token_threshold: # If the message is too long
if _i > 1: # If this is not the most recent Assistant message
message_id = msg.get('message_id') # Get the message_id
if message_id:
msg["content"] = self._compress_message(msg["content"], message_id, token_threshold * 3)
else:
logger.warning(f"UNEXPECTED: Message has no message_id {str(msg)[:100]}")
else:
msg["content"] = self._safe_truncate(msg["content"], int(max_tokens * 2))
return messages
def _remove_meta_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Remove meta messages from the messages."""
result: List[Dict[str, Any]] = []
for msg in messages:
msg_content = msg.get('content')
# Try to parse msg_content as JSON if it's a string
if isinstance(msg_content, str):
try: msg_content = json.loads(msg_content)
except json.JSONDecodeError: pass
if isinstance(msg_content, dict):
# Create a copy to avoid modifying the original
msg_content_copy = msg_content.copy()
if "tool_execution" in msg_content_copy:
tool_execution = msg_content_copy["tool_execution"].copy()
if "arguments" in tool_execution:
del tool_execution["arguments"]
msg_content_copy["tool_execution"] = tool_execution
# Create a new message dict with the modified content
new_msg = msg.copy()
new_msg["content"] = json.dumps(msg_content_copy)
result.append(new_msg)
else:
result.append(msg)
return result
def _compress_messages(self, messages: List[Dict[str, Any]], llm_model: str, max_tokens: Optional[int] = 41000, token_threshold: Optional[int] = 4096, max_iterations: int = 5) -> List[Dict[str, Any]]:
"""Compress the messages.
token_threshold: must be a power of 2
"""
if 'sonnet' in llm_model.lower():
max_tokens = 200 * 1000 - 64000 - 28000
elif 'gpt' in llm_model.lower():
max_tokens = 128 * 1000 - 28000
elif 'gemini' in llm_model.lower():
max_tokens = 1000 * 1000 - 300000
elif 'deepseek' in llm_model.lower():
max_tokens = 128 * 1000 - 28000
else:
max_tokens = 41 * 1000 - 10000
result = messages
result = self._remove_meta_messages(result)
uncompressed_total_token_count = token_counter(model=llm_model, messages=result)
result = self._compress_tool_result_messages(result, llm_model, max_tokens, token_threshold)
result = self._compress_user_messages(result, llm_model, max_tokens, token_threshold)
result = self._compress_assistant_messages(result, llm_model, max_tokens, token_threshold)
compressed_token_count = token_counter(model=llm_model, messages=result)
logger.info(f"_compress_messages: {uncompressed_total_token_count} -> {compressed_token_count}") # Log the token compression for debugging later
if max_iterations <= 0:
logger.warning(f"_compress_messages: Max iterations reached, omitting messages")
result = self._compress_messages_by_omitting_messages(messages, llm_model, max_tokens)
return result
if (compressed_token_count > max_tokens):
logger.warning(f"Further token compression is needed: {compressed_token_count} > {max_tokens}")
result = self._compress_messages(messages, llm_model, max_tokens, int(token_threshold / 2), max_iterations - 1)
return self._middle_out_messages(result)
def _compress_messages_by_omitting_messages(
self,
messages: List[Dict[str, Any]],
llm_model: str,
max_tokens: Optional[int] = 41000,
removal_batch_size: int = 10,
min_messages_to_keep: int = 10
) -> List[Dict[str, Any]]:
"""Compress the messages by omitting messages from the middle.
Args:
messages: List of messages to compress
llm_model: Model name for token counting
max_tokens: Maximum allowed tokens
removal_batch_size: Number of messages to remove per iteration
min_messages_to_keep: Minimum number of messages to preserve
"""
if not messages:
return messages
result = messages
result = self._remove_meta_messages(result)
# Early exit if no compression needed
initial_token_count = token_counter(model=llm_model, messages=result)
max_allowed_tokens = max_tokens or (100 * 1000)
if initial_token_count <= max_allowed_tokens:
return result
# Separate system message (assumed to be first) from conversation messages
system_message = messages[0] if messages and messages[0].get('role') == 'system' else None
conversation_messages = result[1:] if system_message else result
safety_limit = 500
current_token_count = initial_token_count
while current_token_count > max_allowed_tokens and safety_limit > 0:
safety_limit -= 1
if len(conversation_messages) <= min_messages_to_keep:
logger.warning(f"Cannot compress further: only {len(conversation_messages)} messages remain (min: {min_messages_to_keep})")
break
# Calculate removal strategy based on current message count
if len(conversation_messages) > (removal_batch_size * 2):
# Remove from middle, keeping recent and early context
middle_start = len(conversation_messages) // 2 - (removal_batch_size // 2)
middle_end = middle_start + removal_batch_size
conversation_messages = conversation_messages[:middle_start] + conversation_messages[middle_end:]
else:
# Remove from earlier messages, preserving recent context
messages_to_remove = min(removal_batch_size, len(conversation_messages) // 2)
if messages_to_remove > 0:
conversation_messages = conversation_messages[messages_to_remove:]
else:
# Can't remove any more messages
break
# Recalculate token count
messages_to_count = ([system_message] + conversation_messages) if system_message else conversation_messages
current_token_count = token_counter(model=llm_model, messages=messages_to_count)
# Prepare final result
final_messages = ([system_message] + conversation_messages) if system_message else conversation_messages
final_token_count = token_counter(model=llm_model, messages=final_messages)
logger.info(f"_compress_messages_by_omitting_messages: {initial_token_count} -> {final_token_count} tokens ({len(messages)} -> {len(final_messages)} messages)")
return final_messages
def _middle_out_messages(self, messages: List[Dict[str, Any]], max_messages: int = 320) -> List[Dict[str, Any]]:
"""Remove messages from the middle of the list, keeping max_messages total."""
if len(messages) <= max_messages:
return messages
# Keep half from the beginning and half from the end
keep_start = max_messages // 2
keep_end = max_messages - keep_start
return messages[:keep_start] + messages[-keep_end:]
def add_tool(self, tool_class: Type[Tool], function_names: Optional[List[str]] = None, **kwargs):
"""Add a tool to the ThreadManager."""
self.tool_registry.register_tool(tool_class, function_names, **kwargs)
async def add_message(
self,
thread_id: str,
type: str,
content: Union[Dict[str, Any], List[Any], str],
is_llm_message: bool = False,
metadata: Optional[Dict[str, Any]] = None,
agent_id: Optional[str] = None,
agent_version_id: Optional[str] = None
):
"""Add a message to the thread in the database.
Args:
thread_id: The ID of the thread to add the message to.
type: The type of the message (e.g., 'text', 'image_url', 'tool_call', 'tool', 'user', 'assistant').
content: The content of the message. Can be a dictionary, list, or string.
It will be stored as JSONB in the database.
is_llm_message: Flag indicating if the message originated from the LLM.
Defaults to False (user message).
metadata: Optional dictionary for additional message metadata.
Defaults to None, stored as an empty JSONB object if None.
agent_id: Optional ID of the agent associated with this message.
agent_version_id: Optional ID of the specific agent version used.
"""
logger.debug(f"Adding message of type '{type}' to thread {thread_id} (agent: {agent_id}, version: {agent_version_id})")
client = await self.db.client
# Prepare data for insertion
data_to_insert = {
'thread_id': thread_id,
'type': type,
'content': content,
'is_llm_message': is_llm_message,
'metadata': metadata or {},
}
# Add agent information if provided
if agent_id:
data_to_insert['agent_id'] = agent_id
if agent_version_id:
data_to_insert['agent_version_id'] = agent_version_id
try:
# Add returning='representation' to get the inserted row data including the id
result = await client.table('messages').insert(data_to_insert, returning='representation').execute()
logger.info(f"Successfully added message to thread {thread_id}")
if result.data and len(result.data) > 0 and isinstance(result.data[0], dict) and 'message_id' in result.data[0]:
return result.data[0]
else:
logger.error(f"Insert operation failed or did not return expected data structure for thread {thread_id}. Result data: {result.data}")
return None
except Exception as e:
logger.error(f"Failed to add message to thread {thread_id}: {str(e)}", exc_info=True)
raise
async def get_llm_messages(self, thread_id: str) -> List[Dict[str, Any]]:
"""Get all messages for a thread.
This method uses the SQL function which handles context truncation
by considering summary messages.
Args:
thread_id: The ID of the thread to get messages for.
Returns:
List of message objects.
"""
logger.debug(f"Getting messages for thread {thread_id}")
client = await self.db.client
try:
# result = await client.rpc('get_llm_formatted_messages', {'p_thread_id': thread_id}).execute()
# Fetch messages in batches of 1000 to avoid overloading the database
all_messages = []
batch_size = 1000
offset = 0
while True:
result = await client.table('messages').select('message_id, content').eq('thread_id', thread_id).eq('is_llm_message', True).order('created_at').range(offset, offset + batch_size - 1).execute()
if not result.data or len(result.data) == 0:
break
all_messages.extend(result.data)
# If we got fewer than batch_size records, we've reached the end
if len(result.data) < batch_size:
break
offset += batch_size
# Use all_messages instead of result.data in the rest of the method
result_data = all_messages
# Parse the returned data which might be stringified JSON
if not result_data:
return []
# Return properly parsed JSON objects
messages = []
for item in result_data:
if isinstance(item['content'], str):
try:
parsed_item = json.loads(item['content'])
parsed_item['message_id'] = item['message_id']
messages.append(parsed_item)
except json.JSONDecodeError:
logger.error(f"Failed to parse message: {item['content']}")
else:
content = item['content']
content['message_id'] = item['message_id']
messages.append(content)
return messages
except Exception as e:
logger.error(f"Failed to get messages for thread {thread_id}: {str(e)}", exc_info=True)
return []
async def run_thread(
self,
thread_id: str,
system_prompt: Dict[str, Any],
stream: bool = True,
temporary_message: Optional[Dict[str, Any]] = None,
llm_model: str = "gpt-4o",
llm_temperature: float = 0,
llm_max_tokens: Optional[int] = None,
processor_config: Optional[ProcessorConfig] = None,
tool_choice: ToolChoice = "auto",
native_max_auto_continues: int = 25,
max_xml_tool_calls: int = 0,
include_xml_examples: bool = False,
enable_thinking: Optional[bool] = False,
reasoning_effort: Optional[str] = 'low',
enable_context_manager: bool = True,
generation: Optional[StatefulGenerationClient] = None,
) -> Union[Dict[str, Any], AsyncGenerator]:
"""Run a conversation thread with LLM integration and tool execution.
Args:
thread_id: The ID of the thread to run
system_prompt: System message to set the assistant's behavior
stream: Use streaming API for the LLM response
temporary_message: Optional temporary user message for this run only
llm_model: The name of the LLM model to use
llm_temperature: Temperature parameter for response randomness (0-1)
llm_max_tokens: Maximum tokens in the LLM response
processor_config: Configuration for the response processor
tool_choice: Tool choice preference ("auto", "required", "none")
native_max_auto_continues: Maximum number of automatic continuations when
finish_reason="tool_calls" (0 disables auto-continue)
max_xml_tool_calls: Maximum number of XML tool calls to allow (0 = no limit)
include_xml_examples: Whether to include XML tool examples in the system prompt
enable_thinking: Whether to enable thinking before making a decision
reasoning_effort: The effort level for reasoning
enable_context_manager: Whether to enable automatic context summarization.
Returns:
An async generator yielding response chunks or error dict
"""
logger.info(f"Starting thread execution for thread {thread_id}")
logger.info(f"Using model: {llm_model}")
# Log parameters
logger.info(f"Parameters: model={llm_model}, temperature={llm_temperature}, max_tokens={llm_max_tokens}")
logger.info(f"Auto-continue: max={native_max_auto_continues}, XML tool limit={max_xml_tool_calls}")
# Log model info
logger.info(f"🤖 Thread {thread_id}: Using model {llm_model}")
# Apply max_xml_tool_calls if specified and not already set in config
if max_xml_tool_calls > 0 and not processor_config.max_xml_tool_calls:
processor_config.max_xml_tool_calls = max_xml_tool_calls
# Create a working copy of the system prompt to potentially modify
working_system_prompt = system_prompt.copy()
# Add XML examples to system prompt if requested, do this only ONCE before the loop
if include_xml_examples and processor_config.xml_tool_calling:
xml_examples = self.tool_registry.get_xml_examples()
if xml_examples:
examples_content = """
--- XML TOOL CALLING ---
In this environment you have access to a set of tools you can use to answer the user's question. The tools are specified in XML format.
Format your tool calls using the specified XML tags. Place parameters marked as 'attribute' within the opening tag (e.g., `<tag attribute='value'>`). Place parameters marked as 'content' between the opening and closing tags. Place parameters marked as 'element' within their own child tags (e.g., `<tag><element>value</element></tag>`). Refer to the examples provided below for the exact structure of each tool.
String and scalar parameters should be specified as attributes, while content goes between tags.
Note that spaces for string values are not stripped. The output is parsed with regular expressions.
Here are the XML tools available with examples:
"""
for tag_name, example in xml_examples.items():
examples_content += f"<{tag_name}> Example: {example}\\n"
# # Save examples content to a file
# try:
# with open('xml_examples.txt', 'w') as f:
# f.write(examples_content)
# logger.debug("Saved XML examples to xml_examples.txt")
# except Exception as e:
# logger.error(f"Failed to save XML examples to file: {e}")
system_content = working_system_prompt.get('content')
if isinstance(system_content, str):
working_system_prompt['content'] += examples_content
logger.debug("Appended XML examples to string system prompt content.")
elif isinstance(system_content, list):
appended = False
for item in working_system_prompt['content']: # Modify the copy
if isinstance(item, dict) and item.get('type') == 'text' and 'text' in item:
item['text'] += examples_content
logger.debug("Appended XML examples to the first text block in list system prompt content.")
appended = True
break
if not appended:
logger.warning("System prompt content is a list but no text block found to append XML examples.")
else:
logger.warning(f"System prompt content is of unexpected type ({type(system_content)}), cannot add XML examples.")
# Control whether we need to auto-continue due to tool_calls finish reason
auto_continue = True
auto_continue_count = 0
# Define inner function to handle a single run
async def _run_once(temp_msg=None):
try:
# Ensure processor_config is available in this scope
nonlocal processor_config
# Note: processor_config is now guaranteed to exist due to check above
# 1. Get messages from thread for LLM call
messages = await self.get_llm_messages(thread_id)
# 2. Check token count before proceeding
token_count = 0
try:
# Use the potentially modified working_system_prompt for token counting
token_count = token_counter(model=llm_model, messages=[working_system_prompt] + messages)
token_threshold = self.context_manager.token_threshold
logger.info(f"Thread {thread_id} token count: {token_count}/{token_threshold} ({(token_count/token_threshold)*100:.1f}%)")
# if token_count >= token_threshold and enable_context_manager:
# logger.info(f"Thread token count ({token_count}) exceeds threshold ({token_threshold}), summarizing...")
# summarized = await self.context_manager.check_and_summarize_if_needed(
# thread_id=thread_id,
# add_message_callback=self.add_message,
# model=llm_model,
# force=True
# )
# if summarized:
# logger.info("Summarization complete, fetching updated messages with summary")
# messages = await self.get_llm_messages(thread_id)
# # Recount tokens after summarization, using the modified prompt
# new_token_count = token_counter(model=llm_model, messages=[working_system_prompt] + messages)
# logger.info(f"After summarization: token count reduced from {token_count} to {new_token_count}")
# else:
# logger.warning("Summarization failed or wasn't needed - proceeding with original messages")
# elif not enable_context_manager:
# logger.info("Automatic summarization disabled. Skipping token count check and summarization.")
except Exception as e:
logger.error(f"Error counting tokens or summarizing: {str(e)}")
# 3. Prepare messages for LLM call + add temporary message if it exists
# Use the working_system_prompt which may contain the XML examples
prepared_messages = [working_system_prompt]
# Find the last user message index
last_user_index = -1
for i, msg in enumerate(messages):
if msg.get('role') == 'user':
last_user_index = i
# Insert temporary message before the last user message if it exists
if temp_msg and last_user_index >= 0:
prepared_messages.extend(messages[:last_user_index])
prepared_messages.append(temp_msg)
prepared_messages.extend(messages[last_user_index:])
logger.debug("Added temporary message before the last user message")
else:
# If no user message or no temporary message, just add all messages
prepared_messages.extend(messages)
if temp_msg:
prepared_messages.append(temp_msg)
logger.debug("Added temporary message to the end of prepared messages")
# 4. Prepare tools for LLM call
openapi_tool_schemas = None
if processor_config.native_tool_calling:
openapi_tool_schemas = self.tool_registry.get_openapi_schemas()
logger.debug(f"Retrieved {len(openapi_tool_schemas) if openapi_tool_schemas else 0} OpenAPI tool schemas")
prepared_messages = self._compress_messages(prepared_messages, llm_model)
# 5. Make LLM API call
logger.debug("Making LLM API call")
try:
if generation:
generation.update(
input=prepared_messages,
start_time=datetime.datetime.now(datetime.timezone.utc),
model=llm_model,
model_parameters={
"max_tokens": llm_max_tokens,
"temperature": llm_temperature,
"enable_thinking": enable_thinking,
"reasoning_effort": reasoning_effort,
"tool_choice": tool_choice,
"tools": openapi_tool_schemas,
}
)
llm_response = await make_llm_api_call(
prepared_messages, # Pass the potentially modified messages
llm_model,
temperature=llm_temperature,
max_tokens=llm_max_tokens,
tools=openapi_tool_schemas,
tool_choice=tool_choice if processor_config.native_tool_calling else None,
stream=stream,
enable_thinking=enable_thinking,
reasoning_effort=reasoning_effort
)
logger.debug("Successfully received raw LLM API response stream/object")
except Exception as e:
logger.error(f"Failed to make LLM API call: {str(e)}", exc_info=True)
raise
# 6. Process LLM response using the ResponseProcessor
if stream:
logger.debug("Processing streaming response")
response_generator = self.response_processor.process_streaming_response(
llm_response=llm_response,
thread_id=thread_id,
config=processor_config,
prompt_messages=prepared_messages,
llm_model=llm_model,
)
return response_generator
else:
logger.debug("Processing non-streaming response")
# Pass through the response generator without try/except to let errors propagate up
response_generator = self.response_processor.process_non_streaming_response(
llm_response=llm_response,
thread_id=thread_id,
config=processor_config,
prompt_messages=prepared_messages,
llm_model=llm_model,
)
return response_generator # Return the generator
except Exception as e:
logger.error(f"Error in run_thread: {str(e)}", exc_info=True)
# Return the error as a dict to be handled by the caller
return {
"type": "status",
"status": "error",
"message": str(e)
}
# Define a wrapper generator that handles auto-continue logic
async def auto_continue_wrapper():
nonlocal auto_continue, auto_continue_count
while auto_continue and (native_max_auto_continues == 0 or auto_continue_count < native_max_auto_continues):
# Reset auto_continue for this iteration
auto_continue = False
# Run the thread once, passing the potentially modified system prompt
# Pass temp_msg only on the first iteration
try:
response_gen = await _run_once(temporary_message if auto_continue_count == 0 else None)
# Handle error responses
if isinstance(response_gen, dict) and "status" in response_gen and response_gen["status"] == "error":
logger.error(f"Error in auto_continue_wrapper: {response_gen.get('message', 'Unknown error')}")
yield response_gen
return # Exit the generator on error
# Process each chunk
try:
async for chunk in response_gen:
# Check if this is a finish reason chunk with tool_calls or xml_tool_limit_reached
if chunk.get('type') == 'finish':
if chunk.get('finish_reason') == 'tool_calls':
# Only auto-continue if enabled (max > 0)
if native_max_auto_continues > 0:
logger.info(f"Detected finish_reason='tool_calls', auto-continuing ({auto_continue_count + 1}/{native_max_auto_continues})")
auto_continue = True
auto_continue_count += 1
# Don't yield the finish chunk to avoid confusing the client
continue
elif chunk.get('finish_reason') == 'xml_tool_limit_reached':
# Don't auto-continue if XML tool limit was reached
logger.info(f"Detected finish_reason='xml_tool_limit_reached', stopping auto-continue")
auto_continue = False
# Still yield the chunk to inform the client
# Otherwise just yield the chunk normally
yield chunk
# If not auto-continuing, we're done
if not auto_continue:
break
except Exception as e:
# If there's an exception, log it, yield an error status, and stop execution
logger.error(f"Error in auto_continue_wrapper generator: {str(e)}", exc_info=True)
yield {
"type": "status",
"status": "error",
"message": f"Error in thread processing: {str(e)}"
}
return # Exit the generator on any error
except Exception as outer_e:
# Catch exceptions from _run_once itself
logger.error(f"Error executing thread: {str(outer_e)}", exc_info=True)
yield {
"type": "status",
"status": "error",
"message": f"Error executing thread: {str(outer_e)}"
}
return # Exit immediately on exception from _run_once
# If we've reached the max auto-continues, log a warning
if auto_continue and auto_continue_count >= native_max_auto_continues:
logger.warning(f"Reached maximum auto-continue limit ({native_max_auto_continues}), stopping.")
yield {
"type": "content",
"content": f"\n[Agent reached maximum auto-continue limit of {native_max_auto_continues}]"
}
# If auto-continue is disabled (max=0), just run once
if native_max_auto_continues == 0:
logger.info("Auto-continue is disabled (native_max_auto_continues=0)")
# Pass the potentially modified system prompt and temp message
return await _run_once(temporary_message)
# Otherwise return the auto-continue wrapper generator
return auto_continue_wrapper()
app/agentpress/tool.py
+240
View File
@@ -0,0 +1,240 @@
"""
Core tool system providing the foundation for creating and managing tools.
This module defines the base classes and decorators for creating tools in AgentPress:
- Tool base class for implementing tool functionality
- Schema decorators for OpenAPI and XML tool definitions
- Result containers for standardized tool outputs
"""
from typing import Dict, Any, Union, Optional, List
from dataclasses import dataclass, field
from abc import ABC
import json
import inspect
from enum import Enum
from utils.logger import logger
class SchemaType(Enum):
"""Enumeration of supported schema types for tool definitions."""
OPENAPI = "openapi"
XML = "xml"
CUSTOM = "custom"
@dataclass
class XMLNodeMapping:
"""Maps an XML node to a function parameter.
Attributes:
param_name (str): Name of the function parameter
node_type (str): Type of node ("element", "attribute", or "content")
path (str): XPath-like path to the node ("." means root element)
required (bool): Whether the parameter is required (defaults to True)
"""
param_name: str
node_type: str = "element"
path: str = "."
required: bool = True
@dataclass
class XMLTagSchema:
"""Schema definition for XML tool tags.
Attributes:
tag_name (str): Root tag name for the tool
mappings (List[XMLNodeMapping]): Parameter mappings for the tag
example (str, optional): Example showing tag usage
Methods:
add_mapping: Add a new parameter mapping to the schema
"""
tag_name: str
mappings: List[XMLNodeMapping] = field(default_factory=list)
example: Optional[str] = None
def add_mapping(self, param_name: str, node_type: str = "element", path: str = ".", required: bool = True) -> None:
"""Add a new node mapping to the schema.
Args:
param_name: Name of the function parameter
node_type: Type of node ("element", "attribute", or "content")
path: XPath-like path to the node
required: Whether the parameter is required
"""
self.mappings.append(XMLNodeMapping(
param_name=param_name,
node_type=node_type,
path=path,
required=required
))
logger.debug(f"Added XML mapping for parameter '{param_name}' with type '{node_type}' at path '{path}', required={required}")
@dataclass
class ToolSchema:
"""Container for tool schemas with type information.
Attributes:
schema_type (SchemaType): Type of schema (OpenAPI, XML, or Custom)
schema (Dict[str, Any]): The actual schema definition
xml_schema (XMLTagSchema, optional): XML-specific schema if applicable
"""
schema_type: SchemaType
schema: Dict[str, Any]
xml_schema: Optional[XMLTagSchema] = None
@dataclass
class ToolResult:
"""Container for tool execution results.
Attributes:
success (bool): Whether the tool execution succeeded
output (str): Output message or error description
"""
success: bool
output: str
class Tool(ABC):
"""Abstract base class for all tools.
Provides the foundation for implementing tools with schema registration
and result handling capabilities.
Attributes:
_schemas (Dict[str, List[ToolSchema]]): Registered schemas for tool methods
Methods:
get_schemas: Get all registered tool schemas
success_response: Create a successful result
fail_response: Create a failed result
"""
def __init__(self):
"""Initialize tool with empty schema registry."""
self._schemas: Dict[str, List[ToolSchema]] = {}
logger.debug(f"Initializing tool class: {self.__class__.__name__}")
self._register_schemas()
def _register_schemas(self):
"""Register schemas from all decorated methods."""
for name, method in inspect.getmembers(self, predicate=inspect.ismethod):
if hasattr(method, 'tool_schemas'):
self._schemas[name] = method.tool_schemas
logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}")
def get_schemas(self) -> Dict[str, List[ToolSchema]]:
"""Get all registered tool schemas.
Returns:
Dict mapping method names to their schema definitions
"""
return self._schemas
def success_response(self, data: Union[Dict[str, Any], str]) -> ToolResult:
"""Create a successful tool result.
Args:
data: Result data (dictionary or string)
Returns:
ToolResult with success=True and formatted output
"""
if isinstance(data, str):
text = data
else:
text = json.dumps(data, indent=2)
logger.debug(f"Created success response for {self.__class__.__name__}")
return ToolResult(success=True, output=text)
def fail_response(self, msg: str) -> ToolResult:
"""Create a failed tool result.
Args:
msg: Error message describing the failure
Returns:
ToolResult with success=False and error message
"""
logger.debug(f"Tool {self.__class__.__name__} returned failed result: {msg}")
return ToolResult(success=False, output=msg)
def _add_schema(func, schema: ToolSchema):
"""Helper to add schema to a function."""
if not hasattr(func, 'tool_schemas'):
func.tool_schemas = []
func.tool_schemas.append(schema)
logger.debug(f"Added {schema.schema_type.value} schema to function {func.__name__}")
return func
def openapi_schema(schema: Dict[str, Any]):
"""Decorator for OpenAPI schema tools."""
def decorator(func):
logger.debug(f"Applying OpenAPI schema to function {func.__name__}")
return _add_schema(func, ToolSchema(
schema_type=SchemaType.OPENAPI,
schema=schema
))
return decorator
def xml_schema(
tag_name: str,
mappings: List[Dict[str, Any]] = None,
example: str = None
):
"""
Decorator for XML schema tools with improved node mapping.
Args:
tag_name: Name of the root XML tag
mappings: List of mapping definitions, each containing:
- param_name: Name of the function parameter
- node_type: "element", "attribute", or "content"
- path: Path to the node (default "." for root)
- required: Whether the parameter is required (default True)
example: Optional example showing how to use the XML tag
Example:
@xml_schema(
tag_name="str-replace",
mappings=[
{"param_name": "file_path", "node_type": "attribute", "path": "."},
{"param_name": "old_str", "node_type": "element", "path": "old_str"},
{"param_name": "new_str", "node_type": "element", "path": "new_str"}
],
example='''
<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
@@ -0,0 +1,152 @@
from typing import Dict, Type, Any, List, Optional, Callable
from agentpress.tool import Tool, SchemaType
from utils.logger import logger
class ToolRegistry:
"""Registry for managing and accessing tools.
Maintains a collection of tool instances and their schemas, allowing for
selective registration of tool functions and easy access to tool capabilities.
Attributes:
tools (Dict[str, Dict[str, Any]]): OpenAPI-style tools and schemas
xml_tools (Dict[str, Dict[str, Any]]): XML-style tools and schemas
Methods:
register_tool: Register a tool with optional function filtering
get_tool: Get a specific tool by name
get_xml_tool: Get a tool by XML tag name
get_openapi_schemas: Get OpenAPI schemas for function calling
get_xml_examples: Get examples of XML tool usage
"""
def __init__(self):
"""Initialize a new ToolRegistry instance."""
self.tools = {}
self.xml_tools = {}
logger.debug("Initialized new ToolRegistry instance")
def register_tool(self, tool_class: Type[Tool], function_names: Optional[List[str]] = None, **kwargs):
"""Register a tool with optional function filtering.
Args:
tool_class: The tool class to register
function_names: Optional list of specific functions to register
**kwargs: Additional arguments passed to tool initialization
Notes:
- If function_names is None, all functions are registered
- Handles both OpenAPI and XML schema registration
"""
logger.debug(f"Registering tool class: {tool_class.__name__}")
tool_instance = tool_class(**kwargs)
schemas = tool_instance.get_schemas()
logger.debug(f"Available schemas for {tool_class.__name__}: {list(schemas.keys())}")
registered_openapi = 0
registered_xml = 0
for func_name, schema_list in schemas.items():
if function_names is None or func_name in function_names:
for schema in schema_list:
if schema.schema_type == SchemaType.OPENAPI:
self.tools[func_name] = {
"instance": tool_instance,
"schema": schema
}
registered_openapi += 1
logger.debug(f"Registered OpenAPI function {func_name} from {tool_class.__name__}")
if schema.schema_type == SchemaType.XML and schema.xml_schema:
self.xml_tools[schema.xml_schema.tag_name] = {
"instance": tool_instance,
"method": func_name,
"schema": schema
}
registered_xml += 1
logger.debug(f"Registered XML tag {schema.xml_schema.tag_name} -> {func_name} from {tool_class.__name__}")
logger.debug(f"Tool registration complete for {tool_class.__name__}: {registered_openapi} OpenAPI functions, {registered_xml} XML tags")
def get_available_functions(self) -> Dict[str, Callable]:
"""Get all available tool functions.
Returns:
Dict mapping function names to their implementations
"""
available_functions = {}
# Get OpenAPI tool functions
for tool_name, tool_info in self.tools.items():
tool_instance = tool_info['instance']
function_name = tool_name
function = getattr(tool_instance, function_name)
available_functions[function_name] = function
# Get XML tool functions
for tag_name, tool_info in self.xml_tools.items():
tool_instance = tool_info['instance']
method_name = tool_info['method']
function = getattr(tool_instance, method_name)
available_functions[method_name] = function
logger.debug(f"Retrieved {len(available_functions)} available functions")
return available_functions
def get_tool(self, tool_name: str) -> Dict[str, Any]:
"""Get a specific tool by name.
Args:
tool_name: Name of the tool function
Returns:
Dict containing tool instance and schema, or empty dict if not found
"""
tool = self.tools.get(tool_name, {})
if not tool:
logger.warning(f"Tool not found: {tool_name}")
return tool
def get_xml_tool(self, tag_name: str) -> Dict[str, Any]:
"""Get tool info by XML tag name.
Args:
tag_name: XML tag name for the tool
Returns:
Dict containing tool instance, method name, and schema
"""
tool = self.xml_tools.get(tag_name, {})
if not tool:
logger.warning(f"XML tool not found for tag: {tag_name}")
return tool
def get_openapi_schemas(self) -> List[Dict[str, Any]]:
"""Get OpenAPI schemas for function calling.
Returns:
List of OpenAPI-compatible schema definitions
"""
schemas = [
tool_info['schema'].schema
for tool_info in self.tools.values()
if tool_info['schema'].schema_type == SchemaType.OPENAPI
]
logger.debug(f"Retrieved {len(schemas)} OpenAPI schemas")
return schemas
def get_xml_examples(self) -> Dict[str, str]:
"""Get all XML tag examples.
Returns:
Dict mapping tag names to their example usage
"""
examples = {}
for tool_info in self.xml_tools.values():
schema = tool_info['schema']
if schema.xml_schema and schema.xml_schema.example:
examples[schema.xml_schema.tag_name] = schema.xml_schema.example
logger.debug(f"Retrieved {len(examples)} XML examples")
return examples
app/agentpress/utils/__init__.py
+1
View File
@@ -0,0 +1 @@
# Utils module for AgentPress
app/agentpress/utils/json_helpers.py
+174
View File
@@ -0,0 +1,174 @@
"""
JSON helper utilities for handling both legacy (string) and new (dict/list) formats.
These utilities help with the transition from storing JSON as strings to storing
them as proper JSONB objects in the database.
"""
import json
from typing import Any, Union, Dict, List
def ensure_dict(value: Union[str, Dict[str, Any], None], default: Dict[str, Any] = None) -> Dict[str, Any]:
"""
Ensure a value is a dictionary.
Handles:
- None -> returns default or {}
- Dict -> returns as-is
- JSON string -> parses and returns dict
- Other -> returns default or {}
Args:
value: The value to ensure is a dict
default: Default value if conversion fails
Returns:
A dictionary
"""
if default is None:
default = {}
if value is None:
return default
if isinstance(value, dict):
return value
if isinstance(value, str):
try:
parsed = json.loads(value)
if isinstance(parsed, dict):
return parsed
return default
except (json.JSONDecodeError, TypeError):
return default
return default
def ensure_list(value: Union[str, List[Any], None], default: List[Any] = None) -> List[Any]:
"""
Ensure a value is a list.
Handles:
- None -> returns default or []
- List -> returns as-is
- JSON string -> parses and returns list
- Other -> returns default or []
Args:
value: The value to ensure is a list
default: Default value if conversion fails
Returns:
A list
"""
if default is None:
default = []
if value is None:
return default
if isinstance(value, list):
return value
if isinstance(value, str):
try:
parsed = json.loads(value)
if isinstance(parsed, list):
return parsed
return default
except (json.JSONDecodeError, TypeError):
return default
return default
def safe_json_parse(value: Union[str, Dict, List, Any], default: Any = None) -> Any:
"""
Safely parse a value that might be JSON string or already parsed.
This handles the transition period where some data might be stored as
JSON strings (old format) and some as proper objects (new format).
Args:
value: The value to parse
default: Default value if parsing fails
Returns:
Parsed value or default
"""
if value is None:
return default
# If it's already a dict or list, return as-is
if isinstance(value, (dict, list)):
return value
# If it's a string, try to parse it
if isinstance(value, str):
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
# If it's not valid JSON, return the string itself
return value
# For any other type, return as-is
return value
def to_json_string(value: Any) -> str:
"""
Convert a value to a JSON string if needed.
This is used for backwards compatibility when yielding data that
expects JSON strings.
Args:
value: The value to convert
Returns:
JSON string representation
"""
if isinstance(value, str):
# If it's already a string, check if it's valid JSON
try:
json.loads(value)
return value # It's already a JSON string
except (json.JSONDecodeError, TypeError):
# It's a plain string, encode it as JSON
return json.dumps(value)
# For all other types, convert to JSON
return json.dumps(value)
def format_for_yield(message_object: Dict[str, Any]) -> Dict[str, Any]:
"""
Format a message object for yielding, ensuring content and metadata are JSON strings.
This maintains backward compatibility with clients expecting JSON strings
while the database now stores proper objects.
Args:
message_object: The message object from the database
Returns:
Message object with content and metadata as JSON strings
"""
if not message_object:
return message_object
# Create a copy to avoid modifying the original
formatted = message_object.copy()
# Ensure content is a JSON string
if 'content' in formatted and not isinstance(formatted['content'], str):
formatted['content'] = json.dumps(formatted['content'])
# Ensure metadata is a JSON string
if 'metadata' in formatted and not isinstance(formatted['metadata'], str):
formatted['metadata'] = json.dumps(formatted['metadata'])
return formatted
app/agentpress/xml_tool_parser.py
+300
View File
@@ -0,0 +1,300 @@
"""
XML Tool Call Parser Module
This module provides a reliable XML tool call parsing system that supports
the Cursor-style format with structured function_calls blocks.
"""
import re
import xml.etree.ElementTree as ET
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import json
import logging
logger = logging.getLogger(__name__)
@dataclass
class XMLToolCall:
"""Represents a parsed XML tool call."""
function_name: str
parameters: Dict[str, Any]
raw_xml: str
parsing_details: Dict[str, Any]
class XMLToolParser:
"""
Parser for XML tool calls using the Cursor-style format:
<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/tool_base.py
+1 -1
View File
@@ -4,7 +4,7 @@ from typing import Optional
from agentpress.thread_manager import ThreadManager
from app.tool.base import BaseTool, ToolResult
from daytona_sdk import Sandbox
from sandbox.sandbox import get_or_start_sandbox
from daytona.sandbox import get_or_start_sandbox
from utils.logger import logger
from utils.files_utils import clean_path
app/services/billing.py
+969
View File
@@ -0,0 +1,969 @@
"""
Stripe Billing API implementation for Suna on top of Basejump. ONLY HAS SUPPOT FOR USER ACCOUNTS – no team accounts. As we are using the user_id as account_id as is the case with personal accounts. In personal accounts, the account_id equals the user_id. In team accounts, the account_id is unique.
stripe listen --forward-to localhost:8000/api/billing/webhook
"""
from fastapi import APIRouter, HTTPException, Depends, Request
from typing import Optional, Dict, Tuple
import stripe
from datetime import datetime, timezone
from utils.logger import logger
from utils.config import config, EnvMode
from services.supabase import DBConnection
from utils.auth_utils import get_current_user_id_from_jwt
from pydantic import BaseModel
from utils.constants import MODEL_ACCESS_TIERS, MODEL_NAME_ALIASES
import os
# Initialize Stripe
stripe.api_key = config.STRIPE_SECRET_KEY
# Initialize router
router = APIRouter(prefix="/billing", tags=["billing"])
SUBSCRIPTION_TIERS = {
config.STRIPE_FREE_TIER_ID: {'name': 'free', 'minutes': 60},
config.STRIPE_TIER_2_20_ID: {'name': 'tier_2_20', 'minutes': 120}, # 2 hours
config.STRIPE_TIER_6_50_ID: {'name': 'tier_6_50', 'minutes': 360}, # 6 hours
config.STRIPE_TIER_12_100_ID: {'name': 'tier_12_100', 'minutes': 720}, # 12 hours
config.STRIPE_TIER_25_200_ID: {'name': 'tier_25_200', 'minutes': 1500}, # 25 hours
config.STRIPE_TIER_50_400_ID: {'name': 'tier_50_400', 'minutes': 3000}, # 50 hours
config.STRIPE_TIER_125_800_ID: {'name': 'tier_125_800', 'minutes': 7500}, # 125 hours
config.STRIPE_TIER_200_1000_ID: {'name': 'tier_200_1000', 'minutes': 12000}, # 200 hours
}
# Pydantic models for request/response validation
class CreateCheckoutSessionRequest(BaseModel):
price_id: str
success_url: str
cancel_url: str
tolt_referral: Optional[str] = None
class CreatePortalSessionRequest(BaseModel):
return_url: str
class SubscriptionStatus(BaseModel):
status: str # e.g., 'active', 'trialing', 'past_due', 'scheduled_downgrade', 'no_subscription'
plan_name: Optional[str] = None
price_id: Optional[str] = None # Added price ID
current_period_end: Optional[datetime] = None
cancel_at_period_end: bool = False
trial_end: Optional[datetime] = None
minutes_limit: Optional[int] = None
current_usage: Optional[float] = None
# Fields for scheduled changes
has_schedule: bool = False
scheduled_plan_name: Optional[str] = None
scheduled_price_id: Optional[str] = None # Added scheduled price ID
scheduled_change_date: Optional[datetime] = None
# Helper functions
async def get_stripe_customer_id(client, user_id: str) -> Optional[str]:
"""Get the Stripe customer ID for a user."""
result = await client.schema('basejump').from_('billing_customers') \
.select('id') \
.eq('account_id', user_id) \
.execute()
if result.data and len(result.data) > 0:
return result.data[0]['id']
return None
async def create_stripe_customer(client, user_id: str, email: str) -> str:
"""Create a new Stripe customer for a user."""
# Create customer in Stripe
customer = stripe.Customer.create(
email=email,
metadata={"user_id": user_id}
)
# Store customer ID in Supabase
await client.schema('basejump').from_('billing_customers').insert({
'id': customer.id,
'account_id': user_id,
'email': email,
'provider': 'stripe'
}).execute()
return customer.id
async def get_user_subscription(user_id: str) -> Optional[Dict]:
"""Get the current subscription for a user from Stripe."""
try:
# Get customer ID
db = DBConnection()
client = await db.client
customer_id = await get_stripe_customer_id(client, user_id)
if not customer_id:
return None
# Get all active subscriptions for the customer
subscriptions = stripe.Subscription.list(
customer=customer_id,
status='active'
)
# print("Found subscriptions:", subscriptions)
# Check if we have any subscriptions
if not subscriptions or not subscriptions.get('data'):
return None
# Filter subscriptions to only include our product's subscriptions
our_subscriptions = []
for sub in subscriptions['data']:
# Get the first subscription item
if sub.get('items') and sub['items'].get('data') and len(sub['items']['data']) > 0:
item = sub['items']['data'][0]
if item.get('price') and item['price'].get('id') in [
config.STRIPE_FREE_TIER_ID,
config.STRIPE_TIER_2_20_ID,
config.STRIPE_TIER_6_50_ID,
config.STRIPE_TIER_12_100_ID,
config.STRIPE_TIER_25_200_ID,
config.STRIPE_TIER_50_400_ID,
config.STRIPE_TIER_125_800_ID,
config.STRIPE_TIER_200_1000_ID
]:
our_subscriptions.append(sub)
if not our_subscriptions:
return None
# If there are multiple active subscriptions, we need to handle this
if len(our_subscriptions) > 1:
logger.warning(f"User {user_id} has multiple active subscriptions: {[sub['id'] for sub in our_subscriptions]}")
# Get the most recent subscription
most_recent = max(our_subscriptions, key=lambda x: x['created'])
# Cancel all other subscriptions
for sub in our_subscriptions:
if sub['id'] != most_recent['id']:
try:
stripe.Subscription.modify(
sub['id'],
cancel_at_period_end=True
)
logger.info(f"Cancelled subscription {sub['id']} for user {user_id}")
except Exception as e:
logger.error(f"Error cancelling subscription {sub['id']}: {str(e)}")
return most_recent
return our_subscriptions[0]
except Exception as e:
logger.error(f"Error getting subscription from Stripe: {str(e)}")
return None
async def calculate_monthly_usage(client, user_id: str) -> float:
"""Calculate total agent run minutes for the current month for a user."""
# Get start of current month in UTC
now = datetime.now(timezone.utc)
start_of_month = datetime(now.year, now.month, 1, tzinfo=timezone.utc)
# First get all threads for this user
threads_result = await client.table('threads') \
.select('thread_id') \
.eq('account_id', user_id) \
.execute()
if not threads_result.data:
return 0.0
thread_ids = [t['thread_id'] for t in threads_result.data]
# Then get all agent runs for these threads in current month
runs_result = await client.table('agent_runs') \
.select('started_at, completed_at') \
.in_('thread_id', thread_ids) \
.gte('started_at', start_of_month.isoformat()) \
.execute()
if not runs_result.data:
return 0.0
# Calculate total minutes
total_seconds = 0
now_ts = now.timestamp()
for run in runs_result.data:
start_time = datetime.fromisoformat(run['started_at'].replace('Z', '+00:00')).timestamp()
if run['completed_at']:
end_time = datetime.fromisoformat(run['completed_at'].replace('Z', '+00:00')).timestamp()
if start_time < end_time - 7200:
continue
else:
# if the start time is more than an hour ago, don't consider that time in total. else use the current time
if start_time < now_ts - 3600:
continue
else:
end_time = now_ts
total_seconds += (end_time - start_time)
return total_seconds / 60 # Convert to minutes
async def get_allowed_models_for_user(client, user_id: str):
"""
Get the list of models allowed for a user based on their subscription tier.
Returns:
List of model names allowed for the user's subscription tier.
"""
subscription = await get_user_subscription(user_id)
tier_name = 'free'
if subscription:
price_id = None
if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0:
price_id = subscription['items']['data'][0]['price']['id']
else:
price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID)
# Get tier info for this price_id
tier_info = SUBSCRIPTION_TIERS.get(price_id)
if tier_info:
tier_name = tier_info['name']
# Return allowed models for this tier
return MODEL_ACCESS_TIERS.get(tier_name, MODEL_ACCESS_TIERS['free']) # Default to free tier if unknown
async def can_use_model(client, user_id: str, model_name: str):
if config.ENV_MODE == EnvMode.LOCAL:
logger.info("Running in local development mode - billing checks are disabled")
return True, "Local development mode - billing disabled", {
"price_id": "local_dev",
"plan_name": "Local Development",
"minutes_limit": "no limit"
}
allowed_models = await get_allowed_models_for_user(client, user_id)
resolved_model = MODEL_NAME_ALIASES.get(model_name, model_name)
if resolved_model in allowed_models:
return True, "Model access allowed", allowed_models
return False, f"Your current subscription plan does not include access to {model_name}. Please upgrade your subscription or choose from your available models: {', '.join(allowed_models)}", allowed_models
async def check_billing_status(client, user_id: str) -> Tuple[bool, str, Optional[Dict]]:
"""
Check if a user can run agents based on their subscription and usage.
Returns:
Tuple[bool, str, Optional[Dict]]: (can_run, message, subscription_info)
"""
if config.ENV_MODE == EnvMode.LOCAL:
logger.info("Running in local development mode - billing checks are disabled")
return True, "Local development mode - billing disabled", {
"price_id": "local_dev",
"plan_name": "Local Development",
"minutes_limit": "no limit"
}
# Get current subscription
subscription = await get_user_subscription(user_id)
# print("Current subscription:", subscription)
# If no subscription, they can use free tier
if not subscription:
subscription = {
'price_id': config.STRIPE_FREE_TIER_ID, # Free tier
'plan_name': 'free'
}
# Extract price ID from subscription items
price_id = None
if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0:
price_id = subscription['items']['data'][0]['price']['id']
else:
price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID)
# Get tier info - default to free tier if not found
tier_info = SUBSCRIPTION_TIERS.get(price_id)
if not tier_info:
logger.warning(f"Unknown subscription tier: {price_id}, defaulting to free tier")
tier_info = SUBSCRIPTION_TIERS[config.STRIPE_FREE_TIER_ID]
# Calculate current month's usage
current_usage = await calculate_monthly_usage(client, user_id)
# Check if within limits
if current_usage >= tier_info['minutes']:
return False, f"Monthly limit of {tier_info['minutes']} minutes reached. Please upgrade your plan or wait until next month.", subscription
return True, "OK", subscription
# API endpoints
@router.post("/create-checkout-session")
async def create_checkout_session(
request: CreateCheckoutSessionRequest,
current_user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Create a Stripe Checkout session or modify an existing subscription."""
try:
# Get Supabase client
db = DBConnection()
client = await db.client
# Get user email from auth.users
user_result = await client.auth.admin.get_user_by_id(current_user_id)
if not user_result: raise HTTPException(status_code=404, detail="User not found")
email = user_result.user.email
# Get or create Stripe customer
customer_id = await get_stripe_customer_id(client, current_user_id)
if not customer_id: customer_id = await create_stripe_customer(client, current_user_id, email)
# Get the target price and product ID
try:
price = stripe.Price.retrieve(request.price_id, expand=['product'])
product_id = price['product']['id']
except stripe.error.InvalidRequestError:
raise HTTPException(status_code=400, detail=f"Invalid price ID: {request.price_id}")
# Verify the price belongs to our product
if product_id != config.STRIPE_PRODUCT_ID:
raise HTTPException(status_code=400, detail="Price ID does not belong to the correct product.")
# Check for existing subscription for our product
existing_subscription = await get_user_subscription(current_user_id)
# print("Existing subscription for product:", existing_subscription)
if existing_subscription:
# --- Handle Subscription Change (Upgrade or Downgrade) ---
try:
subscription_id = existing_subscription['id']
subscription_item = existing_subscription['items']['data'][0]
current_price_id = subscription_item['price']['id']
# Skip if already on this plan
if current_price_id == request.price_id:
return {
"subscription_id": subscription_id,
"status": "no_change",
"message": "Already subscribed to this plan.",
"details": {
"is_upgrade": None,
"effective_date": None,
"current_price": round(price['unit_amount'] / 100, 2) if price.get('unit_amount') else 0,
"new_price": round(price['unit_amount'] / 100, 2) if price.get('unit_amount') else 0,
}
}
# Get current and new price details
current_price = stripe.Price.retrieve(current_price_id)
new_price = price # Already retrieved
is_upgrade = new_price['unit_amount'] > current_price['unit_amount']
if is_upgrade:
# --- Handle Upgrade --- Immediate modification
updated_subscription = stripe.Subscription.modify(
subscription_id,
items=[{
'id': subscription_item['id'],
'price': request.price_id,
}],
proration_behavior='always_invoice', # Prorate and charge immediately
billing_cycle_anchor='now' # Reset billing cycle
)
# Update active status in database to true (customer has active subscription)
await client.schema('basejump').from_('billing_customers').update(
{'active': True}
).eq('id', customer_id).execute()
logger.info(f"Updated customer {customer_id} active status to TRUE after subscription upgrade")
latest_invoice = None
if updated_subscription.get('latest_invoice'):
latest_invoice = stripe.Invoice.retrieve(updated_subscription['latest_invoice'])
return {
"subscription_id": updated_subscription['id'],
"status": "updated",
"message": "Subscription upgraded successfully",
"details": {
"is_upgrade": True,
"effective_date": "immediate",
"current_price": round(current_price['unit_amount'] / 100, 2) if current_price.get('unit_amount') else 0,
"new_price": round(new_price['unit_amount'] / 100, 2) if new_price.get('unit_amount') else 0,
"invoice": {
"id": latest_invoice['id'] if latest_invoice else None,
"status": latest_invoice['status'] if latest_invoice else None,
"amount_due": round(latest_invoice['amount_due'] / 100, 2) if latest_invoice else 0,
"amount_paid": round(latest_invoice['amount_paid'] / 100, 2) if latest_invoice else 0
} if latest_invoice else None
}
}
else:
# --- Handle Downgrade --- Use Subscription Schedule
try:
current_period_end_ts = subscription_item['current_period_end']
# Retrieve the subscription again to get the schedule ID if it exists
# This ensures we have the latest state before creating/modifying schedule
sub_with_schedule = stripe.Subscription.retrieve(subscription_id)
schedule_id = sub_with_schedule.get('schedule')
# Get the current phase configuration from the schedule or subscription
if schedule_id:
schedule = stripe.SubscriptionSchedule.retrieve(schedule_id)
# Find the current phase in the schedule
# This logic assumes simple schedules; might need refinement for complex ones
current_phase = None
for phase in reversed(schedule['phases']):
if phase['start_date'] <= datetime.now(timezone.utc).timestamp():
current_phase = phase
break
if not current_phase: # Fallback if logic fails
current_phase = schedule['phases'][-1]
else:
# If no schedule, the current subscription state defines the current phase
current_phase = {
'items': existing_subscription['items']['data'], # Use original items data
'start_date': existing_subscription['current_period_start'], # Use sub start if no schedule
# Add other relevant fields if needed for create/modify
}
# Prepare the current phase data for the update/create
# Ensure items is formatted correctly for the API
current_phase_items_for_api = []
for item in current_phase.get('items', []):
price_data = item.get('price')
quantity = item.get('quantity')
price_id = None
# Safely extract price ID whether it's an object or just the ID string
if isinstance(price_data, dict):
price_id = price_data.get('id')
elif isinstance(price_data, str):
price_id = price_data
if price_id and quantity is not None:
current_phase_items_for_api.append({'price': price_id, 'quantity': quantity})
else:
logger.warning(f"Skipping item in current phase due to missing price ID or quantity: {item}")
if not current_phase_items_for_api:
raise ValueError("Could not determine valid items for the current phase.")
current_phase_update_data = {
'items': current_phase_items_for_api,
'start_date': current_phase['start_date'], # Preserve original start date
'end_date': current_period_end_ts, # End this phase at period end
'proration_behavior': 'none'
# Include other necessary fields from current_phase if modifying?
# e.g., 'billing_cycle_anchor', 'collection_method'? Usually inherited.
}
# Define the new (downgrade) phase
new_downgrade_phase_data = {
'items': [{'price': request.price_id, 'quantity': 1}],
'start_date': current_period_end_ts, # Start immediately after current phase ends
'proration_behavior': 'none'
# iterations defaults to 1, meaning it runs for one billing cycle
# then schedule ends based on end_behavior
}
# Update or Create Schedule
if schedule_id:
# Update existing schedule, replacing all future phases
# print(f"Updating existing schedule {schedule_id}")
logger.info(f"Updating existing schedule {schedule_id} for subscription {subscription_id}")
logger.debug(f"Current phase data: {current_phase_update_data}")
logger.debug(f"New phase data: {new_downgrade_phase_data}")
updated_schedule = stripe.SubscriptionSchedule.modify(
schedule_id,
phases=[current_phase_update_data, new_downgrade_phase_data],
end_behavior='release'
)
logger.info(f"Successfully updated schedule {updated_schedule['id']}")
else:
# Create a new schedule using the defined phases
print(f"Creating new schedule for subscription {subscription_id}")
logger.info(f"Creating new schedule for subscription {subscription_id}")
# Deep debug logging - write subscription details to help diagnose issues
logger.debug(f"Subscription details: {subscription_id}, current_period_end_ts: {current_period_end_ts}")
logger.debug(f"Current price: {current_price_id}, New price: {request.price_id}")
try:
updated_schedule = stripe.SubscriptionSchedule.create(
from_subscription=subscription_id,
phases=[
{
'start_date': current_phase['start_date'],
'end_date': current_period_end_ts,
'proration_behavior': 'none',
'items': [
{
'price': current_price_id,
'quantity': 1
}
]
},
{
'start_date': current_period_end_ts,
'proration_behavior': 'none',
'items': [
{
'price': request.price_id,
'quantity': 1
}
]
}
],
end_behavior='release'
)
# Don't try to link the schedule - that's handled by from_subscription
logger.info(f"Created new schedule {updated_schedule['id']} from subscription {subscription_id}")
# print(f"Created new schedule {updated_schedule['id']} from subscription {subscription_id}")
# Verify the schedule was created correctly
fetched_schedule = stripe.SubscriptionSchedule.retrieve(updated_schedule['id'])
logger.info(f"Schedule verification - Status: {fetched_schedule.get('status')}, Phase Count: {len(fetched_schedule.get('phases', []))}")
logger.debug(f"Schedule details: {fetched_schedule}")
except Exception as schedule_error:
logger.exception(f"Failed to create schedule: {str(schedule_error)}")
raise schedule_error # Re-raise to be caught by the outer try-except
return {
"subscription_id": subscription_id,
"schedule_id": updated_schedule['id'],
"status": "scheduled",
"message": "Subscription downgrade scheduled",
"details": {
"is_upgrade": False,
"effective_date": "end_of_period",
"current_price": round(current_price['unit_amount'] / 100, 2) if current_price.get('unit_amount') else 0,
"new_price": round(new_price['unit_amount'] / 100, 2) if new_price.get('unit_amount') else 0,
"effective_at": datetime.fromtimestamp(current_period_end_ts, tz=timezone.utc).isoformat()
}
}
except Exception as e:
logger.exception(f"Error handling subscription schedule for sub {subscription_id}: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error handling subscription schedule: {str(e)}")
except Exception as e:
logger.exception(f"Error updating subscription {existing_subscription.get('id') if existing_subscription else 'N/A'}: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error updating subscription: {str(e)}")
else:
session = stripe.checkout.Session.create(
customer=customer_id,
payment_method_types=['card'],
line_items=[{'price': request.price_id, 'quantity': 1}],
mode='subscription',
success_url=request.success_url,
cancel_url=request.cancel_url,
metadata={
'user_id': current_user_id,
'product_id': product_id,
'tolt_referral': request.tolt_referral
},
allow_promotion_codes=True
)
# Update customer status to potentially active (will be confirmed by webhook)
# This ensures customer is marked as active once payment is completed
await client.schema('basejump').from_('billing_customers').update(
{'active': True}
).eq('id', customer_id).execute()
logger.info(f"Updated customer {customer_id} active status to TRUE after creating checkout session")
return {"session_id": session['id'], "url": session['url'], "status": "new"}
except Exception as e:
logger.exception(f"Error creating checkout session: {str(e)}")
# Check if it's a Stripe error with more details
if hasattr(e, 'json_body') and e.json_body and 'error' in e.json_body:
error_detail = e.json_body['error'].get('message', str(e))
else:
error_detail = str(e)
raise HTTPException(status_code=500, detail=f"Error creating checkout session: {error_detail}")
@router.post("/create-portal-session")
async def create_portal_session(
request: CreatePortalSessionRequest,
current_user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Create a Stripe Customer Portal session for subscription management."""
try:
# Get Supabase client
db = DBConnection()
client = await db.client
# Get customer ID
customer_id = await get_stripe_customer_id(client, current_user_id)
if not customer_id:
raise HTTPException(status_code=404, detail="No billing customer found")
# Ensure the portal configuration has subscription_update enabled
try:
# First, check if we have a configuration that already enables subscription update
configurations = stripe.billing_portal.Configuration.list(limit=100)
active_config = None
# Look for a configuration with subscription_update enabled
for config in configurations.get('data', []):
features = config.get('features', {})
subscription_update = features.get('subscription_update', {})
if subscription_update.get('enabled', False):
active_config = config
logger.info(f"Found existing portal configuration with subscription_update enabled: {config['id']}")
break
# If no config with subscription_update found, create one or update the active one
if not active_config:
# Find the active configuration or create a new one
if configurations.get('data', []):
default_config = configurations['data'][0]
logger.info(f"Updating default portal configuration: {default_config['id']} to enable subscription_update")
active_config = stripe.billing_portal.Configuration.update(
default_config['id'],
features={
'subscription_update': {
'enabled': True,
'proration_behavior': 'create_prorations',
'default_allowed_updates': ['price']
},
# Preserve other features that may already be enabled
'customer_update': default_config.get('features', {}).get('customer_update', {'enabled': True, 'allowed_updates': ['email', 'address']}),
'invoice_history': {'enabled': True},
'payment_method_update': {'enabled': True}
}
)
else:
# Create a new configuration with subscription_update enabled
logger.info("Creating new portal configuration with subscription_update enabled")
active_config = stripe.billing_portal.Configuration.create(
business_profile={
'headline': 'Subscription Management',
'privacy_policy_url': config.FRONTEND_URL + '/privacy',
'terms_of_service_url': config.FRONTEND_URL + '/terms'
},
features={
'subscription_update': {
'enabled': True,
'proration_behavior': 'create_prorations',
'default_allowed_updates': ['price']
},
'customer_update': {
'enabled': True,
'allowed_updates': ['email', 'address']
},
'invoice_history': {'enabled': True},
'payment_method_update': {'enabled': True}
}
)
# Log the active configuration for debugging
logger.info(f"Using portal configuration: {active_config['id']} with subscription_update: {active_config.get('features', {}).get('subscription_update', {}).get('enabled', False)}")
except Exception as config_error:
logger.warning(f"Error configuring portal: {config_error}. Continuing with default configuration.")
# Create portal session using the proper configuration if available
portal_params = {
"customer": customer_id,
"return_url": request.return_url
}
# Add configuration_id if we found or created one with subscription_update enabled
if active_config:
portal_params["configuration"] = active_config['id']
# Create the session
session = stripe.billing_portal.Session.create(**portal_params)
return {"url": session.url}
except Exception as e:
logger.error(f"Error creating portal session: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@router.get("/subscription")
async def get_subscription(
current_user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Get the current subscription status for the current user, including scheduled changes."""
try:
# Get subscription from Stripe (this helper already handles filtering/cleanup)
subscription = await get_user_subscription(current_user_id)
# print("Subscription data for status:", subscription)
if not subscription:
# Default to free tier status if no active subscription for our product
free_tier_id = config.STRIPE_FREE_TIER_ID
free_tier_info = SUBSCRIPTION_TIERS.get(free_tier_id)
return SubscriptionStatus(
status="no_subscription",
plan_name=free_tier_info.get('name', 'free') if free_tier_info else 'free',
price_id=free_tier_id,
minutes_limit=free_tier_info.get('minutes') if free_tier_info else 0
)
# Extract current plan details
current_item = subscription['items']['data'][0]
current_price_id = current_item['price']['id']
current_tier_info = SUBSCRIPTION_TIERS.get(current_price_id)
if not current_tier_info:
# Fallback if somehow subscribed to an unknown price within our product
logger.warning(f"User {current_user_id} subscribed to unknown price {current_price_id}. Defaulting info.")
current_tier_info = {'name': 'unknown', 'minutes': 0}
# Calculate current usage
db = DBConnection()
client = await db.client
current_usage = await calculate_monthly_usage(client, current_user_id)
status_response = SubscriptionStatus(
status=subscription['status'], # 'active', 'trialing', etc.
plan_name=subscription['plan'].get('nickname') or current_tier_info['name'],
price_id=current_price_id,
current_period_end=datetime.fromtimestamp(current_item['current_period_end'], tz=timezone.utc),
cancel_at_period_end=subscription['cancel_at_period_end'],
trial_end=datetime.fromtimestamp(subscription['trial_end'], tz=timezone.utc) if subscription.get('trial_end') else None,
minutes_limit=current_tier_info['minutes'],
current_usage=round(current_usage, 2),
has_schedule=False # Default
)
# Check for an attached schedule (indicates pending downgrade)
schedule_id = subscription.get('schedule')
if schedule_id:
try:
schedule = stripe.SubscriptionSchedule.retrieve(schedule_id)
# Find the *next* phase after the current one
next_phase = None
current_phase_end = current_item['current_period_end']
for phase in schedule.get('phases', []):
# Check if this phase starts exactly when the current one ends
if phase.get('start_date') == current_phase_end:
next_phase = phase
break # Found the immediate next phase
if next_phase:
scheduled_item = next_phase['items'][0] # Assuming single item
scheduled_price_id = scheduled_item['price'] # Price ID might be string here
scheduled_tier_info = SUBSCRIPTION_TIERS.get(scheduled_price_id)
status_response.has_schedule = True
status_response.status = 'scheduled_downgrade' # Override status
status_response.scheduled_plan_name = scheduled_tier_info.get('name', 'unknown') if scheduled_tier_info else 'unknown'
status_response.scheduled_price_id = scheduled_price_id
status_response.scheduled_change_date = datetime.fromtimestamp(next_phase['start_date'], tz=timezone.utc)
except Exception as schedule_error:
logger.error(f"Error retrieving or parsing schedule {schedule_id} for sub {subscription['id']}: {schedule_error}")
# Proceed without schedule info if retrieval fails
return status_response
except Exception as e:
logger.exception(f"Error getting subscription status for user {current_user_id}: {str(e)}") # Use logger.exception
raise HTTPException(status_code=500, detail="Error retrieving subscription status.")
@router.get("/check-status")
async def check_status(
current_user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Check if the user can run agents based on their subscription and usage."""
try:
# Get Supabase client
db = DBConnection()
client = await db.client
can_run, message, subscription = await check_billing_status(client, current_user_id)
return {
"can_run": can_run,
"message": message,
"subscription": subscription
}
except Exception as e:
logger.error(f"Error checking billing status: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@router.post("/webhook")
async def stripe_webhook(request: Request):
"""Handle Stripe webhook events."""
try:
# Get the webhook secret from config
webhook_secret = config.STRIPE_WEBHOOK_SECRET
# Get the webhook payload
payload = await request.body()
sig_header = request.headers.get('stripe-signature')
# Verify webhook signature
try:
event = stripe.Webhook.construct_event(
payload, sig_header, webhook_secret
)
except ValueError as e:
raise HTTPException(status_code=400, detail="Invalid payload")
except stripe.error.SignatureVerificationError as e:
raise HTTPException(status_code=400, detail="Invalid signature")
# Handle the event
if event.type in ['customer.subscription.created', 'customer.subscription.updated', 'customer.subscription.deleted']:
# Extract the subscription and customer information
subscription = event.data.object
customer_id = subscription.get('customer')
if not customer_id:
logger.warning(f"No customer ID found in subscription event: {event.type}")
return {"status": "error", "message": "No customer ID found"}
# Get database connection
db = DBConnection()
client = await db.client
if event.type == 'customer.subscription.created' or event.type == 'customer.subscription.updated':
# Check if subscription is active
if subscription.get('status') in ['active', 'trialing']:
# Update customer's active status to true
await client.schema('basejump').from_('billing_customers').update(
{'active': True}
).eq('id', customer_id).execute()
logger.info(f"Webhook: Updated customer {customer_id} active status to TRUE based on {event.type}")
else:
# Subscription is not active (e.g., past_due, canceled, etc.)
# Check if customer has any other active subscriptions before updating status
has_active = len(stripe.Subscription.list(
customer=customer_id,
status='active',
limit=1
).get('data', [])) > 0
if not has_active:
await client.schema('basejump').from_('billing_customers').update(
{'active': False}
).eq('id', customer_id).execute()
logger.info(f"Webhook: Updated customer {customer_id} active status to FALSE based on {event.type}")
elif event.type == 'customer.subscription.deleted':
# Check if customer has any other active subscriptions
has_active = len(stripe.Subscription.list(
customer=customer_id,
status='active',
limit=1
).get('data', [])) > 0
if not has_active:
# If no active subscriptions left, set active to false
await client.schema('basejump').from_('billing_customers').update(
{'active': False}
).eq('id', customer_id).execute()
logger.info(f"Webhook: Updated customer {customer_id} active status to FALSE after subscription deletion")
logger.info(f"Processed {event.type} event for customer {customer_id}")
return {"status": "success"}
except Exception as e:
logger.error(f"Error processing webhook: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@router.get("/available-models")
async def get_available_models(
current_user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Get the list of models available to the user based on their subscription tier."""
try:
# Get Supabase client
db = DBConnection()
client = await db.client
# Check if we're in local development mode
if config.ENV_MODE == EnvMode.LOCAL:
logger.info("Running in local development mode - billing checks are disabled")
# In local mode, return all models from MODEL_NAME_ALIASES
model_info = []
for short_name, full_name in MODEL_NAME_ALIASES.items():
# Skip entries where the key is a full name to avoid duplicates
# if short_name == full_name or '/' in short_name:
# continue
model_info.append({
"id": full_name,
"display_name": short_name,
"short_name": short_name,
"requires_subscription": False # Always false in local dev mode
})
return {
"models": model_info,
"subscription_tier": "Local Development",
"total_models": len(model_info)
}
# For non-local mode, get list of allowed models for this user
allowed_models = await get_allowed_models_for_user(client, current_user_id)
free_tier_models = MODEL_ACCESS_TIERS.get('free', [])
# Get subscription info for context
subscription = await get_user_subscription(current_user_id)
# Determine tier name from subscription
tier_name = 'free'
if subscription:
price_id = None
if subscription.get('items') and subscription['items'].get('data') and len(subscription['items']['data']) > 0:
price_id = subscription['items']['data'][0]['price']['id']
else:
price_id = subscription.get('price_id', config.STRIPE_FREE_TIER_ID)
# Get tier info for this price_id
tier_info = SUBSCRIPTION_TIERS.get(price_id)
if tier_info:
tier_name = tier_info['name']
# Get all unique full model names from MODEL_NAME_ALIASES
all_models = set()
model_aliases = {}
for short_name, full_name in MODEL_NAME_ALIASES.items():
# Add all unique full model names
all_models.add(full_name)
# Only include short names that don't match their full names for aliases
if short_name != full_name and not short_name.startswith("openai/") and not short_name.startswith("anthropic/") and not short_name.startswith("openrouter/") and not short_name.startswith("xai/"):
if full_name not in model_aliases:
model_aliases[full_name] = short_name
# Create model info with display names for ALL models
model_info = []
for model in all_models:
display_name = model_aliases.get(model, model.split('/')[-1] if '/' in model else model)
# Check if model requires subscription (not in free tier)
requires_sub = model not in free_tier_models
# Check if model is available with current subscription
is_available = model in allowed_models
model_info.append({
"id": model,
"display_name": display_name,
"short_name": model_aliases.get(model),
"requires_subscription": requires_sub,
"is_available": is_available
})
return {
"models": model_info,
"subscription_tier": tier_name,
"total_models": len(model_info)
}
except Exception as e:
logger.error(f"Error getting available models: {str(e)}")
raise HTTPException(status_code=500, detail=f"Error getting available models: {str(e)}")
app/services/docker/redis.conf
+1
View File
@@ -0,0 +1 @@
timeout 120
app/services/email.py
+192
View File
@@ -0,0 +1,192 @@
import os
import logging
from typing import Optional
import mailtrap as mt
from utils.config import config
logger = logging.getLogger(__name__)
class EmailService:
def __init__(self):
self.api_token = os.getenv('MAILTRAP_API_TOKEN')
self.sender_email = os.getenv('MAILTRAP_SENDER_EMAIL', 'dom@kortix.ai')
self.sender_name = os.getenv('MAILTRAP_SENDER_NAME', 'Suna Team')
if not self.api_token:
logger.warning("MAILTRAP_API_TOKEN not found in environment variables")
self.client = None
else:
self.client = mt.MailtrapClient(token=self.api_token)
def send_welcome_email(self, user_email: str, user_name: Optional[str] = None) -> bool:
if not self.client:
logger.error("Cannot send email: MAILTRAP_API_TOKEN not configured")
return False
if not user_name:
user_name = user_email.split('@')[0].title()
subject = "🎉 Welcome to Suna — Let's Get Started "
html_content = self._get_welcome_email_template(user_name)
text_content = self._get_welcome_email_text(user_name)
return self._send_email(
to_email=user_email,
to_name=user_name,
subject=subject,
html_content=html_content,
text_content=text_content
)
def _send_email(
self,
to_email: str,
to_name: str,
subject: str,
html_content: str,
text_content: str
) -> bool:
try:
mail = mt.Mail(
sender=mt.Address(email=self.sender_email, name=self.sender_name),
to=[mt.Address(email=to_email, name=to_name)],
subject=subject,
text=text_content,
html=html_content,
category="welcome"
)
response = self.client.send(mail)
logger.info(f"Welcome email sent to {to_email}. Response: {response}")
return True
except Exception as e:
logger.error(f"Error sending email to {to_email}: {str(e)}")
return False
def _get_welcome_email_template(self, user_name: str) -> str:
return f"""<!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
@@ -0,0 +1,70 @@
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel, EmailStr
from typing import Optional
import asyncio
from services.email import email_service
from utils.logger import logger
router = APIRouter()
class SendWelcomeEmailRequest(BaseModel):
email: EmailStr
name: Optional[str] = None
class EmailResponse(BaseModel):
success: bool
message: str
@router.post("/send-welcome-email", response_model=EmailResponse)
async def send_welcome_email(request: SendWelcomeEmailRequest):
try:
logger.info(f"Sending welcome email to {request.email}")
success = email_service.send_welcome_email(
user_email=request.email,
user_name=request.name
)
if success:
return EmailResponse(
success=True,
message="Welcome email sent successfully"
)
else:
return EmailResponse(
success=False,
message="Failed to send welcome email"
)
except Exception as e:
logger.error(f"Error sending welcome email to {request.email}: {str(e)}")
raise HTTPException(
status_code=500,
detail="Internal server error while sending email"
)
@router.post("/send-welcome-email-background", response_model=EmailResponse)
async def send_welcome_email_background(request: SendWelcomeEmailRequest):
try:
logger.info(f"Queuing welcome email for {request.email}")
def send_email():
return email_service.send_welcome_email(
user_email=request.email,
user_name=request.name
)
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(send_email)
return EmailResponse(
success=True,
message="Welcome email queued for sending"
)
except Exception as e:
logger.error(f"Error queuing welcome email for {request.email}: {str(e)}")
raise HTTPException(
status_code=500,
detail="Internal server error while queuing email"
)
app/services/langfuse.py
+12
View File
@@ -0,0 +1,12 @@
import os
from langfuse import Langfuse
public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
secret_key = os.getenv("LANGFUSE_SECRET_KEY")
host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")
enabled = False
if public_key and secret_key:
enabled = True
langfuse = Langfuse(enabled=enabled)
app/services/llm.py
+411
View File
@@ -0,0 +1,411 @@
"""
LLM API interface for making calls to various language models.
This module provides a unified interface for making API calls to different LLM providers
(OpenAI, Anthropic, Groq, etc.) using LiteLLM. It includes support for:
- Streaming responses
- Tool calls and function calling
- Retry logic with exponential backoff
- Model-specific configurations
- Comprehensive error handling and logging
"""
from typing import Union, Dict, Any, Optional, AsyncGenerator, List
import os
import json
import asyncio
from openai import OpenAIError
import litellm
from utils.logger import logger
from utils.config import config
# litellm.set_verbose=True
litellm.modify_params=True
# Constants
MAX_RETRIES = 2
RATE_LIMIT_DELAY = 30
RETRY_DELAY = 0.1
class LLMError(Exception):
"""Base exception for LLM-related errors."""
pass
class LLMRetryError(LLMError):
"""Exception raised when retries are exhausted."""
pass
def setup_api_keys() -> None:
"""Set up API keys from environment variables."""
providers = ['OPENAI', 'ANTHROPIC', 'GROQ', 'OPENROUTER']
for provider in providers:
key = getattr(config, f'{provider}_API_KEY')
if key:
logger.debug(f"API key set for provider: {provider}")
else:
logger.warning(f"No API key found for provider: {provider}")
# Set up OpenRouter API base if not already set
if config.OPENROUTER_API_KEY and config.OPENROUTER_API_BASE:
os.environ['OPENROUTER_API_BASE'] = config.OPENROUTER_API_BASE
logger.debug(f"Set OPENROUTER_API_BASE to {config.OPENROUTER_API_BASE}")
# Set up AWS Bedrock credentials
aws_access_key = config.AWS_ACCESS_KEY_ID
aws_secret_key = config.AWS_SECRET_ACCESS_KEY
aws_region = config.AWS_REGION_NAME
if aws_access_key and aws_secret_key and aws_region:
logger.debug(f"AWS credentials set for Bedrock in region: {aws_region}")
# Configure LiteLLM to use AWS credentials
os.environ['AWS_ACCESS_KEY_ID'] = aws_access_key
os.environ['AWS_SECRET_ACCESS_KEY'] = aws_secret_key
os.environ['AWS_REGION_NAME'] = aws_region
else:
logger.warning(f"Missing AWS credentials for Bedrock integration - access_key: {bool(aws_access_key)}, secret_key: {bool(aws_secret_key)}, region: {aws_region}")
async def handle_error(error: Exception, attempt: int, max_attempts: int) -> None:
"""Handle API errors with appropriate delays and logging."""
delay = RATE_LIMIT_DELAY if isinstance(error, litellm.exceptions.RateLimitError) else RETRY_DELAY
logger.warning(f"Error on attempt {attempt + 1}/{max_attempts}: {str(error)}")
logger.debug(f"Waiting {delay} seconds before retry...")
await asyncio.sleep(delay)
def prepare_params(
messages: List[Dict[str, Any]],
model_name: str,
temperature: float = 0,
max_tokens: Optional[int] = None,
response_format: Optional[Any] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: str = "auto",
api_key: Optional[str] = None,
api_base: Optional[str] = None,
stream: bool = False,
top_p: Optional[float] = None,
model_id: Optional[str] = None,
enable_thinking: Optional[bool] = False,
reasoning_effort: Optional[str] = 'low'
) -> Dict[str, Any]:
"""Prepare parameters for the API call."""
params = {
"model": model_name,
"messages": messages,
"temperature": temperature,
"response_format": response_format,
"top_p": top_p,
"stream": stream,
}
if api_key:
params["api_key"] = api_key
if api_base:
params["api_base"] = api_base
if model_id:
params["model_id"] = model_id
# Handle token limits
if max_tokens is not None:
# For Claude 3.7 in Bedrock, do not set max_tokens or max_tokens_to_sample
# as it causes errors with inference profiles
if model_name.startswith("bedrock/") and "claude-3-7" in model_name:
logger.debug(f"Skipping max_tokens for Claude 3.7 model: {model_name}")
# Do not add any max_tokens parameter for Claude 3.7
else:
param_name = "max_completion_tokens" if 'o1' in model_name else "max_tokens"
params[param_name] = max_tokens
# Add tools if provided
if tools:
params.update({
"tools": tools,
"tool_choice": tool_choice
})
logger.debug(f"Added {len(tools)} tools to API parameters")
# # Add Claude-specific headers
if "claude" in model_name.lower() or "anthropic" in model_name.lower():
params["extra_headers"] = {
# "anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"
"anthropic-beta": "output-128k-2025-02-19"
}
params["fallbacks"] = [{
"model": "openrouter/anthropic/claude-sonnet-4",
"messages": messages,
}]
logger.debug("Added Claude-specific headers")
# Add OpenRouter-specific parameters
if model_name.startswith("openrouter/"):
logger.debug(f"Preparing OpenRouter parameters for model: {model_name}")
# Add optional site URL and app name from config
site_url = config.OR_SITE_URL
app_name = config.OR_APP_NAME
if site_url or app_name:
extra_headers = params.get("extra_headers", {})
if site_url:
extra_headers["HTTP-Referer"] = site_url
if app_name:
extra_headers["X-Title"] = app_name
params["extra_headers"] = extra_headers
logger.debug(f"Added OpenRouter site URL and app name to headers")
# Add Bedrock-specific parameters
if model_name.startswith("bedrock/"):
logger.debug(f"Preparing AWS Bedrock parameters for model: {model_name}")
if not model_id and "anthropic.claude-3-7-sonnet" in model_name:
params["model_id"] = "arn:aws:bedrock:us-west-2:935064898258:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
logger.debug(f"Auto-set model_id for Claude 3.7 Sonnet: {params['model_id']}")
# Apply Anthropic prompt caching (minimal implementation)
# Check model name *after* potential modifications (like adding bedrock/ prefix)
effective_model_name = params.get("model", model_name) # Use model from params if set, else original
if "claude" in effective_model_name.lower() or "anthropic" in effective_model_name.lower():
messages = params["messages"] # Direct reference, modification affects params
# Ensure messages is a list
if not isinstance(messages, list):
return params # Return early if messages format is unexpected
# 1. Process the first message if it's a system prompt with string content
if messages and messages[0].get("role") == "system":
content = messages[0].get("content")
if isinstance(content, str):
# Wrap the string content in the required list structure
messages[0]["content"] = [
{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}
]
elif isinstance(content, list):
# If content is already a list, check if the first text block needs cache_control
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
if "cache_control" not in item:
item["cache_control"] = {"type": "ephemeral"}
break # Apply to the first text block only for system prompt
# 2. Find and process relevant user and assistant messages (limit to 4 max)
last_user_idx = -1
second_last_user_idx = -1
last_assistant_idx = -1
for i in range(len(messages) - 1, -1, -1):
role = messages[i].get("role")
if role == "user":
if last_user_idx == -1:
last_user_idx = i
elif second_last_user_idx == -1:
second_last_user_idx = i
elif role == "assistant":
if last_assistant_idx == -1:
last_assistant_idx = i
# Stop searching if we've found all needed messages (system, last user, second last user, last assistant)
found_count = sum(idx != -1 for idx in [last_user_idx, second_last_user_idx, last_assistant_idx])
if found_count >= 3:
break
# Helper function to apply cache control
def apply_cache_control(message_idx: int, message_role: str):
if message_idx == -1:
return
message = messages[message_idx]
content = message.get("content")
if isinstance(content, str):
message["content"] = [
{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}
]
elif isinstance(content, list):
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
if "cache_control" not in item:
item["cache_control"] = {"type": "ephemeral"}
# Apply cache control to the identified messages (max 4: system, last user, second last user, last assistant)
# System message is always at index 0 if present
apply_cache_control(0, "system")
apply_cache_control(last_user_idx, "last user")
apply_cache_control(second_last_user_idx, "second last user")
apply_cache_control(last_assistant_idx, "last assistant")
# Add reasoning_effort for Anthropic models if enabled
use_thinking = enable_thinking if enable_thinking is not None else False
is_anthropic = "anthropic" in effective_model_name.lower() or "claude" in effective_model_name.lower()
if is_anthropic and use_thinking:
effort_level = reasoning_effort if reasoning_effort else 'low'
params["reasoning_effort"] = effort_level
params["temperature"] = 1.0 # Required by Anthropic when reasoning_effort is used
logger.info(f"Anthropic thinking enabled with reasoning_effort='{effort_level}'")
return params
async def make_llm_api_call(
messages: List[Dict[str, Any]],
model_name: str,
response_format: Optional[Any] = None,
temperature: float = 0,
max_tokens: Optional[int] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_choice: str = "auto",
api_key: Optional[str] = None,
api_base: Optional[str] = None,
stream: bool = False,
top_p: Optional[float] = None,
model_id: Optional[str] = None,
enable_thinking: Optional[bool] = False,
reasoning_effort: Optional[str] = 'low'
) -> Union[Dict[str, Any], AsyncGenerator]:
"""
Make an API call to a language model using LiteLLM.
Args:
messages: List of message dictionaries for the conversation
model_name: Name of the model to use (e.g., "gpt-4", "claude-3", "openrouter/openai/gpt-4", "bedrock/anthropic.claude-3-sonnet-20240229-v1:0")
response_format: Desired format for the response
temperature: Sampling temperature (0-1)
max_tokens: Maximum tokens in the response
tools: List of tool definitions for function calling
tool_choice: How to select tools ("auto" or "none")
api_key: Override default API key
api_base: Override default API base URL
stream: Whether to stream the response
top_p: Top-p sampling parameter
model_id: Optional ARN for Bedrock inference profiles
enable_thinking: Whether to enable thinking
reasoning_effort: Level of reasoning effort
Returns:
Union[Dict[str, Any], AsyncGenerator]: API response or stream
Raises:
LLMRetryError: If API call fails after retries
LLMError: For other API-related errors
"""
# debug <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
@@ -0,0 +1,129 @@
import os
import sys
import json
import asyncio
import subprocess
from typing import Dict, Any
from concurrent.futures import ThreadPoolExecutor
from fastapi import HTTPException # type: ignore
from utils.logger import logger
from mcp import ClientSession
from mcp.client.sse import sse_client # type: ignore
from mcp.client.streamable_http import streamablehttp_client # type: ignore
async def connect_streamable_http_server(url):
async with streamablehttp_client(url) as (
read_stream,
write_stream,
_,
):
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
tool_result = await session.list_tools()
print(f"Connected via HTTP ({len(tool_result.tools)} tools)")
tools_info = []
for tool in tool_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"inputSchema": tool.inputSchema
}
tools_info.append(tool_info)
return tools_info
async def discover_custom_tools(request_type: str, config: Dict[str, Any]):
logger.info(f"Received custom MCP discovery request: type={request_type}")
logger.debug(f"Request config: {config}")
tools = []
server_name = None
if request_type == 'http':
if 'url' not in config:
raise HTTPException(status_code=400, detail="HTTP configuration must include 'url' field")
url = config['url']
try:
async with asyncio.timeout(15):
tools_info = await connect_streamable_http_server(url)
for tool_info in tools_info:
tools.append({
"name": tool_info["name"],
"description": tool_info["description"],
"inputSchema": tool_info["inputSchema"]
})
except asyncio.TimeoutError:
raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond")
except Exception as e:
logger.error(f"Error connecting to HTTP MCP server: {e}")
raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}")
elif request_type == 'sse':
if 'url' not in config:
raise HTTPException(status_code=400, detail="SSE configuration must include 'url' field")
url = config['url']
headers = config.get('headers', {})
try:
async with asyncio.timeout(15):
try:
async with sse_client(url, headers=headers) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
tools_info = []
for tool in tools_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"input_schema": tool.inputSchema
}
tools_info.append(tool_info)
for tool_info in tools_info:
tools.append({
"name": tool_info["name"],
"description": tool_info["description"],
"inputSchema": tool_info["input_schema"]
})
except TypeError as e:
if "unexpected keyword argument" in str(e):
async with sse_client(url) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
tools_info = []
for tool in tools_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"input_schema": tool.inputSchema
}
tools_info.append(tool_info)
for tool_info in tools_info:
tools.append({
"name": tool_info["name"],
"description": tool_info["description"],
"inputSchema": tool_info["input_schema"]
})
else:
raise
except asyncio.TimeoutError:
raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond")
except Exception as e:
logger.error(f"Error connecting to SSE MCP server: {e}")
raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}")
else:
raise HTTPException(status_code=400, detail="Invalid server type. Must be 'http' or 'sse'")
response_data = {"tools": tools, "count": len(tools)}
if server_name:
response_data["serverName"] = server_name
logger.info(f"Returning {len(tools)} tools for server {server_name}")
return response_data
app/services/mcp_temp.py
+299
View File
@@ -0,0 +1,299 @@
import os
import sys
import json
import asyncio
import subprocess
from typing import Dict, Any
from concurrent.futures import ThreadPoolExecutor
from fastapi import HTTPException # type: ignore
from utils.logger import logger
from mcp import ClientSession
from mcp.client.sse import sse_client # type: ignore
from mcp.client.streamable_http import streamablehttp_client # type: ignore
windows_executor = ThreadPoolExecutor(max_workers=4)
# def run_mcp_stdio_sync(command, args, env_vars, timeout=30):
# try:
# env = os.environ.copy()
# env.update(env_vars)
# full_command = [command] + args
# process = subprocess.Popen(
# full_command,
# stdin=subprocess.PIPE,
# stdout=subprocess.PIPE,
# stderr=subprocess.PIPE,
# env=env,
# text=True,
# bufsize=0,
# creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == "win32" else 0
# )
# init_request = {
# "jsonrpc": "2.0",
# "id": 1,
# "method": "initialize",
# "params": {
# "protocolVersion": "2024-11-05",
# "capabilities": {},
# "clientInfo": {"name": "mcp-client", "version": "1.0.0"}
# }
# }
# process.stdin.write(json.dumps(init_request) + "\n")
# process.stdin.flush()
# init_response_line = process.stdout.readline().strip()
# if not init_response_line:
# raise Exception("No response from MCP server during initialization")
# init_response = json.loads(init_response_line)
# init_notification = {
# "jsonrpc": "2.0",
# "method": "notifications/initialized"
# }
# process.stdin.write(json.dumps(init_notification) + "\n")
# process.stdin.flush()
# tools_request = {
# "jsonrpc": "2.0",
# "id": 2,
# "method": "tools/list",
# "params": {}
# }
# process.stdin.write(json.dumps(tools_request) + "\n")
# process.stdin.flush()
# tools_response_line = process.stdout.readline().strip()
# if not tools_response_line:
# raise Exception("No response from MCP server for tools list")
# tools_response = json.loads(tools_response_line)
# tools_info = []
# if "result" in tools_response and "tools" in tools_response["result"]:
# for tool in tools_response["result"]["tools"]:
# tool_info = {
# "name": tool["name"],
# "description": tool.get("description", ""),
# "input_schema": tool.get("inputSchema", {})
# }
# tools_info.append(tool_info)
# return {
# "status": "connected",
# "transport": "stdio",
# "tools": tools_info
# }
# except subprocess.TimeoutExpired:
# return {
# "status": "error",
# "error": f"Process timeout after {timeout} seconds",
# "tools": []
# }
# except json.JSONDecodeError as e:
# return {
# "status": "error",
# "error": f"Invalid JSON response: {str(e)}",
# "tools": []
# }
# except Exception as e:
# return {
# "status": "error",
# "error": str(e),
# "tools": []
# }
# finally:
# try:
# if 'process' in locals():
# process.terminate()
# process.wait(timeout=5)
# except:
# pass
# async def connect_stdio_server_windows(server_name, server_config, all_tools, timeout):
# """Windows-compatible stdio connection using subprocess"""
# logger.info(f"Connecting to {server_name} using Windows subprocess method")
# command = server_config["command"]
# args = server_config.get("args", [])
# env_vars = server_config.get("env", {})
# loop = asyncio.get_event_loop()
# result = await loop.run_in_executor(
# windows_executor,
# run_mcp_stdio_sync,
# command,
# args,
# env_vars,
# timeout
# )
# all_tools[server_name] = result
# if result["status"] == "connected":
# logger.info(f" {server_name}: Connected via Windows subprocess ({len(result['tools'])} tools)")
# else:
# logger.error(f" {server_name}: Error - {result['error']}")
# async def list_mcp_tools_mixed_windows(config, timeout=15):
# all_tools = {}
# if "mcpServers" not in config:
# return all_tools
# mcp_servers = config["mcpServers"]
# for server_name, server_config in mcp_servers.items():
# logger.info(f"Connecting to MCP server: {server_name}")
# if server_config.get("disabled", False):
# all_tools[server_name] = {"status": "disabled", "tools": []}
# logger.info(f" {server_name}: Disabled")
# continue
# try:
# await connect_stdio_server_windows(server_name, server_config, all_tools, timeout)
# except asyncio.TimeoutError:
# all_tools[server_name] = {
# "status": "error",
# "error": f"Connection timeout after {timeout} seconds",
# "tools": []
# }
# logger.error(f" {server_name}: Timeout after {timeout} seconds")
# except Exception as e:
# error_msg = str(e)
# all_tools[server_name] = {
# "status": "error",
# "error": error_msg,
# "tools": []
# }
# logger.error(f" {server_name}: Error - {error_msg}")
# import traceback
# logger.debug(f"Full traceback for {server_name}: {traceback.format_exc()}")
# return all_tools
async def discover_custom_tools(request_type: str, config: Dict[str, Any]):
logger.info(f"Received custom MCP discovery request: type={request_type}")
logger.debug(f"Request config: {config}")
tools = []
server_name = None
# if request_type == 'json':
# try:
# all_tools = await list_mcp_tools_mixed_windows(config, timeout=30)
# if "mcpServers" in config and config["mcpServers"]:
# server_name = list(config["mcpServers"].keys())[0]
# if server_name in all_tools:
# server_info = all_tools[server_name]
# if server_info["status"] == "connected":
# tools = server_info["tools"]
# logger.info(f"Found {len(tools)} tools for server {server_name}")
# else:
# error_msg = server_info.get("error", "Unknown error")
# logger.error(f"Server {server_name} failed: {error_msg}")
# raise HTTPException(
# status_code=400,
# detail=f"Failed to connect to MCP server '{server_name}': {error_msg}"
# )
# else:
# logger.error(f"Server {server_name} not found in results")
# raise HTTPException(status_code=400, detail=f"Server '{server_name}' not found in results")
# else:
# logger.error("No MCP servers configured")
# raise HTTPException(status_code=400, detail="No MCP servers configured")
# except HTTPException:
# raise
# except Exception as e:
# logger.error(f"Error connecting to stdio MCP server: {e}")
# import traceback
# logger.error(f"Full traceback: {traceback.format_exc()}")
# raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}")
# if request_type == 'http':
# if 'url' not in config:
# raise HTTPException(status_code=400, detail="HTTP configuration must include 'url' field")
# url = config['url']
# await connect_streamable_http_server(url)
# tools = await connect_streamable_http_server(url)
# elif request_type == 'sse':
# if 'url' not in config:
# raise HTTPException(status_code=400, detail="SSE configuration must include 'url' field")
# url = config['url']
# headers = config.get('headers', {})
# try:
# async with asyncio.timeout(15):
# try:
# async with sse_client(url, headers=headers) as (read, write):
# async with ClientSession(read, write) as session:
# await session.initialize()
# tools_result = await session.list_tools()
# tools_info = []
# for tool in tools_result.tools:
# tool_info = {
# "name": tool.name,
# "description": tool.description,
# "input_schema": tool.inputSchema
# }
# tools_info.append(tool_info)
# for tool_info in tools_info:
# tools.append({
# "name": tool_info["name"],
# "description": tool_info["description"],
# "inputSchema": tool_info["input_schema"]
# })
# except TypeError as e:
# if "unexpected keyword argument" in str(e):
# async with sse_client(url) as (read, write):
# async with ClientSession(read, write) as session:
# await session.initialize()
# tools_result = await session.list_tools()
# tools_info = []
# for tool in tools_result.tools:
# tool_info = {
# "name": tool.name,
# "description": tool.description,
# "input_schema": tool.inputSchema
# }
# tools_info.append(tool_info)
# for tool_info in tools_info:
# tools.append({
# "name": tool_info["name"],
# "description": tool_info["description"],
# "inputSchema": tool_info["input_schema"]
# })
# else:
# raise
# except asyncio.TimeoutError:
# raise HTTPException(status_code=408, detail="Connection timeout - server took too long to respond")
# except Exception as e:
# logger.error(f"Error connecting to SSE MCP server: {e}")
# raise HTTPException(status_code=400, detail=f"Failed to connect to MCP server: {str(e)}")
# else:
# raise HTTPException(status_code=400, detail="Invalid server type. Must be 'json' or 'sse'")
# response_data = {"tools": tools, "count": len(tools)}
# if server_name:
# response_data["serverName"] = server_name
# logger.info(f"Returning {len(tools)} tools for server {server_name}")
# return response_data
app/services/redis.py
+153
View File
@@ -0,0 +1,153 @@
import redis.asyncio as redis
import os
from dotenv import load_dotenv
import asyncio
from utils.logger import logger
from typing import List, Any
from utils.retry import retry
# Redis client
client: redis.Redis | None = None
_initialized = False
_init_lock = asyncio.Lock()
# Constants
REDIS_KEY_TTL = 3600 * 24 # 24 hour TTL as safety mechanism
def initialize():
"""Initialize Redis connection using environment variables."""
global client
# Load environment variables if not already loaded
load_dotenv()
# Get Redis configuration
redis_host = os.getenv("REDIS_HOST", "redis")
redis_port = int(os.getenv("REDIS_PORT", 6379))
redis_password = os.getenv("REDIS_PASSWORD", "")
# Convert string 'True'/'False' to boolean
redis_ssl_str = os.getenv("REDIS_SSL", "False")
redis_ssl = redis_ssl_str.lower() == "true"
logger.info(f"Initializing Redis connection to {redis_host}:{redis_port}")
# Create Redis client with basic configuration
client = redis.Redis(
host=redis_host,
port=redis_port,
password=redis_password,
ssl=redis_ssl,
decode_responses=True,
socket_timeout=5.0,
socket_connect_timeout=5.0,
retry_on_timeout=True,
health_check_interval=30,
)
return client
async def initialize_async():
"""Initialize Redis connection asynchronously."""
global client, _initialized
async with _init_lock:
if not _initialized:
logger.info("Initializing Redis connection")
initialize()
try:
await client.ping()
logger.info("Successfully connected to Redis")
_initialized = True
except Exception as e:
logger.error(f"Failed to connect to Redis: {e}")
client = None
_initialized = False
raise
return client
async def close():
"""Close Redis connection."""
global client, _initialized
if client:
logger.info("Closing Redis connection")
await client.aclose()
client = None
_initialized = False
logger.info("Redis connection closed")
async def get_client():
"""Get the Redis client, initializing if necessary."""
global client, _initialized
if client is None or not _initialized:
await retry(lambda: initialize_async())
return client
# Basic Redis operations
async def set(key: str, value: str, ex: int = None, nx: bool = False):
"""Set a Redis key."""
redis_client = await get_client()
return await redis_client.set(key, value, ex=ex, nx=nx)
async def get(key: str, default: str = None):
"""Get a Redis key."""
redis_client = await get_client()
result = await redis_client.get(key)
return result if result is not None else default
async def delete(key: str):
"""Delete a Redis key."""
redis_client = await get_client()
return await redis_client.delete(key)
async def publish(channel: str, message: str):
"""Publish a message to a Redis channel."""
redis_client = await get_client()
return await redis_client.publish(channel, message)
async def create_pubsub():
"""Create a Redis pubsub object."""
redis_client = await get_client()
return redis_client.pubsub()
# List operations
async def rpush(key: str, *values: Any):
"""Append one or more values to a list."""
redis_client = await get_client()
return await redis_client.rpush(key, *values)
async def lrange(key: str, start: int, end: int) -> List[str]:
"""Get a range of elements from a list."""
redis_client = await get_client()
return await redis_client.lrange(key, start, end)
async def llen(key: str) -> int:
"""Get the length of a list."""
redis_client = await get_client()
return await redis_client.llen(key)
# Key management
async def expire(key: str, time: int):
"""Set a key's time to live in seconds."""
redis_client = await get_client()
return await redis_client.expire(key, time)
async def keys(pattern: str) -> List[str]:
"""Get keys matching a pattern."""
redis_client = await get_client()
return await redis_client.keys(pattern)
app/services/supabase.py
+113
View File
@@ -0,0 +1,113 @@
"""
Centralized database connection management for AgentPress using Supabase.
"""
from typing import Optional
from supabase import create_async_client, AsyncClient
from utils.logger import logger
from utils.config import config
import base64
import uuid
from datetime import datetime
class DBConnection:
"""Singleton database connection manager using Supabase."""
_instance: Optional['DBConnection'] = None
_initialized = False
_client: Optional[AsyncClient] = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
"""No initialization needed in __init__ as it's handled in __new__"""
pass
async def initialize(self):
"""Initialize the database connection."""
if self._initialized:
return
try:
supabase_url = config.SUPABASE_URL
# Use service role key preferentially for backend operations
supabase_key = config.SUPABASE_SERVICE_ROLE_KEY or config.SUPABASE_ANON_KEY
if not supabase_url or not supabase_key:
logger.error("Missing required environment variables for Supabase connection")
raise RuntimeError("SUPABASE_URL and a key (SERVICE_ROLE_KEY or ANON_KEY) environment variables must be set.")
logger.debug("Initializing Supabase connection")
self._client = await create_async_client(supabase_url, supabase_key)
self._initialized = True
key_type = "SERVICE_ROLE_KEY" if config.SUPABASE_SERVICE_ROLE_KEY else "ANON_KEY"
logger.debug(f"Database connection initialized with Supabase using {key_type}")
except Exception as e:
logger.error(f"Database initialization error: {e}")
raise RuntimeError(f"Failed to initialize database connection: {str(e)}")
@classmethod
async def disconnect(cls):
"""Disconnect from the database."""
if cls._client:
logger.info("Disconnecting from Supabase database")
await cls._client.close()
cls._initialized = False
logger.info("Database disconnected successfully")
@property
async def client(self) -> AsyncClient:
"""Get the Supabase client instance."""
if not self._initialized:
logger.debug("Supabase client not initialized, initializing now")
await self.initialize()
if not self._client:
logger.error("Database client is None after initialization")
raise RuntimeError("Database not initialized")
return self._client
async def upload_base64_image(self, base64_data: str, bucket_name: str = "browser-screenshots") -> str:
"""Upload a base64 encoded image to Supabase storage and return the URL.
Args:
base64_data (str): Base64 encoded image data (with or without data URL prefix)
bucket_name (str): Name of the storage bucket to upload to
Returns:
str: Public URL of the uploaded image
"""
try:
# Remove data URL prefix if present
if base64_data.startswith('data:'):
base64_data = base64_data.split(',')[1]
# Decode base64 data
image_data = base64.b64decode(base64_data)
# Generate unique filename
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
unique_id = str(uuid.uuid4())[:8]
filename = f"image_{timestamp}_{unique_id}.png"
# Upload to Supabase storage
client = await self.client
storage_response = await client.storage.from_(bucket_name).upload(
filename,
image_data,
{"content-type": "image/png"}
)
# Get public URL
public_url = await client.storage.from_(bucket_name).get_public_url(filename)
logger.debug(f"Successfully uploaded image to {public_url}")
return public_url
except Exception as e:
logger.error(f"Error uploading base64 image: {e}")
raise RuntimeError(f"Failed to upload image: {str(e)}")
app/services/transcription.py
+76
View File
@@ -0,0 +1,76 @@
import os
import openai
import tempfile
from fastapi import APIRouter, UploadFile, File, HTTPException, Depends
from pydantic import BaseModel
from typing import Optional
from utils.logger import logger
from utils.auth_utils import get_current_user_id_from_jwt
router = APIRouter(tags=["transcription"])
class TranscriptionResponse(BaseModel):
text: str
@router.post("/transcription", response_model=TranscriptionResponse)
async def transcribe_audio(
audio_file: UploadFile = File(...),
user_id: str = Depends(get_current_user_id_from_jwt)
):
"""Transcribe audio file to text using OpenAI Whisper."""
try:
# Validate file type - OpenAI supports these formats
allowed_types = [
'audio/mp3', 'audio/mpeg', 'audio/mp4', 'audio/m4a',
'audio/wav', 'audio/webm', 'audio/mpga'
]
logger.info(f"Received audio file: {audio_file.filename}, content_type: {audio_file.content_type}")
if audio_file.content_type not in allowed_types:
raise HTTPException(
status_code=400,
detail=f"Unsupported file type: {audio_file.content_type}. Supported types: {', '.join(allowed_types)}"
)
# Check file size (25MB limit)
content = await audio_file.read()
if len(content) > 25 * 1024 * 1024: # 25MB
raise HTTPException(status_code=400, detail="File size exceeds 25MB limit")
# Reset file pointer
await audio_file.seek(0)
# Initialize OpenAI client
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
# Create a temporary file with the correct extension
file_extension = audio_file.filename.split('.')[-1] if audio_file.filename and '.' in audio_file.filename else 'webm'
with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{file_extension}') as temp_file:
temp_file.write(content)
temp_file_path = temp_file.name
try:
# Transcribe audio using the temporary file
# OpenAI Whisper API has built-in limits: 25MB file size and handles duration limits internally
with open(temp_file_path, 'rb') as f:
transcription = client.audio.transcriptions.create(
model="gpt-4o-mini-transcribe",
file=f,
response_format="text"
)
logger.info(f"Successfully transcribed audio for user {user_id}")
return TranscriptionResponse(text=transcription)
finally:
# Clean up temporary file
try:
os.unlink(temp_file_path)
except Exception as e:
logger.warning(f"Failed to delete temporary file {temp_file_path}: {e}")
except Exception as e:
logger.error(f"Error transcribing audio for user {user_id}: {str(e)}")
raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
app/tool/base.py
+102 -3
View File
@@ -1,16 +1,75 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
from pydantic import BaseModel, Field
from typing import Any, Dict, Optional, Union
import inspect
import json
from utils.logger import logger
# class BaseTool(ABC, BaseModel):
# name: str
# description: str
# parameters: Optional[dict] = None
# class Config:
# arbitrary_types_allowed = True
# async def __call__(self, **kwargs) -> Any:
# """Execute the tool with given parameters."""
# return await self.execute(**kwargs)
# @abstractmethod
# async def execute(self, **kwargs) -> Any:
# """Execute the tool with given parameters."""
# def to_param(self) -> Dict:
# """Convert tool to function call format."""
# return {
# "type": "function",
# "function": {
# "name": self.name,
# "description": self.description,
# "parameters": self.parameters,
# },
# }
class BaseTool(ABC, BaseModel):
"""Consolidated base class for all tools combining BaseModel and Tool functionality.
Provides:
- Pydantic model validation
- Schema registration
- Standardized result handling
- Abstract execution interface
Attributes:
name (str): Tool name
description (str): Tool description
parameters (dict): Tool parameters schema
_schemas (Dict[str, List[ToolSchema]]): Registered method schemas
"""
name: str
description: str
parameters: Optional[dict] = None
# _schemas: Dict[str, List[ToolSchema]] = {}
class Config:
arbitrary_types_allowed = True
underscore_attrs_are_private = False
# def __init__(self, **data):
# """Initialize tool with model validation and schema registration."""
# super().__init__(**data)
# logger.debug(f"Initializing tool class: {self.__class__.__name__}")
# self._register_schemas()
# def _register_schemas(self):
# """Register schemas from all decorated methods."""
# for name, method in inspect.getmembers(self, predicate=inspect.ismethod):
# if hasattr(method, 'tool_schemas'):
# self._schemas[name] = method.tool_schemas
# logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}")
async def __call__(self, **kwargs) -> Any:
"""Execute the tool with given parameters."""
@@ -21,7 +80,11 @@ class BaseTool(ABC, BaseModel):
"""Execute the tool with given parameters."""
def to_param(self) -> Dict:
"""Convert tool to function call format."""
"""Convert tool to function call format.
Returns:
Dictionary with tool metadata in OpenAI function calling format
"""
return {
"type": "function",
"function": {
@@ -31,6 +94,42 @@ class BaseTool(ABC, BaseModel):
},
}
# def get_schemas(self) -> Dict[str, List[ToolSchema]]:
# """Get all registered tool schemas.
# Returns:
# Dict mapping method names to their schema definitions
# """
# return self._schemas
def success_response(self, data: Union[Dict[str, Any], str]) -> ToolResult:
"""Create a successful tool result.
Args:
data: Result data (dictionary or string)
Returns:
ToolResult with success=True and formatted output
"""
if isinstance(data, str):
text = data
else:
text = json.dumps(data, indent=2)
logger.debug(f"Created success response for {self.__class__.__name__}")
return ToolResult(output=text)
def fail_response(self, msg: str) -> ToolResult:
"""Create a failed tool result.
Args:
msg: Error message describing the failure
Returns:
ToolResult with success=False and error message
"""
logger.debug(f"Tool {self.__class__.__name__} returned failed result: {msg}")
return ToolResult(error=msg)
class ToolResult(BaseModel):
"""Represents the result of a tool execution."""
app/tool/data_providers_tool.py
-188
View File
@@ -1,188 +0,0 @@
import json
from typing import Union, Dict, Any
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
from agent.tools.data_providers.LinkedinProvider import LinkedinProvider
from agent.tools.data_providers.YahooFinanceProvider import YahooFinanceProvider
from agent.tools.data_providers.AmazonProvider import AmazonProvider
from agent.tools.data_providers.ZillowProvider import ZillowProvider
from agent.tools.data_providers.TwitterProvider import TwitterProvider
class DataProvidersTool(Tool):
"""Tool for making requests to various data providers."""
def __init__(self):
super().__init__()
self.register_data_providers = {
"linkedin": LinkedinProvider(),
"yahoo_finance": YahooFinanceProvider(),
"amazon": AmazonProvider(),
"zillow": ZillowProvider(),
"twitter": TwitterProvider()
}
@openapi_schema({
"type": "function",
"function": {
"name": "get_data_provider_endpoints",
"description": "Get available endpoints for a specific data provider",
"parameters": {
"type": "object",
"properties": {
"service_name": {
"type": "string",
"description": "The name of the data provider (e.g., 'linkedin', 'twitter', 'zillow', 'amazon', 'yahoo_finance')"
}
},
"required": ["service_name"]
}
}
})
@xml_schema(
tag_name="get-data-provider-endpoints",
mappings=[
{"param_name": "service_name", "node_type": "attribute", "path": "."}
],
example='''
<!--
The get-data-provider-endpoints tool returns available endpoints for a specific data provider.
Use this tool when you need to discover what endpoints are available.
-->
<!-- Example to get LinkedIn API endpoints -->
<function_calls>
<invoke name="get_data_provider_endpoints">
<parameter name="service_name">linkedin</parameter>
</invoke>
</function_calls>
'''
)
async def get_data_provider_endpoints(
self,
service_name: str
) -> ToolResult:
"""
Get available endpoints for a specific data provider.
Parameters:
- service_name: The name of the data provider (e.g., 'linkedin')
"""
try:
if not service_name:
return self.fail_response("Data provider name is required.")
if service_name not in self.register_data_providers:
return self.fail_response(f"Data provider '{service_name}' not found. Available data providers: {list(self.register_data_providers.keys())}")
endpoints = self.register_data_providers[service_name].get_endpoints()
return self.success_response(endpoints)
except Exception as e:
error_message = str(e)
simplified_message = f"Error getting data provider endpoints: {error_message[:200]}"
if len(error_message) > 200:
simplified_message += "..."
return self.fail_response(simplified_message)
@openapi_schema({
"type": "function",
"function": {
"name": "execute_data_provider_call",
"description": "Execute a call to a specific data provider endpoint",
"parameters": {
"type": "object",
"properties": {
"service_name": {
"type": "string",
"description": "The name of the API service (e.g., 'linkedin')"
},
"route": {
"type": "string",
"description": "The key of the endpoint to call"
},
"payload": {
"type": "object",
"description": "The payload to send with the API call"
}
},
"required": ["service_name", "route"]
}
}
})
@xml_schema(
tag_name="execute-data-provider-call",
mappings=[
{"param_name": "service_name", "node_type": "attribute", "path": "service_name"},
{"param_name": "route", "node_type": "attribute", "path": "route"},
{"param_name": "payload", "node_type": "content", "path": "."}
],
example='''
<!--
The execute-data-provider-call tool makes a request to a specific data provider endpoint.
Use this tool when you need to call an data provider endpoint with specific parameters.
The route must be a valid endpoint key obtained from get-data-provider-endpoints tool!!
-->
<!-- Example to call linkedIn service with the specific route person -->
<function_calls>
<invoke name="execute_data_provider_call">
<parameter name="service_name">linkedin</parameter>
<parameter name="route">person</parameter>
<parameter name="payload">{"link": "https://www.linkedin.com/in/johndoe/"}</parameter>
</invoke>
</function_calls>
'''
)
async def execute_data_provider_call(
self,
service_name: str,
route: str,
payload: Union[Dict[str, Any], str, None] = None
) -> ToolResult:
"""
Execute a call to a specific data provider endpoint.
Parameters:
- service_name: The name of the data provider (e.g., 'linkedin')
- route: The key of the endpoint to call
- payload: The payload to send with the data provider call (dict or JSON string)
"""
try:
# Handle payload - it can be either a dict or a JSON string
if isinstance(payload, str):
try:
payload = json.loads(payload)
except json.JSONDecodeError as e:
return self.fail_response(f"Invalid JSON in payload: {str(e)}")
elif payload is None:
payload = {}
# If payload is already a dict, use it as-is
if not service_name:
return self.fail_response("service_name is required.")
if not route:
return self.fail_response("route is required.")
if service_name not in self.register_data_providers:
return self.fail_response(f"API '{service_name}' not found. Available APIs: {list(self.register_data_providers.keys())}")
data_provider = self.register_data_providers[service_name]
if route == service_name:
return self.fail_response(f"route '{route}' is the same as service_name '{service_name}'. YOU FUCKING IDIOT!")
if route not in data_provider.get_endpoints().keys():
return self.fail_response(f"Endpoint '{route}' not found in {service_name} data provider.")
result = data_provider.call_endpoint(route, payload)
return self.success_response(result)
except Exception as e:
error_message = str(e)
print(error_message)
simplified_message = f"Error executing data provider call: {error_message[:200]}"
if len(error_message) > 200:
simplified_message += "..."
return self.fail_response(simplified_message)
app/tool/expand_msg_tool.py
-103
View File
@@ -1,103 +0,0 @@
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
from agentpress.thread_manager import ThreadManager
import json
class ExpandMessageTool(Tool):
"""Tool for expanding a previous message to the user."""
def __init__(self, thread_id: str, thread_manager: ThreadManager):
super().__init__()
self.thread_manager = thread_manager
self.thread_id = thread_id
@openapi_schema({
"type": "function",
"function": {
"name": "expand_message",
"description": "Expand a message from the previous conversation with the user. Use this tool to expand a message that was truncated in the earlier conversation.",
"parameters": {
"type": "object",
"properties": {
"message_id": {
"type": "string",
"description": "The ID of the message to expand. Must be a UUID."
}
},
"required": ["message_id"]
}
}
})
@xml_schema(
tag_name="expand-message",
mappings=[
{"param_name": "message_id", "node_type": "attribute", "path": "."}
],
example='''
<!-- Example 1: Expand a message that was truncated in the previous conversation -->
<function_calls>
<invoke name="expand_message">
<parameter name="message_id">ecde3a4c-c7dc-4776-ae5c-8209517c5576</parameter>
</invoke>
</function_calls>
<!-- Example 2: Expand a message to create reports or analyze truncated data -->
<function_calls>
<invoke name="expand_message">
<parameter name="message_id">f47ac10b-58cc-4372-a567-0e02b2c3d479</parameter>
</invoke>
</function_calls>
<!-- Example 3: Expand a message when you need the full content for analysis -->
<function_calls>
<invoke name="expand_message">
<parameter name="message_id">550e8400-e29b-41d4-a716-446655440000</parameter>
</invoke>
</function_calls>
'''
)
async def expand_message(self, message_id: str) -> ToolResult:
"""Expand a message from the previous conversation with the user.
Args:
message_id: The ID of the message to expand
Returns:
ToolResult indicating the message was successfully expanded
"""
try:
client = await self.thread_manager.db.client
message = await client.table('messages').select('*').eq('message_id', message_id).eq('thread_id', self.thread_id).execute()
if not message.data or len(message.data) == 0:
return self.fail_response(f"Message with ID {message_id} not found in thread {self.thread_id}")
message_data = message.data[0]
message_content = message_data['content']
final_content = message_content
if isinstance(message_content, dict) and 'content' in message_content:
final_content = message_content['content']
elif isinstance(message_content, str):
try:
parsed_content = json.loads(message_content)
if isinstance(parsed_content, dict) and 'content' in parsed_content:
final_content = parsed_content['content']
except json.JSONDecodeError:
pass
return self.success_response({"status": "Message expanded successfully.", "message": final_content})
except Exception as e:
return self.fail_response(f"Error expanding message: {str(e)}")
if __name__ == "__main__":
import asyncio
async def test_expand_message_tool():
expand_message_tool = ExpandMessageTool()
# Test expand message
expand_message_result = await expand_message_tool.expand_message(
message_id="004ab969-ef9a-4656-8aba-e392345227cd"
)
print("Expand message result:", expand_message_result)
asyncio.run(test_expand_message_tool())
app/tool/mcp_tool_wrapper.py
-715
View File
@@ -1,715 +0,0 @@
"""
MCP Tool Wrapper for AgentPress
This module provides a generic tool wrapper that handles all MCP (Model Context Protocol)
server tool calls through dynamically generated individual function methods.
"""
import json
from typing import Any, Dict, List, Optional
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema, ToolSchema, SchemaType
from mcp_local.client import MCPManager
from utils.logger import logger
import inspect
from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client
from mcp import StdioServerParameters
import asyncio
class MCPToolWrapper(Tool):
"""
A generic tool wrapper that dynamically creates individual methods for each MCP tool.
This tool creates separate function calls for each MCP tool while routing them all
through the same underlying implementation.
"""
def __init__(self, mcp_configs: Optional[List[Dict[str, Any]]] = None):
"""
Initialize the MCP tool wrapper.
Args:
mcp_configs: List of MCP configurations from agent's configured_mcps
"""
# Don't call super().__init__() yet - we need to set up dynamic methods first
self.mcp_manager = MCPManager()
self.mcp_configs = mcp_configs or []
self._initialized = False
self._dynamic_tools = {}
self._schemas: Dict[str, List[ToolSchema]] = {}
self._custom_tools = {} # Store custom MCP tools separately
# Now initialize the parent class which will call _register_schemas
super().__init__()
async def _ensure_initialized(self):
"""Ensure MCP servers are initialized."""
if not self._initialized:
# Initialize standard MCP servers from Smithery
standard_configs = [cfg for cfg in self.mcp_configs if not cfg.get('isCustom', False)]
custom_configs = [cfg for cfg in self.mcp_configs if cfg.get('isCustom', False)]
# Initialize standard MCPs through MCPManager
if standard_configs:
for config in standard_configs:
try:
logger.info(f"Attempting to connect to MCP server: {config['qualifiedName']}")
await self.mcp_manager.connect_server(config)
logger.info(f"Successfully connected to MCP server: {config['qualifiedName']}")
except Exception as e:
logger.error(f"Failed to connect to MCP server {config['qualifiedName']}: {e}")
import traceback
logger.error(f"Full traceback: {traceback.format_exc()}")
# Initialize custom MCPs directly
if custom_configs:
await self._initialize_custom_mcps(custom_configs)
# Create dynamic tools for all connected servers
await self._create_dynamic_tools()
self._initialized = True
async def _connect_sse_server(self, server_name, server_config, all_tools, timeout):
url = server_config["url"]
headers = server_config.get("headers", {})
async with asyncio.timeout(timeout):
try:
async with sse_client(url, headers=headers) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
tools_info = []
for tool in tools_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"input_schema": tool.inputSchema
}
tools_info.append(tool_info)
all_tools[server_name] = {
"status": "connected",
"transport": "sse",
"url": url,
"tools": tools_info
}
logger.info(f" {server_name}: Connected via SSE ({len(tools_info)} tools)")
except TypeError as e:
if "unexpected keyword argument" in str(e):
async with sse_client(url) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
tools_info = []
for tool in tools_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"input_schema": tool.inputSchema
}
tools_info.append(tool_info)
all_tools[server_name] = {
"status": "connected",
"transport": "sse",
"url": url,
"tools": tools_info
}
logger.info(f" {server_name}: Connected via SSE ({len(tools_info)} tools)")
else:
raise
async def _connect_streamable_http_server(self, url):
async with streamablehttp_client(url) as (
read_stream,
write_stream,
_,
):
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
tool_result = await session.list_tools()
print(f"Connected via HTTP ({len(tool_result.tools)} tools)")
tools_info = []
for tool in tool_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"inputSchema": tool.inputSchema
}
tools_info.append(tool_info)
return tools_info
async def _connect_stdio_server(self, server_name, server_config, all_tools, timeout):
"""Connect to a stdio-based MCP server."""
server_params = StdioServerParameters(
command=server_config["command"],
args=server_config.get("args", []),
env=server_config.get("env", {})
)
async with asyncio.timeout(timeout):
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools_result = await session.list_tools()
tools_info = []
for tool in tools_result.tools:
tool_info = {
"name": tool.name,
"description": tool.description,
"input_schema": tool.inputSchema
}
tools_info.append(tool_info)
all_tools[server_name] = {
"status": "connected",
"transport": "stdio",
"tools": tools_info
}
logger.info(f" {server_name}: Connected via stdio ({len(tools_info)} tools)")
async def _initialize_custom_mcps(self, custom_configs):
"""Initialize custom MCP servers."""
for config in custom_configs:
try:
logger.info(f"Initializing custom MCP: {config}")
custom_type = config.get('customType', 'sse')
server_config = config.get('config', {})
enabled_tools = config.get('enabledTools', [])
server_name = config.get('name', 'Unknown')
logger.info(f"Initializing custom MCP: {server_name} (type: {custom_type})")
if custom_type == 'sse':
if 'url' not in server_config:
logger.error(f"Custom MCP {server_name}: Missing 'url' in config")
continue
url = server_config['url']
logger.info(f"Initializing custom MCP {url} with SSE type")
try:
# Use the working connect_sse_server method
all_tools = {}
await self._connect_sse_server(server_name, server_config, all_tools, 15)
# Process the results
if server_name in all_tools and all_tools[server_name].get('status') == 'connected':
tools_info = all_tools[server_name].get('tools', [])
tools_registered = 0
for tool_info in tools_info:
tool_name_from_server = tool_info['name']
if not enabled_tools or tool_name_from_server in enabled_tools:
tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}"
self._custom_tools[tool_name] = {
'name': tool_name,
'description': tool_info['description'],
'parameters': tool_info['input_schema'],
'server': server_name,
'original_name': tool_name_from_server,
'is_custom': True,
'custom_type': custom_type,
'custom_config': server_config
}
tools_registered += 1
logger.debug(f"Registered custom tool: {tool_name}")
logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools")
else:
logger.error(f"Failed to connect to custom MCP {server_name}")
except Exception as e:
logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}")
continue
elif custom_type == 'http':
if 'url' not in server_config:
logger.error(f"Custom MCP {server_name}: Missing 'url' in config")
continue
url = server_config['url']
logger.info(f"Initializing custom MCP {url} with HTTP type")
try:
tools_info = await self._connect_streamable_http_server(url)
tools_registered = 0
for tool_info in tools_info:
tool_name_from_server = tool_info['name']
if not enabled_tools or tool_name_from_server in enabled_tools:
tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}"
self._custom_tools[tool_name] = {
'name': tool_name,
'description': tool_info['description'],
'parameters': tool_info['inputSchema'],
'server': server_name,
'original_name': tool_name_from_server,
'is_custom': True,
'custom_type': custom_type,
'custom_config': server_config
}
tools_registered += 1
logger.debug(f"Registered custom tool: {tool_name}")
logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools")
except Exception as e:
logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}")
continue
elif custom_type == 'json':
if 'command' not in server_config:
logger.error(f"Custom MCP {server_name}: Missing 'command' in config")
continue
logger.info(f"Initializing custom MCP {server_name} with JSON/stdio type")
try:
# Use the stdio connection method
all_tools = {}
await self._connect_stdio_server(server_name, server_config, all_tools, 15)
# Process the results
if server_name in all_tools and all_tools[server_name].get('status') == 'connected':
tools_info = all_tools[server_name].get('tools', [])
tools_registered = 0
for tool_info in tools_info:
tool_name_from_server = tool_info['name']
if not enabled_tools or tool_name_from_server in enabled_tools:
tool_name = f"custom_{server_name.replace(' ', '_').lower()}_{tool_name_from_server}"
self._custom_tools[tool_name] = {
'name': tool_name,
'description': tool_info['description'],
'parameters': tool_info['input_schema'],
'server': server_name,
'original_name': tool_name_from_server,
'is_custom': True,
'custom_type': custom_type,
'custom_config': server_config
}
tools_registered += 1
logger.debug(f"Registered custom tool: {tool_name}")
logger.info(f"Successfully initialized custom MCP {server_name} with {tools_registered} tools")
else:
logger.error(f"Failed to connect to custom MCP {server_name}")
except Exception as e:
logger.error(f"Custom MCP {server_name}: Connection failed - {str(e)}")
continue
else:
logger.error(f"Custom MCP {server_name}: Unsupported type '{custom_type}', supported types are 'sse', 'http' and 'json'")
continue
except Exception as e:
logger.error(f"Failed to initialize custom MCP {config.get('name', 'Unknown')}: {e}")
continue
async def initialize_and_register_tools(self, tool_registry=None):
"""Initialize MCP tools and optionally update the tool registry.
This method should be called after the tool has been registered to dynamically
add the MCP tool schemas to the registry.
Args:
tool_registry: Optional ToolRegistry instance to update with new schemas
"""
await self._ensure_initialized()
if tool_registry and self._dynamic_tools:
logger.info(f"Updating tool registry with {len(self._dynamic_tools)} MCP tools")
for method_name, schemas in self._schemas.items():
if method_name not in ['call_mcp_tool']: # Skip the fallback method
pass
async def _create_dynamic_tools(self):
"""Create dynamic tool methods for each available MCP tool."""
try:
# Get standard MCP tools
available_tools = self.mcp_manager.get_all_tools_openapi()
logger.info(f"MCPManager returned {len(available_tools)} tools")
for tool_info in available_tools:
tool_name = tool_info.get('name', '')
logger.info(f"Processing tool: {tool_name}")
if tool_name:
# Create a dynamic method for this tool with proper OpenAI schema
self._create_dynamic_method(tool_name, tool_info)
# Get custom MCP tools
logger.info(f"Processing {len(self._custom_tools)} custom MCP tools")
for tool_name, tool_info in self._custom_tools.items():
logger.info(f"Processing custom tool: {tool_name}")
# Convert custom tool info to the expected format
openapi_tool_info = {
"name": tool_name,
"description": tool_info['description'],
"parameters": tool_info['parameters']
}
self._create_dynamic_method(tool_name, openapi_tool_info)
logger.info(f"Created {len(self._dynamic_tools)} dynamic MCP tool methods")
except Exception as e:
logger.error(f"Error creating dynamic MCP tools: {e}")
def _create_dynamic_method(self, tool_name: str, tool_info: Dict[str, Any]):
"""Create a dynamic method for a specific MCP tool with proper OpenAI schema."""
if tool_name.startswith("custom_"):
if tool_name in self._custom_tools:
clean_tool_name = self._custom_tools[tool_name]['original_name']
server_name = self._custom_tools[tool_name]['server']
else:
parts = tool_name.split("_")
if len(parts) >= 3:
clean_tool_name = "_".join(parts[2:])
server_name = parts[1] if len(parts) > 1 else "unknown"
else:
clean_tool_name = tool_name
server_name = "unknown"
else:
parts = tool_name.split("_", 2)
clean_tool_name = parts[2] if len(parts) > 2 else tool_name
server_name = parts[1] if len(parts) > 1 else "unknown"
method_name = clean_tool_name.replace('-', '_')
logger.info(f"Creating dynamic method for tool '{tool_name}': clean_tool_name='{clean_tool_name}', method_name='{method_name}', server='{server_name}'")
original_full_name = tool_name
# Create the dynamic method
async def dynamic_tool_method(**kwargs) -> ToolResult:
"""Dynamically created method for MCP tool."""
# Use the original full tool name for execution
return await self._execute_mcp_tool(original_full_name, kwargs)
# Set the method name to match the tool name
dynamic_tool_method.__name__ = method_name
dynamic_tool_method.__qualname__ = f"{self.__class__.__name__}.{method_name}"
# Build a more descriptive description
base_description = tool_info.get("description", f"MCP tool from {server_name}")
full_description = f"{base_description} (MCP Server: {server_name})"
# Create the OpenAI schema for this tool
openapi_function_schema = {
"type": "function",
"function": {
"name": method_name, # Use the clean method name for function calling
"description": full_description,
"parameters": tool_info.get("parameters", {
"type": "object",
"properties": {},
"required": []
})
}
}
# Create a ToolSchema object
tool_schema = ToolSchema(
schema_type=SchemaType.OPENAPI,
schema=openapi_function_schema
)
# Add the schema to our schemas dict
self._schemas[method_name] = [tool_schema]
# Also add the schema to the method itself (for compatibility)
dynamic_tool_method.tool_schemas = [tool_schema]
# Store the method and its info
self._dynamic_tools[tool_name] = {
'method': dynamic_tool_method,
'method_name': method_name,
'original_tool_name': tool_name,
'clean_tool_name': clean_tool_name,
'server_name': server_name,
'info': tool_info,
'schema': tool_schema
}
# Add the method to this instance
setattr(self, method_name, dynamic_tool_method)
logger.debug(f"Created dynamic method '{method_name}' for MCP tool '{tool_name}' from server '{server_name}'")
def _register_schemas(self):
"""Register schemas from all decorated methods and dynamic tools."""
# First register static schemas from decorated methods
for name, method in inspect.getmembers(self, predicate=inspect.ismethod):
if hasattr(method, 'tool_schemas'):
self._schemas[name] = method.tool_schemas
logger.debug(f"Registered schemas for method '{name}' in {self.__class__.__name__}")
# Note: Dynamic schemas will be added after async initialization
logger.debug(f"Initial registration complete for MCPToolWrapper")
def get_schemas(self) -> Dict[str, List[ToolSchema]]:
"""Get all registered tool schemas including dynamic ones."""
# Return all schemas including dynamically added ones
return self._schemas
def __getattr__(self, name: str):
"""Handle calls to dynamically created MCP tool methods."""
# Look for exact method name match first
for tool_data in self._dynamic_tools.values():
if tool_data['method_name'] == name:
return tool_data['method']
# Try with underscore/hyphen conversion
name_with_hyphens = name.replace('_', '-')
for tool_name, tool_data in self._dynamic_tools.items():
if tool_data['method_name'] == name or tool_name == name_with_hyphens:
return tool_data['method']
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
async def get_available_tools(self) -> List[Dict[str, Any]]:
"""Get all available MCP tools in OpenAPI format."""
await self._ensure_initialized()
return self.mcp_manager.get_all_tools_openapi()
async def _execute_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> ToolResult:
"""Execute an MCP tool call."""
await self._ensure_initialized()
logger.info(f"Executing MCP tool {tool_name} with arguments {arguments}")
try:
# Check if it's a custom MCP tool first
if tool_name in self._custom_tools:
tool_info = self._custom_tools[tool_name]
return await self._execute_custom_mcp_tool(tool_name, arguments, tool_info)
else:
# Use standard MCP manager for Smithery servers
result = await self.mcp_manager.execute_tool(tool_name, arguments)
if isinstance(result, dict):
if result.get('isError', False):
return self.fail_response(result.get('content', 'Tool execution failed'))
else:
return self.success_response(result.get('content', result))
else:
return self.success_response(result)
except Exception as e:
logger.error(f"Error executing MCP tool {tool_name}: {str(e)}")
return self.fail_response(f"Error executing tool: {str(e)}")
async def _execute_custom_mcp_tool(self, tool_name: str, arguments: Dict[str, Any], tool_info: Dict[str, Any]) -> ToolResult:
"""Execute a custom MCP tool call."""
try:
custom_type = tool_info['custom_type']
custom_config = tool_info['custom_config']
original_tool_name = tool_info['original_name']
if custom_type == 'sse':
# Execute SSE-based custom MCP using the same pattern as _connect_sse_server
url = custom_config['url']
headers = custom_config.get('headers', {})
async with asyncio.timeout(30): # 30 second timeout for tool execution
try:
# Try with headers first (same pattern as _connect_sse_server)
async with sse_client(url, headers=headers) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(original_tool_name, arguments)
# Handle the result properly
if hasattr(result, 'content'):
content = result.content
if isinstance(content, list):
# Extract text from content list
text_parts = []
for item in content:
if hasattr(item, 'text'):
text_parts.append(item.text)
else:
text_parts.append(str(item))
content_str = "\n".join(text_parts)
elif hasattr(content, 'text'):
content_str = content.text
else:
content_str = str(content)
return self.success_response(content_str)
else:
return self.success_response(str(result))
except TypeError as e:
if "unexpected keyword argument" in str(e):
# Fallback: try without headers (exact pattern from _connect_sse_server)
async with sse_client(url) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(original_tool_name, arguments)
# Handle the result properly
if hasattr(result, 'content'):
content = result.content
if isinstance(content, list):
# Extract text from content list
text_parts = []
for item in content:
if hasattr(item, 'text'):
text_parts.append(item.text)
else:
text_parts.append(str(item))
content_str = "\n".join(text_parts)
elif hasattr(content, 'text'):
content_str = content.text
else:
content_str = str(content)
return self.success_response(content_str)
else:
return self.success_response(str(result))
else:
raise
elif custom_type == 'http':
# Execute HTTP-based custom MCP
url = custom_config['url']
async with asyncio.timeout(30): # 30 second timeout for tool execution
async with streamablehttp_client(url) as (read, write, _):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(original_tool_name, arguments)
# Handle the result properly
if hasattr(result, 'content'):
content = result.content
if isinstance(content, list):
# Extract text from content list
text_parts = []
for item in content:
if hasattr(item, 'text'):
text_parts.append(item.text)
else:
text_parts.append(str(item))
content_str = "\n".join(text_parts)
elif hasattr(content, 'text'):
content_str = content.text
else:
content_str = str(content)
return self.success_response(content_str)
else:
return self.success_response(str(result))
elif custom_type == 'json':
# Execute stdio-based custom MCP using the same pattern as _connect_stdio_server
server_params = StdioServerParameters(
command=custom_config["command"],
args=custom_config.get("args", []),
env=custom_config.get("env", {})
)
async with asyncio.timeout(30): # 30 second timeout for tool execution
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
result = await session.call_tool(original_tool_name, arguments)
# Handle the result properly
if hasattr(result, 'content'):
content = result.content
if isinstance(content, list):
# Extract text from content list
text_parts = []
for item in content:
if hasattr(item, 'text'):
text_parts.append(item.text)
else:
text_parts.append(str(item))
content_str = "\n".join(text_parts)
elif hasattr(content, 'text'):
content_str = content.text
else:
content_str = str(content)
return self.success_response(content_str)
else:
return self.success_response(str(result))
else:
return self.fail_response(f"Unsupported custom MCP type: {custom_type}")
except asyncio.TimeoutError:
return self.fail_response(f"Tool execution timeout for {tool_name}")
except Exception as e:
logger.error(f"Error executing custom MCP tool {tool_name}: {str(e)}")
return self.fail_response(f"Error executing custom tool: {str(e)}")
# Keep the original call_mcp_tool method as a fallback
@openapi_schema({
"type": "function",
"function": {
"name": "call_mcp_tool",
"description": "Execute a tool from any connected MCP server. This is a fallback wrapper that forwards calls to MCP tools. The tool_name should be in the format 'mcp_{server}_{tool}' where {server} is the MCP server's qualified name and {tool} is the specific tool name.",
"parameters": {
"type": "object",
"properties": {
"tool_name": {
"type": "string",
"description": "The full MCP tool name in format 'mcp_{server}_{tool}', e.g., 'mcp_exa_web_search_exa'"
},
"arguments": {
"type": "object",
"description": "The arguments to pass to the MCP tool, as a JSON object. The required arguments depend on the specific tool being called.",
"additionalProperties": True
}
},
"required": ["tool_name", "arguments"]
}
}
})
@xml_schema(
tag_name="call-mcp-tool",
mappings=[
{"param_name": "tool_name", "node_type": "attribute", "path": "."},
{"param_name": "arguments", "node_type": "content", "path": "."}
],
example='''
<function_calls>
<invoke name="call_mcp_tool">
<parameter name="tool_name">mcp_exa_web_search_exa</parameter>
<parameter name="arguments">{"query": "latest developments in AI", "num_results": 10}</parameter>
</invoke>
</function_calls>
'''
)
async def call_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> ToolResult:
"""
Execute an MCP tool call (fallback method).
Args:
tool_name: The full MCP tool name (e.g., "mcp_exa_web_search_exa")
arguments: The arguments to pass to the tool
Returns:
ToolResult with the tool execution result
"""
return await self._execute_mcp_tool(tool_name, arguments)
async def cleanup(self):
"""Disconnect all MCP servers."""
if self._initialized:
try:
await self.mcp_manager.disconnect_all()
except Exception as e:
logger.error(f"Error during MCP cleanup: {str(e)}")
finally:
self._initialized = False
app/tool/message_tool.py
-270
View File
@@ -1,270 +0,0 @@
from typing import List, Optional, Union
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
from utils.logger import logger
class MessageTool(Tool):
"""Tool for user communication and interaction.
This tool provides methods for asking questions, with support for
attachments and user takeover suggestions.
"""
def __init__(self):
super().__init__()
# Commented out as we are just doing this via prompt as there is no need to call it as a tool
@openapi_schema({
"type": "function",
"function": {
"name": "ask",
"description": "Ask user a question and wait for response. Use for: 1) Requesting clarification on ambiguous requirements, 2) Seeking confirmation before proceeding with high-impact changes, 3) Gathering additional information needed to complete a task, 4) Offering options and requesting user preference, 5) Validating assumptions when critical to task success. IMPORTANT: Use this tool only when user input is essential to proceed. Always provide clear context and options when applicable. Include relevant attachments when the question relates to specific files or resources.",
"parameters": {
"type": "object",
"properties": {
"text": {
"type": "string",
"description": "Question text to present to user - should be specific and clearly indicate what information you need. Include: 1) Clear question or request, 2) Context about why the input is needed, 3) Available options if applicable, 4) Impact of different choices, 5) Any relevant constraints or considerations."
},
"attachments": {
"anyOf": [
{"type": "string"},
{"items": {"type": "string"}, "type": "array"}
],
"description": "(Optional) List of files or URLs to attach to the question. Include when: 1) Question relates to specific files or configurations, 2) User needs to review content before answering, 3) Options or choices are documented in files, 4) Supporting evidence or context is needed. Always use relative paths to /workspace directory."
}
},
"required": ["text"]
}
}
})
@xml_schema(
tag_name="ask",
mappings=[
{"param_name": "text", "node_type": "content", "path": "."},
{"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False}
],
example='''
<function_calls>
<invoke name="ask">
<parameter name="text">I'm planning to bake the chocolate cake for your birthday party. The recipe mentions "rich frosting" but doesn't specify what type. Could you clarify your preferences? For example:
1. Would you prefer buttercream or cream cheese frosting?
2. Do you want any specific flavor added to the frosting (vanilla, coffee, etc.)?
3. Should I add any decorative toppings like sprinkles or fruit?
4. Do you have any dietary restrictions I should be aware of?
This information will help me make sure the cake meets your expectations for the celebration.</parameter>
<parameter name="attachments">recipes/chocolate_cake.txt,photos/cake_examples.jpg</parameter>
</invoke>
</function_calls>
'''
)
async def ask(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult:
"""Ask the user a question and wait for a response.
Args:
text: The question to present to the user
attachments: Optional file paths or URLs to attach to the question
Returns:
ToolResult indicating the question was successfully sent
"""
try:
# Convert single attachment to list for consistent handling
if attachments and isinstance(attachments, str):
attachments = [attachments]
return self.success_response({"status": "Awaiting user response..."})
except Exception as e:
return self.fail_response(f"Error asking user: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "web_browser_takeover",
"description": "Request user takeover of browser interaction. Use this tool when: 1) The page requires complex human interaction that automated tools cannot handle, 2) Authentication or verification steps require human input, 3) The page has anti-bot measures that prevent automated access, 4) Complex form filling or navigation is needed, 5) The page requires human verification (CAPTCHA, etc.). IMPORTANT: This tool should be used as a last resort after web-search and crawl-webpage have failed, and when direct browser tools are insufficient. Always provide clear context about why takeover is needed and what actions the user should take.",
"parameters": {
"type": "object",
"properties": {
"text": {
"type": "string",
"description": "Instructions for the user about what actions to take in the browser. Include: 1) Clear explanation of why takeover is needed, 2) Specific steps the user should take, 3) What information to look for or extract, 4) How to indicate when they're done, 5) Any important context about the current page state."
},
"attachments": {
"anyOf": [
{"type": "string"},
{"items": {"type": "string"}, "type": "array"}
],
"description": "(Optional) List of files or URLs to attach to the takeover request. Include when: 1) Screenshots or visual references are needed, 2) Previous search results or crawled content is relevant, 3) Supporting documentation is required. Always use relative paths to /workspace directory."
}
},
"required": ["text"]
}
}
})
@xml_schema(
tag_name="web-browser-takeover",
mappings=[
{"param_name": "text", "node_type": "content", "path": "."},
{"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False}
],
example='''
<function_calls>
<invoke name="web_browser_takeover">
<parameter name="text">I've encountered a CAPTCHA verification on the page. Please:
1. Solve the CAPTCHA puzzle
2. Let me know once you've completed it
3. I'll then continue with the automated process
If you encounter any issues or need to take additional steps, please let me know.</parameter>
</invoke>
</function_calls>
'''
)
async def web_browser_takeover(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult:
"""Request user takeover of browser interaction.
Args:
text: Instructions for the user about what actions to take
attachments: Optional file paths or URLs to attach to the request
Returns:
ToolResult indicating the takeover request was successfully sent
"""
try:
# Convert single attachment to list for consistent handling
if attachments and isinstance(attachments, str):
attachments = [attachments]
return self.success_response({"status": "Awaiting user browser takeover..."})
except Exception as e:
return self.fail_response(f"Error requesting browser takeover: {str(e)}")
# @openapi_schema({
# "type": "function",
# "function": {
# "name": "inform",
# "description": "Inform the user about progress, completion of a major step, or important context. Use this tool: 1) To provide updates between major sections of work, 2) After accomplishing significant milestones, 3) When transitioning to a new phase of work, 4) To confirm actions were completed successfully, 5) To provide context about upcoming steps. IMPORTANT: Use FREQUENTLY throughout execution to provide UI context to the user. The user CANNOT respond to this tool - they can only respond to the 'ask' tool. Use this tool to keep the user informed without requiring their input.",
# "parameters": {
# "type": "object",
# "properties": {
# "text": {
# "type": "string",
# "description": "Information to present to the user. Include: 1) Clear statement of what has been accomplished or what is happening, 2) Relevant context or impact, 3) Brief indication of next steps if applicable."
# },
# "attachments": {
# "anyOf": [
# {"type": "string"},
# {"items": {"type": "string"}, "type": "array"}
# ],
# "description": "(Optional) List of files or URLs to attach to the information. Include when: 1) Information relates to specific files or resources, 2) Showing intermediate results or outputs, 3) Providing supporting documentation. Always use relative paths to /workspace directory."
# }
# },
# "required": ["text"]
# }
# }
# })
# @xml_schema(
# tag_name="inform",
# mappings=[
# {"param_name": "text", "node_type": "content", "path": "."},
# {"param_name": "attachments", "node_type": "attribute", "path": ".", "required": False}
# ],
# example='''
# Inform the user about progress, completion of a major step, or important context. Use this tool: 1) To provide updates between major sections of work, 2) After accomplishing significant milestones, 3) When transitioning to a new phase of work, 4) To confirm actions were completed successfully, 5) To provide context about upcoming steps. IMPORTANT: Use FREQUENTLY throughout execution to provide UI context to the user. The user CANNOT respond to this tool - they can only respond to the 'ask' tool. Use this tool to keep the user informed without requiring their input."
# <!-- Use inform FREQUENTLY to provide UI context and progress updates - THE USER CANNOT RESPOND to this tool -->
# <!-- The user can ONLY respond to the ask tool, not to inform -->
# <!-- Examples of when to use inform: -->
# <!-- 1. Completing major milestones -->
# <!-- 2. Transitioning between work phases -->
# <!-- 3. Confirming important actions -->
# <!-- 4. Providing context about upcoming steps -->
# <!-- 5. Sharing significant intermediate results -->
# <!-- 6. Providing regular UI updates throughout execution -->
# <inform attachments="analysis_results.csv,summary_chart.png">
# I've completed the data analysis of the sales figures. Key findings include:
# - Q4 sales were 28% higher than Q3
# - Product line A showed the strongest performance
# - Three regions missed their targets
# I'll now proceed with creating the executive summary report based on these findings.
# </inform>
# '''
# )
# async def inform(self, text: str, attachments: Optional[Union[str, List[str]]] = None) -> ToolResult:
# """Inform the user about progress or important updates without requiring a response.
# Args:
# text: The information to present to the user
# attachments: Optional file paths or URLs to attach
# Returns:
# ToolResult indicating the information was successfully sent
# """
# try:
# # Convert single attachment to list for consistent handling
# if attachments and isinstance(attachments, str):
# attachments = [attachments]
# return self.success_response({"status": "Information sent"})
# except Exception as e:
# return self.fail_response(f"Error informing user: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "complete",
"description": "A special tool to indicate you have completed all tasks and are about to enter complete state. Use ONLY when: 1) All tasks in todo.md are marked complete [x], 2) The user's original request has been fully addressed, 3) There are no pending actions or follow-ups required, 4) You've delivered all final outputs and results to the user. IMPORTANT: This is the ONLY way to properly terminate execution. Never use this tool unless ALL tasks are complete and verified. Always ensure you've provided all necessary outputs and references before using this tool.",
"parameters": {
"type": "object",
"properties": {},
"required": []
}
}
})
@xml_schema(
tag_name="complete",
mappings=[],
example='''
<function_calls>
<invoke name="complete">
</invoke>
</function_calls>
'''
)
async def complete(self) -> ToolResult:
"""Indicate that the agent has completed all tasks and is entering complete state.
Returns:
ToolResult indicating successful transition to complete state
"""
try:
return self.success_response({"status": "complete"})
except Exception as e:
return self.fail_response(f"Error entering complete state: {str(e)}")
if __name__ == "__main__":
import asyncio
async def test_message_tool():
message_tool = MessageTool()
# Test question
ask_result = await message_tool.ask(
text="Would you like to proceed with the next phase?",
attachments="summary.pdf"
)
print("Question result:", ask_result)
# Test inform
inform_result = await message_tool.inform(
text="Completed analysis of data. Processing results now.",
attachments="analysis.pdf"
)
print("Inform result:", inform_result)
asyncio.run(test_message_tool())
app/tool/sb_browser_tool.py
+74 -74
View File
@@ -6,14 +6,14 @@ from PIL import Image
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from agentpress.thread_manager import ThreadManager
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from utils.logger import logger
from utils.s3_upload_utils import upload_base64_image
class SandboxBrowserTool(SandboxToolsBase):
"""Tool for executing tasks in a Daytona sandbox with browser-use capabilities."""
def __init__(self, project_id: str, thread_id: str, thread_manager: ThreadManager):
super().__init__(project_id, thread_manager)
self.thread_id = thread_id
@@ -21,11 +21,11 @@ class SandboxBrowserTool(SandboxToolsBase):
def _validate_base64_image(self, base64_string: str, max_size_mb: int = 10) -> tuple[bool, str]:
"""
Comprehensive validation of base64 image data.
Args:
base64_string (str): The base64 encoded image data
max_size_mb (int): Maximum allowed image size in megabytes
Returns:
tuple[bool, str]: (is_valid, error_message)
"""
@@ -33,94 +33,94 @@ class SandboxBrowserTool(SandboxToolsBase):
# Check if data exists and has reasonable length
if not base64_string or len(base64_string) < 10:
return False, "Base64 string is empty or too short"
# Remove data URL prefix if present (data:image/jpeg;base64,...)
if base64_string.startswith('data:'):
try:
base64_string = base64_string.split(',', 1)[1]
except (IndexError, ValueError):
return False, "Invalid data URL format"
# Check if string contains only valid base64 characters
# Base64 alphabet: A-Z, a-z, 0-9, +, /, = (padding)
import re
if not re.match(r'^[A-Za-z0-9+/]*={0,2}$', base64_string):
return False, "Invalid base64 characters detected"
# Check if base64 string length is valid (must be multiple of 4)
if len(base64_string) % 4 != 0:
return False, "Invalid base64 string length"
# Attempt to decode base64
try:
image_data = base64.b64decode(base64_string, validate=True)
except Exception as e:
return False, f"Base64 decoding failed: {str(e)}"
# Check decoded data size
if len(image_data) == 0:
return False, "Decoded image data is empty"
# Check if decoded data size exceeds limit
max_size_bytes = max_size_mb * 1024 * 1024
if len(image_data) > max_size_bytes:
return False, f"Image size ({len(image_data)} bytes) exceeds limit ({max_size_bytes} bytes)"
# Validate that decoded data is actually a valid image using PIL
try:
image_stream = io.BytesIO(image_data)
with Image.open(image_stream) as img:
# Verify the image by attempting to load it
img.verify()
# Check if image format is supported
supported_formats = {'JPEG', 'PNG', 'GIF', 'BMP', 'WEBP', 'TIFF'}
if img.format not in supported_formats:
return False, f"Unsupported image format: {img.format}"
# Re-open for dimension checks (verify() closes the image)
image_stream.seek(0)
with Image.open(image_stream) as img_check:
width, height = img_check.size
# Check reasonable dimension limits
max_dimension = 8192 # 8K resolution limit
if width > max_dimension or height > max_dimension:
return False, f"Image dimensions ({width}x{height}) exceed limit ({max_dimension}x{max_dimension})"
# Check minimum dimensions
if width < 1 or height < 1:
return False, f"Invalid image dimensions: {width}x{height}"
logger.debug(f"Valid image detected: {img.format}, {width}x{height}, {len(image_data)} bytes")
except Exception as e:
return False, f"Invalid image data: {str(e)}"
return True, "Valid image"
except Exception as e:
logger.error(f"Unexpected error during base64 image validation: {e}")
return False, f"Validation error: {str(e)}"
async def _execute_browser_action(self, endpoint: str, params: dict = None, method: str = "POST") -> ToolResult:
"""Execute a browser automation action through the API
Args:
endpoint (str): The API endpoint to call
params (dict, optional): Parameters to send. Defaults to None.
method (str, optional): HTTP method to use. Defaults to "POST".
Returns:
ToolResult: Result of the execution
"""
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# Build the curl command
url = f"http://localhost:8003/api/automation/{endpoint}"
if method == "GET" and params:
query_params = "&".join([f"{k}={v}" for k, v in params.items()])
url = f"{url}?{query_params}"
@@ -130,19 +130,19 @@ class SandboxBrowserTool(SandboxToolsBase):
if params:
json_data = json.dumps(params)
curl_cmd += f" -d '{json_data}'"
logger.debug("\033[95mExecuting curl command:\033[0m")
logger.debug(f"{curl_cmd}")
response = self.sandbox.process.exec(curl_cmd, timeout=30)
if response.exit_code == 0:
try:
result = json.loads(response.result)
if not "content" in result:
result["content"] = ""
if not "role" in result:
result["role"] = "assistant"
@@ -153,7 +153,7 @@ class SandboxBrowserTool(SandboxToolsBase):
# Comprehensive validation of the base64 image data
screenshot_data = result["screenshot_base64"]
is_valid, validation_message = self._validate_base64_image(screenshot_data)
if is_valid:
logger.debug(f"Screenshot validation passed: {validation_message}")
image_url = await upload_base64_image(screenshot_data)
@@ -162,10 +162,10 @@ class SandboxBrowserTool(SandboxToolsBase):
else:
logger.warning(f"Screenshot validation failed: {validation_message}")
result["image_validation_error"] = validation_message
# Remove base64 data from result to keep it clean
del result["screenshot_base64"]
except Exception as e:
logger.error(f"Failed to process screenshot: {e}")
result["image_upload_error"] = str(e)
@@ -251,10 +251,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_navigate_to(self, url: str) -> ToolResult:
"""Navigate to a specific url
Args:
url (str): The url to navigate to
Returns:
dict: Result of the execution
"""
@@ -290,10 +290,10 @@ class SandboxBrowserTool(SandboxToolsBase):
# )
# async def browser_search_google(self, query: str) -> ToolResult:
# """Search Google with the provided query
# Args:
# query (str): The search query to use
# Returns:
# dict: Result of the execution
# """
@@ -323,7 +323,7 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_go_back(self) -> ToolResult:
"""Navigate back in browser history
Returns:
dict: Result of the execution
"""
@@ -361,10 +361,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_wait(self, seconds: int = 3) -> ToolResult:
"""Wait for the specified number of seconds
Args:
seconds (int, optional): Number of seconds to wait. Defaults to 3.
Returns:
dict: Result of the execution
"""
@@ -403,10 +403,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_click_element(self, index: int) -> ToolResult:
"""Click on an element by index
Args:
index (int): The index of the element to click
Returns:
dict: Result of the execution
"""
@@ -451,11 +451,11 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_input_text(self, index: int, text: str) -> ToolResult:
"""Input text into an element
Args:
index (int): The index of the element to input text into
text (str): The text to input
Returns:
dict: Result of the execution
"""
@@ -494,10 +494,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_send_keys(self, keys: str) -> ToolResult:
"""Send keyboard keys
Args:
keys (str): The keys to send (e.g., 'Enter', 'Escape', 'Control+a')
Returns:
dict: Result of the execution
"""
@@ -536,10 +536,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_switch_tab(self, page_id: int) -> ToolResult:
"""Switch to a different browser tab
Args:
page_id (int): The ID of the tab to switch to
Returns:
dict: Result of the execution
"""
@@ -576,10 +576,10 @@ class SandboxBrowserTool(SandboxToolsBase):
# )
# async def browser_open_tab(self, url: str) -> ToolResult:
# """Open a new browser tab with the specified URL
# Args:
# url (str): The URL to open in the new tab
# Returns:
# dict: Result of the execution
# """
@@ -618,10 +618,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_close_tab(self, page_id: int) -> ToolResult:
"""Close a browser tab
Args:
page_id (int): The ID of the tab to close
Returns:
dict: Result of the execution
"""
@@ -658,23 +658,23 @@ class SandboxBrowserTool(SandboxToolsBase):
# )
# async def browser_extract_content(self, goal: str) -> ToolResult:
# """Extract content from the current page based on the provided goal
# Args:
# goal (str): The extraction goal
# Returns:
# dict: Result of the execution
# """
# logger.debug(f"\033[95mExtracting content with goal: {goal}\033[0m")
# result = await self._execute_browser_action("extract_content", {"goal": goal})
# # Format content for better readability
# if result.get("success"):
# logger.debug(f"\033[92mContent extraction successful\033[0m")
# content = result.data.get("content", "")
# url = result.data.get("url", "")
# title = result.data.get("title", "")
# if content:
# content_preview = content[:200] + "..." if len(content) > 200 else content
# logger.debug(f"\033[95mExtracted content from {title} ({url}):\033[0m")
@@ -684,7 +684,7 @@ class SandboxBrowserTool(SandboxToolsBase):
# logger.debug(f"\033[93mNo content extracted from {url}\033[0m")
# else:
# logger.debug(f"\033[91mFailed to extract content: {result.data.get('error', 'Unknown error')}\033[0m")
# return result
@openapi_schema({
@@ -718,10 +718,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_scroll_down(self, amount: int = None) -> ToolResult:
"""Scroll down the page
Args:
amount (int, optional): Pixel amount to scroll. If None, scrolls one page.
Returns:
dict: Result of the execution
"""
@@ -731,7 +731,7 @@ class SandboxBrowserTool(SandboxToolsBase):
logger.debug(f"\033[95mScrolling down by {amount} pixels\033[0m")
else:
logger.debug(f"\033[95mScrolling down one page\033[0m")
return await self._execute_browser_action("scroll_down", params)
@openapi_schema({
@@ -765,10 +765,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_scroll_up(self, amount: int = None) -> ToolResult:
"""Scroll up the page
Args:
amount (int, optional): Pixel amount to scroll. If None, scrolls one page.
Returns:
dict: Result of the execution
"""
@@ -778,7 +778,7 @@ class SandboxBrowserTool(SandboxToolsBase):
logger.debug(f"\033[95mScrolling up by {amount} pixels\033[0m")
else:
logger.debug(f"\033[95mScrolling up one page\033[0m")
return await self._execute_browser_action("scroll_up", params)
@openapi_schema({
@@ -813,10 +813,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_scroll_to_text(self, text: str) -> ToolResult:
"""Scroll to specific text on the page
Args:
text (str): The text to scroll to
Returns:
dict: Result of the execution
"""
@@ -855,10 +855,10 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_get_dropdown_options(self, index: int) -> ToolResult:
"""Get all options from a dropdown element
Args:
index (int): The index of the dropdown element
Returns:
dict: Result of the execution with the dropdown options
"""
@@ -903,11 +903,11 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_select_dropdown_option(self, index: int, text: str) -> ToolResult:
"""Select an option from a dropdown by text
Args:
index (int): The index of the dropdown element
text (str): The text of the option to select
Returns:
dict: Result of the execution
"""
@@ -969,11 +969,11 @@ class SandboxBrowserTool(SandboxToolsBase):
</function_calls>
'''
)
async def browser_drag_drop(self, element_source: str = None, element_target: str = None,
async def browser_drag_drop(self, element_source: str = None, element_target: str = None,
coord_source_x: int = None, coord_source_y: int = None,
coord_target_x: int = None, coord_target_y: int = None) -> ToolResult:
"""Perform drag and drop operation between elements or coordinates
Args:
element_source (str, optional): The source element selector
element_target (str, optional): The target element selector
@@ -981,12 +981,12 @@ class SandboxBrowserTool(SandboxToolsBase):
coord_source_y (int, optional): The source Y coordinate
coord_target_x (int, optional): The target X coordinate
coord_target_y (int, optional): The target Y coordinate
Returns:
dict: Result of the execution
"""
params = {}
if element_source and element_target:
params["element_source"] = element_source
params["element_target"] = element_target
@@ -999,7 +999,7 @@ class SandboxBrowserTool(SandboxToolsBase):
logger.debug(f"\033[95mDragging from coordinates ({coord_source_x}, {coord_source_y}) to ({coord_target_x}, {coord_target_y})\033[0m")
else:
return self.fail_response("Must provide either element selectors or coordinates for drag and drop")
return await self._execute_browser_action("drag_drop", params)
@openapi_schema({
@@ -1040,13 +1040,13 @@ class SandboxBrowserTool(SandboxToolsBase):
)
async def browser_click_coordinates(self, x: int, y: int) -> ToolResult:
"""Click at specific X,Y coordinates on the page
Args:
x (int): The X coordinate to click
y (int): The Y coordinate to click
Returns:
dict: Result of the execution
"""
logger.debug(f"\033[95mClicking at coordinates: ({x}, {y})\033[0m")
return await self._execute_browser_action("click_coordinates", {"x": x, "y": y})
return await self._execute_browser_action("click_coordinates", {"x": x, "y": y})
app/tool/sb_deploy_tool.py
+19 -19
View File
@@ -1,7 +1,7 @@
import os
from dotenv import load_dotenv
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from utils.files_utils import clean_path
from agentpress.thread_manager import ThreadManager
@@ -48,11 +48,11 @@ class SandboxDeployTool(SandboxToolsBase):
{"param_name": "directory_path", "node_type": "attribute", "path": "directory_path"}
],
example='''
<!--
<!--
IMPORTANT: Only use this tool when:
1. The user explicitly requests permanent deployment to production
2. You have a complete, ready-to-deploy directory
2. You have a complete, ready-to-deploy directory
NOTE: If the same name is used, it will redeploy to the same project as before
-->
@@ -68,11 +68,11 @@ class SandboxDeployTool(SandboxToolsBase):
"""
Deploy a static website (HTML+CSS+JS) from the sandbox to Cloudflare Pages.
Only use this tool when permanent deployment to a production environment is needed.
Args:
name: Name for the deployment, will be used in the URL as {name}.kortix.cloud
directory_path: Path to the directory to deploy, relative to /workspace
Returns:
ToolResult containing:
- Success: Deployment information including URL
@@ -81,10 +81,10 @@ class SandboxDeployTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
directory_path = self.clean_path(directory_path)
full_path = f"{self.workspace_path}/{directory_path}"
# Verify the directory exists
try:
dir_info = self.sandbox.fs.get_file_info(full_path)
@@ -92,26 +92,26 @@ class SandboxDeployTool(SandboxToolsBase):
return self.fail_response(f"'{directory_path}' is not a directory")
except Exception as e:
return self.fail_response(f"Directory '{directory_path}' does not exist: {str(e)}")
# Deploy to Cloudflare Pages directly from the container
try:
# Get Cloudflare API token from environment
if not self.cloudflare_api_token:
return self.fail_response("CLOUDFLARE_API_TOKEN environment variable not set")
# Single command that creates the project if it doesn't exist and then deploys
project_name = f"{self.sandbox_id}-{name}"
deploy_cmd = f'''cd {self.workspace_path} && export CLOUDFLARE_API_TOKEN={self.cloudflare_api_token} &&
(npx wrangler pages deploy {full_path} --project-name {project_name} ||
(npx wrangler pages project create {project_name} --production-branch production &&
deploy_cmd = f'''cd {self.workspace_path} && export CLOUDFLARE_API_TOKEN={self.cloudflare_api_token} &&
(npx wrangler pages deploy {full_path} --project-name {project_name} ||
(npx wrangler pages project create {project_name} --production-branch production &&
npx wrangler pages deploy {full_path} --project-name {project_name}))'''
# Execute the command directly using the sandbox's process.exec method
response = self.sandbox.process.exec(f"/bin/sh -c \"{deploy_cmd}\"",
timeout=300)
print(f"Deployment command output: {response.result}")
if response.exit_code == 0:
return self.success_response({
"message": f"Website deployed successfully",
@@ -127,21 +127,21 @@ class SandboxDeployTool(SandboxToolsBase):
if __name__ == "__main__":
import asyncio
import sys
async def test_deploy():
# Replace these with actual values for testing
sandbox_id = "sandbox-ccb30b35"
password = "test-password"
# Initialize the deploy tool
deploy_tool = SandboxDeployTool(sandbox_id, password)
# Test deployment - replace with actual directory path and site name
result = await deploy_tool.deploy(
name="test-site-1x",
directory_path="website" # Directory containing static site files
)
print(f"Deployment result: {result}")
asyncio.run(test_deploy())
app/tool/sb_expose_tool.py
+11 -11
View File
@@ -1,5 +1,5 @@
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from agentpress.thread_manager import ThreadManager
import asyncio
import time
@@ -13,25 +13,25 @@ class SandboxExposeTool(SandboxToolsBase):
async def _wait_for_sandbox_services(self, timeout: int = 30) -> bool:
"""Wait for sandbox services to be fully started before exposing ports."""
start_time = time.time()
while time.time() - start_time < timeout:
try:
# Check if supervisord is running and managing services
result = self.sandbox.process.exec("supervisorctl status", timeout=10)
if result.exit_code == 0:
# Check if key services are running
status_output = result.output
if "http_server" in status_output and "RUNNING" in status_output:
return True
# If services aren't ready, wait a bit
await asyncio.sleep(2)
except Exception as e:
# If we can't check status, wait a bit and try again
await asyncio.sleep(2)
return False
@openapi_schema({
@@ -85,10 +85,10 @@ class SandboxExposeTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# Convert port to integer if it's a string
port = int(port)
# Validate port number
if not 1 <= port <= 65535:
return self.fail_response(f"Invalid port number: {port}. Must be between 1 and 65535.")
@@ -110,16 +110,16 @@ class SandboxExposeTool(SandboxToolsBase):
# Get the preview link for the specified port
preview_link = self.sandbox.get_preview_link(port)
# Extract the actual URL from the preview link object
url = preview_link.url if hasattr(preview_link, 'url') else str(preview_link)
return self.success_response({
"url": url,
"port": port,
"message": f"Successfully exposed port {port} to the public. Users can now access this service at: {url}"
})
except ValueError:
return self.fail_response(f"Invalid port number: {port}. Must be a valid integer between 1 and 65535.")
except Exception as e:
app/tool/sb_files_tool.py
+35 -35
View File
@@ -1,5 +1,5 @@
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from utils.files_utils import should_exclude_file, clean_path
from agentpress.thread_manager import ThreadManager
from utils.logger import logger
@@ -35,11 +35,11 @@ class SandboxFilesTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
files = self.sandbox.fs.list_files(self.workspace_path)
for file_info in files:
rel_path = file_info.name
# Skip excluded files and directories
if self._should_exclude_file(rel_path) or file_info.is_dir:
continue
@@ -59,7 +59,7 @@ class SandboxFilesTool(SandboxToolsBase):
print(f"Skipping binary file: {rel_path}")
return files_state
except Exception as e:
print(f"Error getting workspace state: {str(e)}")
return {}
@@ -111,7 +111,7 @@ class SandboxFilesTool(SandboxToolsBase):
# This is the file content
def main():
print("Hello, World!")
if __name__ == "__main__":
main()
</parameter>
@@ -123,23 +123,23 @@ class SandboxFilesTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
file_path = self.clean_path(file_path)
full_path = f"{self.workspace_path}/{file_path}"
if self._file_exists(full_path):
return self.fail_response(f"File '{file_path}' already exists. Use update_file to modify existing files.")
# Create parent directories if needed
parent_dir = '/'.join(full_path.split('/')[:-1])
if parent_dir:
self.sandbox.fs.create_folder(parent_dir, "755")
# Write the file content
self.sandbox.fs.upload_file(file_contents.encode(), full_path)
self.sandbox.fs.set_file_permissions(full_path, permissions)
message = f"File '{file_path}' created successfully."
# Check if index.html was created and add 8080 server info (only in root workspace)
if file_path.lower() == 'index.html':
try:
@@ -149,7 +149,7 @@ class SandboxFilesTool(SandboxToolsBase):
message += "\n[Note: Use the provided HTTP server URL above instead of starting a new server]"
except Exception as e:
logger.warning(f"Failed to get website URL for index.html: {str(e)}")
return self.success_response(message)
except Exception as e:
return self.fail_response(f"Error creating file: {str(e)}")
@@ -200,41 +200,41 @@ class SandboxFilesTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
file_path = self.clean_path(file_path)
full_path = f"{self.workspace_path}/{file_path}"
if not self._file_exists(full_path):
return self.fail_response(f"File '{file_path}' does not exist")
content = self.sandbox.fs.download_file(full_path).decode()
old_str = old_str.expandtabs()
new_str = new_str.expandtabs()
occurrences = content.count(old_str)
if occurrences == 0:
return self.fail_response(f"String '{old_str}' not found in file")
if occurrences > 1:
lines = [i+1 for i, line in enumerate(content.split('\n')) if old_str in line]
return self.fail_response(f"Multiple occurrences found in lines {lines}. Please ensure string is unique")
# Perform replacement
new_content = content.replace(old_str, new_str)
self.sandbox.fs.upload_file(new_content.encode(), full_path)
# Show snippet around the edit
replacement_line = content.split(old_str)[0].count('\n')
start_line = max(0, replacement_line - self.SNIPPET_LINES)
end_line = replacement_line + self.SNIPPET_LINES + new_str.count('\n')
snippet = '\n'.join(new_content.split('\n')[start_line:end_line + 1])
# Get preview URL if it's an HTML file
# preview_url = self._get_preview_url(file_path)
message = f"Replacement successful."
# if preview_url:
# message += f"\n\nYou can preview this HTML file at: {preview_url}"
return self.success_response(message)
except Exception as e:
return self.fail_response(f"Error replacing string: {str(e)}")
@@ -288,17 +288,17 @@ class SandboxFilesTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
file_path = self.clean_path(file_path)
full_path = f"{self.workspace_path}/{file_path}"
if not self._file_exists(full_path):
return self.fail_response(f"File '{file_path}' does not exist. Use create_file to create a new file.")
self.sandbox.fs.upload_file(file_contents.encode(), full_path)
self.sandbox.fs.set_file_permissions(full_path, permissions)
message = f"File '{file_path}' completely rewritten successfully."
# Check if index.html was rewritten and add 8080 server info (only in root workspace)
if file_path.lower() == 'index.html':
try:
@@ -308,7 +308,7 @@ class SandboxFilesTool(SandboxToolsBase):
message += "\n[Note: Use the provided HTTP server URL above instead of starting a new server]"
except Exception as e:
logger.warning(f"Failed to get website URL for index.html: {str(e)}")
return self.success_response(message)
except Exception as e:
return self.fail_response(f"Error rewriting file: {str(e)}")
@@ -347,12 +347,12 @@ class SandboxFilesTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
file_path = self.clean_path(file_path)
full_path = f"{self.workspace_path}/{file_path}"
if not self._file_exists(full_path):
return self.fail_response(f"File '{file_path}' does not exist")
self.sandbox.fs.delete_file(full_path)
return self.success_response(f"File '{file_path}' deleted successfully.")
except Exception as e:
@@ -412,12 +412,12 @@ class SandboxFilesTool(SandboxToolsBase):
# )
# async def read_file(self, file_path: str, start_line: int = 1, end_line: Optional[int] = None) -> ToolResult:
# """Read file content with optional line range specification.
# Args:
# file_path: Path to the file relative to /workspace
# start_line: Starting line number (1-based), defaults to 1
# end_line: Ending line number (inclusive), defaults to None (end of file)
# Returns:
# ToolResult containing:
# - Success: File content and metadata
@@ -426,27 +426,27 @@ class SandboxFilesTool(SandboxToolsBase):
# try:
# file_path = self.clean_path(file_path)
# full_path = f"{self.workspace_path}/{file_path}"
# if not self._file_exists(full_path):
# return self.fail_response(f"File '{file_path}' does not exist")
# # Download and decode file content
# content = self.sandbox.fs.download_file(full_path).decode()
# # Split content into lines
# lines = content.split('\n')
# total_lines = len(lines)
# # Handle line range if specified
# if start_line > 1 or end_line is not None:
# # Convert to 0-based indices
# start_idx = max(0, start_line - 1)
# end_idx = end_line if end_line is not None else total_lines
# end_idx = min(end_idx, total_lines) # Ensure we don't exceed file length
# # Extract the requested lines
# content = '\n'.join(lines[start_idx:end_idx])
# return self.success_response({
# "content": content,
# "file_path": file_path,
@@ -454,7 +454,7 @@ class SandboxFilesTool(SandboxToolsBase):
# "end_line": end_line if end_line is not None else total_lines,
# "total_lines": total_lines
# })
# except UnicodeDecodeError:
# return self.fail_response(f"File '{file_path}' appears to be binary and cannot be read as text")
# except Exception as e:
app/tool/sb_shell_tool.py
+39 -39
View File
@@ -2,11 +2,11 @@ from typing import Optional, Dict, Any
import time
from uuid import uuid4
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from agentpress.thread_manager import ThreadManager
class SandboxShellTool(SandboxToolsBase):
"""Tool for executing tasks in a Daytona sandbox with browser-use capabilities.
"""Tool for executing tasks in a Daytona sandbox with browser-use capabilities.
Uses sessions for maintaining state between commands and provides comprehensive process management."""
def __init__(self, project_id: str, thread_manager: ThreadManager):
@@ -108,8 +108,8 @@ class SandboxShellTool(SandboxToolsBase):
'''
)
async def execute_command(
self,
command: str,
self,
command: str,
folder: Optional[str] = None,
session_name: Optional[str] = None,
blocking: bool = False,
@@ -118,61 +118,61 @@ class SandboxShellTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# Set up working directory
cwd = self.workspace_path
if folder:
folder = folder.strip('/')
cwd = f"{self.workspace_path}/{folder}"
# Generate a session name if not provided
if not session_name:
session_name = f"session_{str(uuid4())[:8]}"
# Check if tmux session already exists
check_session = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'")
session_exists = "not_exists" not in check_session.get("output", "")
if not session_exists:
# Create a new tmux session
await self._execute_raw_command(f"tmux new-session -d -s {session_name}")
# Ensure we're in the correct directory and send command to tmux
full_command = f"cd {cwd} && {command}"
wrapped_command = full_command.replace('"', '\\"') # Escape double quotes
# Send command to tmux session
await self._execute_raw_command(f'tmux send-keys -t {session_name} "{wrapped_command}" Enter')
if blocking:
# For blocking execution, wait and capture output
start_time = time.time()
while (time.time() - start_time) < timeout:
# Wait a bit before checking
time.sleep(2)
# Check if session still exists (command might have exited)
check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'ended'")
if "ended" in check_result.get("output", ""):
break
# Get current output and check for common completion indicators
output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -")
current_output = output_result.get("output", "")
# Check for prompt indicators that suggest command completion
last_lines = current_output.split('\n')[-3:]
completion_indicators = ['$', '#', '>', 'Done', 'Completed', 'Finished', '✓']
if any(indicator in line for indicator in completion_indicators for line in last_lines):
break
# Capture final output
output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -")
final_output = output_result.get("output", "")
# Kill the session after capture
await self._execute_raw_command(f"tmux kill-session -t {session_name}")
return self.success_response({
"output": final_output,
"session_name": session_name,
@@ -187,7 +187,7 @@ class SandboxShellTool(SandboxToolsBase):
"message": f"Command sent to tmux session '{session_name}'. Use check_command_output to view results.",
"completed": False
})
except Exception as e:
# Attempt to clean up session in case of error
if session_name:
@@ -201,7 +201,7 @@ class SandboxShellTool(SandboxToolsBase):
"""Execute a raw command directly in the sandbox."""
# Ensure session exists for raw commands
session_id = await self._ensure_session("raw_commands")
# Execute command in session
from sandbox.sandbox import SessionExecuteRequest
req = SessionExecuteRequest(
@@ -209,18 +209,18 @@ class SandboxShellTool(SandboxToolsBase):
var_async=False,
cwd=self.workspace_path
)
response = self.sandbox.process.execute_session_command(
session_id=session_id,
req=req,
timeout=30 # Short timeout for utility commands
)
logs = self.sandbox.process.get_session_command_logs(
session_id=session_id,
command_id=response.cmd_id
)
return {
"output": logs,
"exit_code": response.exit_code
@@ -260,7 +260,7 @@ class SandboxShellTool(SandboxToolsBase):
<parameter name="session_name">dev_server</parameter>
</invoke>
</function_calls>
<!-- Example 2: Check final output and kill session -->
<function_calls>
<invoke name="check_command_output">
@@ -278,29 +278,29 @@ class SandboxShellTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# Check if session exists
check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'")
if "not_exists" in check_result.get("output", ""):
return self.fail_response(f"Tmux session '{session_name}' does not exist.")
# Get output from tmux pane
output_result = await self._execute_raw_command(f"tmux capture-pane -t {session_name} -p -S - -E -")
output = output_result.get("output", "")
# Kill session if requested
if kill_session:
await self._execute_raw_command(f"tmux kill-session -t {session_name}")
termination_status = "Session terminated."
else:
termination_status = "Session still running."
return self.success_response({
"output": output,
"session_name": session_name,
"status": termination_status
})
except Exception as e:
return self.fail_response(f"Error checking command output: {str(e)}")
@@ -341,19 +341,19 @@ class SandboxShellTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# Check if session exists
check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'")
if "not_exists" in check_result.get("output", ""):
return self.fail_response(f"Tmux session '{session_name}' does not exist.")
# Kill the session
await self._execute_raw_command(f"tmux kill-session -t {session_name}")
return self.success_response({
"message": f"Tmux session '{session_name}' terminated successfully."
})
except Exception as e:
return self.fail_response(f"Error terminating command: {str(e)}")
@@ -382,17 +382,17 @@ class SandboxShellTool(SandboxToolsBase):
try:
# Ensure sandbox is initialized
await self._ensure_sandbox()
# List all tmux sessions
result = await self._execute_raw_command("tmux list-sessions 2>/dev/null || echo 'No sessions'")
output = result.get("output", "")
if "No sessions" in output or not output.strip():
return self.success_response({
"message": "No active tmux sessions found.",
"sessions": []
})
# Parse session list
sessions = []
for line in output.split('\n'):
@@ -401,12 +401,12 @@ class SandboxShellTool(SandboxToolsBase):
if parts:
session_name = parts[0].strip()
sessions.append(session_name)
return self.success_response({
"message": f"Found {len(sessions)} active sessions.",
"sessions": sessions
})
except Exception as e:
return self.fail_response(f"Error listing commands: {str(e)}")
@@ -414,10 +414,10 @@ class SandboxShellTool(SandboxToolsBase):
"""Clean up all sessions."""
for session_name in list(self._sessions.keys()):
await self._cleanup_session(session_name)
# Also clean up any tmux sessions
try:
await self._ensure_sandbox()
await self._execute_raw_command("tmux kill-server 2>/dev/null || true")
except:
pass
pass
app/tool/sb_vision_tool.py
+13 -13
View File
@@ -6,7 +6,7 @@ from io import BytesIO
from PIL import Image
from agentpress.tool import ToolResult, openapi_schema, xml_schema
from sandbox.tool_base import SandboxToolsBase
from daytona.tool_base import SandboxToolsBase
from agentpress.thread_manager import ThreadManager
import json
@@ -38,19 +38,19 @@ class SandboxVisionTool(SandboxToolsBase):
def compress_image(self, image_bytes: bytes, mime_type: str, file_path: str) -> Tuple[bytes, str]:
"""Compress an image to reduce its size while maintaining reasonable quality.
Args:
image_bytes: Original image bytes
mime_type: MIME type of the image
file_path: Path to the image file (for logging)
Returns:
Tuple of (compressed_bytes, new_mime_type)
"""
try:
# Open image from bytes
img = Image.open(BytesIO(image_bytes))
# Convert RGBA to RGB if necessary (for JPEG)
if img.mode in ('RGBA', 'LA', 'P'):
# Create a white background
@@ -59,7 +59,7 @@ class SandboxVisionTool(SandboxToolsBase):
img = img.convert('RGBA')
background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
img = background
# Calculate new dimensions while maintaining aspect ratio
width, height = img.size
if width > DEFAULT_MAX_WIDTH or height > DEFAULT_MAX_HEIGHT:
@@ -68,10 +68,10 @@ class SandboxVisionTool(SandboxToolsBase):
new_height = int(height * ratio)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
print(f"[SeeImage] Resized image from {width}x{height} to {new_width}x{new_height}")
# Save to bytes with compression
output = BytesIO()
# Determine output format based on original mime type
if mime_type == 'image/gif':
# Keep GIFs as GIFs to preserve animation
@@ -85,17 +85,17 @@ class SandboxVisionTool(SandboxToolsBase):
# Convert everything else to JPEG for better compression
img.save(output, format='JPEG', quality=DEFAULT_JPEG_QUALITY, optimize=True)
output_mime = 'image/jpeg'
compressed_bytes = output.getvalue()
# Log compression results
original_size = len(image_bytes)
compressed_size = len(compressed_bytes)
compression_ratio = (1 - compressed_size / original_size) * 100
print(f"[SeeImage] Compressed '{file_path}' from {original_size / 1024:.1f}KB to {compressed_size / 1024:.1f}KB ({compression_ratio:.1f}% reduction)")
return compressed_bytes, output_mime
except Exception as e:
print(f"[SeeImage] Failed to compress image: {str(e)}. Using original.")
return image_bytes, mime_type
@@ -173,7 +173,7 @@ class SandboxVisionTool(SandboxToolsBase):
# Compress the image
compressed_bytes, compressed_mime_type = self.compress_image(image_bytes, mime_type, cleaned_path)
# Check if compressed image is still too large
if len(compressed_bytes) > MAX_COMPRESSED_SIZE:
return self.fail_response(f"Image file '{cleaned_path}' is still too large after compression ({len(compressed_bytes) / (1024*1024):.2f}MB). Maximum compressed size is {MAX_COMPRESSED_SIZE / (1024*1024)}MB.")
@@ -203,4 +203,4 @@ class SandboxVisionTool(SandboxToolsBase):
return self.success_response(f"Successfully loaded and compressed the image '{cleaned_path}' (reduced from {file_info.size / 1024:.1f}KB to {len(compressed_bytes) / 1024:.1f}KB).")
except Exception as e:
return self.fail_response(f"An unexpected error occurred while trying to see the image: {str(e)}")
return self.fail_response(f"An unexpected error occurred while trying to see the image: {str(e)}")
app/tool/update_agent_tool.py
-889
View File
@@ -1,889 +0,0 @@
import json
import httpx
from typing import Optional, Dict, Any, List
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
from agentpress.thread_manager import ThreadManager
class UpdateAgentTool(Tool):
"""Tool for updating agent configuration.
This tool is used by the agent builder to update agent properties
based on user requirements.
"""
def __init__(self, thread_manager: ThreadManager, db_connection, agent_id: str):
super().__init__()
self.thread_manager = thread_manager
self.db = db_connection
self.agent_id = agent_id
# Smithery API configuration
self.smithery_api_base_url = "https://registry.smithery.ai"
import os
self.smithery_api_key = os.getenv("SMITHERY_API_KEY")
@openapi_schema({
"type": "function",
"function": {
"name": "update_agent",
"description": "Update the agent's configuration including name, description, system prompt, tools, and MCP servers. Call this whenever the user wants to modify any aspect of the agent.",
"parameters": {
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "The name of the agent. Should be descriptive and indicate the agent's purpose."
},
"description": {
"type": "string",
"description": "A brief description of what the agent does and its capabilities."
},
"system_prompt": {
"type": "string",
"description": "The system instructions that define the agent's behavior, expertise, and approach. This should be comprehensive and well-structured."
},
"agentpress_tools": {
"type": "object",
"description": "Configuration for AgentPress tools. Each key is a tool name, and the value is an object with 'enabled' (boolean) and 'description' (string) properties.",
"additionalProperties": {
"type": "object",
"properties": {
"enabled": {"type": "boolean"},
"description": {"type": "string"}
}
}
},
"configured_mcps": {
"type": "array",
"description": "List of configured MCP servers for external integrations.",
"items": {
"type": "object",
"properties": {
"name": {"type": "string"},
"qualifiedName": {"type": "string"},
"config": {"type": "object"},
"enabledTools": {
"type": "array",
"items": {"type": "string"}
}
}
}
},
"avatar": {
"type": "string",
"description": "Emoji to use as the agent's avatar."
},
"avatar_color": {
"type": "string",
"description": "Hex color code for the agent's avatar background."
}
},
"required": []
}
}
})
@xml_schema(
tag_name="update-agent",
mappings=[
{"param_name": "name", "node_type": "attribute", "path": ".", "required": False},
{"param_name": "description", "node_type": "element", "path": "description", "required": False},
{"param_name": "system_prompt", "node_type": "element", "path": "system_prompt", "required": False},
{"param_name": "agentpress_tools", "node_type": "element", "path": "agentpress_tools", "required": False},
{"param_name": "configured_mcps", "node_type": "element", "path": "configured_mcps", "required": False},
{"param_name": "avatar", "node_type": "attribute", "path": ".", "required": False},
{"param_name": "avatar_color", "node_type": "attribute", "path": ".", "required": False}
],
example='''
<function_calls>
<invoke name="update_agent">
<parameter name="name">Research Assistant</parameter>
<parameter name="description">An AI assistant specialized in conducting research and providing comprehensive analysis</parameter>
<parameter name="system_prompt">You are a research assistant with expertise in gathering, analyzing, and synthesizing information. Your approach is thorough and methodical...</parameter>
<parameter name="agentpress_tools">{"web_search": {"enabled": true, "description": "Search the web for information"}, "sb_files": {"enabled": true, "description": "Read and write files"}}</parameter>
<parameter name="avatar">🔬</parameter>
<parameter name="avatar_color">#4F46E5</parameter>
</invoke>
</function_calls>
'''
)
async def update_agent(
self,
name: Optional[str] = None,
description: Optional[str] = None,
system_prompt: Optional[str] = None,
agentpress_tools: Optional[Dict[str, Dict[str, Any]]] = None,
configured_mcps: Optional[list] = None,
avatar: Optional[str] = None,
avatar_color: Optional[str] = None
) -> ToolResult:
"""Update agent configuration with provided fields.
Args:
name: Agent name
description: Agent description
system_prompt: System instructions for the agent
agentpress_tools: AgentPress tools configuration
configured_mcps: MCP servers configuration
avatar: Emoji avatar
avatar_color: Avatar background color
Returns:
ToolResult with updated agent data or error
"""
try:
client = await self.db.client
update_data = {}
if name is not None:
update_data["name"] = name
if description is not None:
update_data["description"] = description
if system_prompt is not None:
update_data["system_prompt"] = system_prompt
if agentpress_tools is not None:
formatted_tools = {}
for tool_name, tool_config in agentpress_tools.items():
if isinstance(tool_config, dict):
formatted_tools[tool_name] = {
"enabled": tool_config.get("enabled", False),
"description": tool_config.get("description", "")
}
update_data["agentpress_tools"] = formatted_tools
if configured_mcps is not None:
if isinstance(configured_mcps, str):
configured_mcps = json.loads(configured_mcps)
update_data["configured_mcps"] = configured_mcps
if avatar is not None:
update_data["avatar"] = avatar
if avatar_color is not None:
update_data["avatar_color"] = avatar_color
if not update_data:
return self.fail_response("No fields provided to update")
result = await client.table('agents').update(update_data).eq('agent_id', self.agent_id).execute()
if not result.data:
return self.fail_response("Failed to update agent")
return self.success_response({
"message": "Agent updated successfully",
"updated_fields": list(update_data.keys()),
"agent": result.data[0]
})
except Exception as e:
return self.fail_response(f"Error updating agent: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "get_current_agent_config",
"description": "Get the current configuration of the agent being edited. Use this to check what's already configured before making updates.",
"parameters": {
"type": "object",
"properties": {},
"required": []
}
}
})
@xml_schema(
tag_name="get-current-agent-config",
mappings=[],
example='''
<function_calls>
<invoke name="get_current_agent_config">
</invoke>
</function_calls>
'''
)
async def get_current_agent_config(self) -> ToolResult:
"""Get the current agent configuration.
Returns:
ToolResult with current agent configuration
"""
try:
client = await self.db.client
result = await client.table('agents').select('*').eq('agent_id', self.agent_id).execute()
if not result.data:
return self.fail_response("Agent not found")
agent = result.data[0]
config_summary = {
"agent_id": agent["agent_id"],
"name": agent.get("name", "Untitled Agent"),
"description": agent.get("description", "No description set"),
"system_prompt": agent.get("system_prompt", "No system prompt set"),
"avatar": agent.get("avatar", "🤖"),
"avatar_color": agent.get("avatar_color", "#6B7280"),
"agentpress_tools": agent.get("agentpress_tools", {}),
"configured_mcps": agent.get("configured_mcps", []),
"created_at": agent.get("created_at"),
"updated_at": agent.get("updated_at")
}
tools_count = len([t for t, cfg in config_summary["agentpress_tools"].items() if cfg.get("enabled")])
mcps_count = len(config_summary["configured_mcps"])
return self.success_response({
"summary": f"Agent '{config_summary['name']}' has {tools_count} tools enabled and {mcps_count} MCP servers configured.",
"configuration": config_summary
})
except Exception as e:
return self.fail_response(f"Error getting agent configuration: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "search_mcp_servers",
"description": "Search for MCP servers from the Smithery registry based on user requirements. Use this when the user wants to add MCP tools to their agent.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query for finding relevant MCP servers (e.g., 'linear', 'github', 'database', 'search')"
},
"category": {
"type": "string",
"description": "Optional category filter",
"enum": ["AI & Search", "Development & Version Control", "Project Management", "Communication & Collaboration", "Data & Analytics", "Cloud & Infrastructure", "File Storage", "Marketing & Sales", "Customer Support", "Finance", "Automation & Productivity", "Utilities"]
},
"limit": {
"type": "integer",
"description": "Maximum number of servers to return (default: 10)",
"default": 10
}
},
"required": ["query"]
}
}
})
@xml_schema(
tag_name="search-mcp-servers",
mappings=[
{"param_name": "query", "node_type": "attribute", "path": "."},
{"param_name": "category", "node_type": "attribute", "path": "."},
{"param_name": "limit", "node_type": "attribute", "path": "."}
],
example='''
<function_calls>
<invoke name="search_mcp_servers">
<parameter name="query">linear</parameter>
<parameter name="limit">5</parameter>
</invoke>
</function_calls>
'''
)
async def search_mcp_servers(
self,
query: str,
category: Optional[str] = None,
limit: int = 10
) -> ToolResult:
"""Search for MCP servers based on user requirements.
Args:
query: Search query for finding relevant MCP servers
category: Optional category filter
limit: Maximum number of servers to return
Returns:
ToolResult with matching MCP servers
"""
try:
async with httpx.AsyncClient() as client:
headers = {
"Accept": "application/json",
"User-Agent": "Suna-MCP-Integration/1.0"
}
if self.smithery_api_key:
headers["Authorization"] = f"Bearer {self.smithery_api_key}"
params = {
"q": query,
"page": 1,
"pageSize": min(limit * 2, 50) # Get more results to filter
}
response = await client.get(
f"{self.smithery_api_base_url}/servers",
headers=headers,
params=params,
timeout=30.0
)
response.raise_for_status()
data = response.json()
servers = data.get("servers", [])
# Filter by category if specified
if category:
filtered_servers = []
for server in servers:
server_category = self._categorize_server(server)
if server_category == category:
filtered_servers.append(server)
servers = filtered_servers
# Sort by useCount and limit results
servers = sorted(servers, key=lambda x: x.get("useCount", 0), reverse=True)[:limit]
# Format results for user-friendly display
formatted_servers = []
for server in servers:
formatted_servers.append({
"name": server.get("displayName", server.get("qualifiedName", "Unknown")),
"qualifiedName": server.get("qualifiedName"),
"description": server.get("description", "No description available"),
"useCount": server.get("useCount", 0),
"category": self._categorize_server(server),
"homepage": server.get("homepage", ""),
"isDeployed": server.get("isDeployed", False)
})
if not formatted_servers:
return ToolResult(
success=False,
output=json.dumps([], ensure_ascii=False)
)
return ToolResult(
success=True,
output=json.dumps(formatted_servers, ensure_ascii=False)
)
except Exception as e:
return self.fail_response(f"Error searching MCP servers: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "get_mcp_server_tools",
"description": "Get detailed information about a specific MCP server including its available tools. Use this after the user selects a server they want to connect to.",
"parameters": {
"type": "object",
"properties": {
"qualified_name": {
"type": "string",
"description": "The qualified name of the MCP server (e.g., 'exa', '@smithery-ai/github')"
}
},
"required": ["qualified_name"]
}
}
})
@xml_schema(
tag_name="get-mcp-server-tools",
mappings=[
{"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True}
],
example='''
<function_calls>
<invoke name="get_mcp_server_tools">
<parameter name="qualified_name">exa</parameter>
</invoke>
</function_calls>
'''
)
async def get_mcp_server_tools(self, qualified_name: str) -> ToolResult:
"""Get detailed information about a specific MCP server and its tools.
Args:
qualified_name: The qualified name of the MCP server
Returns:
ToolResult with server details and available tools
"""
try:
# First get server metadata from registry
async with httpx.AsyncClient() as client:
headers = {
"Accept": "application/json",
"User-Agent": "Suna-MCP-Integration/1.0"
}
if self.smithery_api_key:
headers["Authorization"] = f"Bearer {self.smithery_api_key}"
# URL encode the qualified name if it contains special characters
from urllib.parse import quote
if '@' in qualified_name or '/' in qualified_name:
encoded_name = quote(qualified_name, safe='')
else:
encoded_name = qualified_name
url = f"{self.smithery_api_base_url}/servers/{encoded_name}"
response = await client.get(
url,
headers=headers,
timeout=30.0
)
response.raise_for_status()
server_data = response.json()
# Now connect to the MCP server to get actual tools using ClientSession
try:
# Import MCP components
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
import base64
import os
# Check if Smithery API key is available
smithery_api_key = os.getenv("SMITHERY_API_KEY")
if not smithery_api_key:
raise ValueError("SMITHERY_API_KEY environment variable is not set")
# Create server URL with empty config for testing
config_json = json.dumps({})
config_b64 = base64.b64encode(config_json.encode()).decode()
server_url = f"https://server.smithery.ai/{qualified_name}/mcp?config={config_b64}&api_key={smithery_api_key}"
# Connect and get tools
async with streamablehttp_client(server_url) as (read_stream, write_stream, _):
async with ClientSession(read_stream, write_stream) as session:
# Initialize the connection
await session.initialize()
# List available tools
tools_result = await session.list_tools()
tools = tools_result.tools if hasattr(tools_result, 'tools') else tools_result
# Format tools for user-friendly display
formatted_tools = []
for tool in tools:
tool_info = {
"name": tool.name,
"description": getattr(tool, 'description', 'No description available'),
}
# Extract parameters from inputSchema if available
if hasattr(tool, 'inputSchema') and tool.inputSchema:
schema = tool.inputSchema
if isinstance(schema, dict):
tool_info["parameters"] = schema.get("properties", {})
tool_info["required_params"] = schema.get("required", [])
else:
tool_info["parameters"] = {}
tool_info["required_params"] = []
else:
tool_info["parameters"] = {}
tool_info["required_params"] = []
formatted_tools.append(tool_info)
# Extract configuration requirements from server metadata
config_requirements = []
security = server_data.get("security", {})
if security:
for key, value in security.items():
if isinstance(value, dict):
config_requirements.append({
"name": key,
"description": value.get("description", f"Configuration for {key}"),
"required": value.get("required", False),
"type": value.get("type", "string")
})
server_info = {
"name": server_data.get("displayName", qualified_name),
"qualifiedName": qualified_name,
"description": server_data.get("description", "No description available"),
"homepage": server_data.get("homepage", ""),
"iconUrl": server_data.get("iconUrl", ""),
"isDeployed": server_data.get("isDeployed", False),
"tools": formatted_tools,
"config_requirements": config_requirements,
"total_tools": len(formatted_tools)
}
return self.success_response({
"message": f"Found {len(formatted_tools)} tools for {server_info['name']}",
"server": server_info
})
except Exception as mcp_error:
# If MCP connection fails, fall back to registry data
tools = server_data.get("tools", [])
formatted_tools = []
for tool in tools:
formatted_tools.append({
"name": tool.get("name", "Unknown"),
"description": tool.get("description", "No description available"),
"parameters": tool.get("inputSchema", {}).get("properties", {}),
"required_params": tool.get("inputSchema", {}).get("required", [])
})
config_requirements = []
security = server_data.get("security", {})
if security:
for key, value in security.items():
if isinstance(value, dict):
config_requirements.append({
"name": key,
"description": value.get("description", f"Configuration for {key}"),
"required": value.get("required", False),
"type": value.get("type", "string")
})
server_info = {
"name": server_data.get("displayName", qualified_name),
"qualifiedName": qualified_name,
"description": server_data.get("description", "No description available"),
"homepage": server_data.get("homepage", ""),
"iconUrl": server_data.get("iconUrl", ""),
"isDeployed": server_data.get("isDeployed", False),
"tools": formatted_tools,
"config_requirements": config_requirements,
"total_tools": len(formatted_tools),
"note": "Tools listed from registry metadata (MCP connection failed - may need configuration)"
}
return self.success_response({
"message": f"Found {len(formatted_tools)} tools for {server_info['name']} (from registry)",
"server": server_info
})
except Exception as e:
return self.fail_response(f"Error getting MCP server tools: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "configure_mcp_server",
"description": "Configure and add an MCP server to the agent with selected tools. Use this after the user has chosen which tools they want from a server.",
"parameters": {
"type": "object",
"properties": {
"qualified_name": {
"type": "string",
"description": "The qualified name of the MCP server"
},
"display_name": {
"type": "string",
"description": "Display name for the server"
},
"enabled_tools": {
"type": "array",
"description": "List of tool names to enable for this server",
"items": {"type": "string"}
},
"config": {
"type": "object",
"description": "Configuration object with API keys and other settings",
"additionalProperties": True
}
},
"required": ["qualified_name", "display_name", "enabled_tools"]
}
}
})
@xml_schema(
tag_name="configure-mcp-server",
mappings=[
{"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True},
{"param_name": "display_name", "node_type": "attribute", "path": ".", "required": True},
{"param_name": "enabled_tools", "node_type": "element", "path": "enabled_tools", "required": True},
{"param_name": "config", "node_type": "element", "path": "config", "required": False}
],
example='''
<function_calls>
<invoke name="configure_mcp_server">
<parameter name="qualified_name">exa</parameter>
<parameter name="display_name">Exa Search</parameter>
<parameter name="enabled_tools">["search", "find_similar"]</parameter>
<parameter name="config">{"exaApiKey": "user-api-key"}</parameter>
</invoke>
</function_calls>
'''
)
async def configure_mcp_server(
self,
qualified_name: str,
display_name: str,
enabled_tools: List[str],
config: Optional[Dict[str, Any]] = None
) -> ToolResult:
"""Configure and add an MCP server to the agent.
Args:
qualified_name: The qualified name of the MCP server
display_name: Display name for the server
enabled_tools: List of tool names to enable
config: Configuration object with API keys and settings
Returns:
ToolResult with configuration status
"""
try:
client = await self.db.client
# Get current agent configuration
result = await client.table('agents').select('configured_mcps').eq('agent_id', self.agent_id).execute()
if not result.data:
return self.fail_response("Agent not found")
current_mcps = result.data[0].get('configured_mcps', [])
# Check if server is already configured
existing_server_index = None
for i, mcp in enumerate(current_mcps):
if mcp.get('qualifiedName') == qualified_name:
existing_server_index = i
break
# Create new MCP configuration
new_mcp_config = {
"name": display_name,
"qualifiedName": qualified_name,
"config": config or {},
"enabledTools": enabled_tools
}
# Update or add the configuration
if existing_server_index is not None:
current_mcps[existing_server_index] = new_mcp_config
action = "updated"
else:
current_mcps.append(new_mcp_config)
action = "added"
# Save to database
update_result = await client.table('agents').update({
'configured_mcps': current_mcps
}).eq('agent_id', self.agent_id).execute()
if not update_result.data:
return self.fail_response("Failed to save MCP configuration")
return self.success_response({
"message": f"Successfully {action} MCP server '{display_name}' with {len(enabled_tools)} tools",
"server": new_mcp_config,
"total_mcp_servers": len(current_mcps),
"action": action
})
except Exception as e:
return self.fail_response(f"Error configuring MCP server: {str(e)}")
@openapi_schema({
"type": "function",
"function": {
"name": "get_popular_mcp_servers",
"description": "Get a list of popular and recommended MCP servers organized by category. Use this to show users popular options when they want to add MCP tools.",
"parameters": {
"type": "object",
"properties": {
"category": {
"type": "string",
"description": "Optional category filter to show only servers from a specific category",
"enum": ["AI & Search", "Development & Version Control", "Project Management", "Communication & Collaboration", "Data & Analytics", "Cloud & Infrastructure", "File Storage", "Marketing & Sales", "Customer Support", "Finance", "Automation & Productivity", "Utilities"]
}
},
"required": []
}
}
})
@xml_schema(
tag_name="get-popular-mcp-servers",
mappings=[
{"param_name": "category", "node_type": "attribute", "path": ".", "required": False}
],
example='''
<function_calls>
<invoke name="get_popular_mcp_servers">
<parameter name="category">AI & Search</parameter>
</invoke>
</function_calls>
'''
)
async def get_popular_mcp_servers(self, category: Optional[str] = None) -> ToolResult:
"""Get popular MCP servers organized by category.
Args:
category: Optional category filter
Returns:
ToolResult with popular MCP servers
"""
try:
async with httpx.AsyncClient() as client:
headers = {
"Accept": "application/json",
"User-Agent": "Suna-MCP-Integration/1.0"
}
if self.smithery_api_key:
headers["Authorization"] = f"Bearer {self.smithery_api_key}"
response = await client.get(
f"{self.smithery_api_base_url}/servers",
headers=headers,
params={"page": 1, "pageSize": 50},
timeout=30.0
)
response.raise_for_status()
data = response.json()
servers = data.get("servers", [])
# Categorize servers
categorized = {}
for server in servers:
server_category = self._categorize_server(server)
if category and server_category != category:
continue
if server_category not in categorized:
categorized[server_category] = []
categorized[server_category].append({
"name": server.get("displayName", server.get("qualifiedName", "Unknown")),
"qualifiedName": server.get("qualifiedName"),
"description": server.get("description", "No description available"),
"useCount": server.get("useCount", 0),
"homepage": server.get("homepage", ""),
"isDeployed": server.get("isDeployed", False)
})
# Sort categories and servers within each category
for cat in categorized:
categorized[cat] = sorted(categorized[cat], key=lambda x: x["useCount"], reverse=True)[:5]
return self.success_response({
"message": f"Found popular MCP servers" + (f" in category '{category}'" if category else ""),
"categorized_servers": categorized,
"total_categories": len(categorized)
})
except Exception as e:
return self.fail_response(f"Error getting popular MCP servers: {str(e)}")
def _categorize_server(self, server: Dict[str, Any]) -> str:
"""Categorize a server based on its qualified name and description."""
qualified_name = server.get("qualifiedName", "").lower()
description = server.get("description", "").lower()
# Category mappings
category_mappings = {
"AI & Search": ["exa", "perplexity", "openai", "anthropic", "duckduckgo", "brave", "google", "search"],
"Development & Version Control": ["github", "gitlab", "bitbucket", "git"],
"Project Management": ["linear", "jira", "asana", "notion", "trello", "monday", "clickup"],
"Communication & Collaboration": ["slack", "discord", "teams", "zoom", "telegram"],
"Data & Analytics": ["postgres", "mysql", "mongodb", "bigquery", "snowflake", "sqlite", "redis", "database"],
"Cloud & Infrastructure": ["aws", "gcp", "azure", "vercel", "netlify", "cloudflare", "docker"],
"File Storage": ["gdrive", "google-drive", "dropbox", "box", "onedrive", "s3", "drive"],
"Marketing & Sales": ["hubspot", "salesforce", "mailchimp", "sendgrid"],
"Customer Support": ["zendesk", "intercom", "freshdesk", "helpscout"],
"Finance": ["stripe", "quickbooks", "xero", "plaid"],
"Automation & Productivity": ["playwright", "puppeteer", "selenium", "desktop-commander", "sequential-thinking", "automation"],
"Utilities": ["filesystem", "memory", "fetch", "time", "weather", "currency", "file"]
}
# Check qualified name and description for category keywords
for category, keywords in category_mappings.items():
for keyword in keywords:
if keyword in qualified_name or keyword in description:
return category
return "Other"
@openapi_schema({
"type": "function",
"function": {
"name": "test_mcp_server_connection",
"description": "Test connectivity to an MCP server with provided configuration. Use this to validate that a server can be connected to before adding it to the agent.",
"parameters": {
"type": "object",
"properties": {
"qualified_name": {
"type": "string",
"description": "The qualified name of the MCP server"
},
"config": {
"type": "object",
"description": "Configuration object with API keys and other settings",
"additionalProperties": True
}
},
"required": ["qualified_name"]
}
}
})
@xml_schema(
tag_name="test-mcp-server-connection",
mappings=[
{"param_name": "qualified_name", "node_type": "attribute", "path": ".", "required": True},
{"param_name": "config", "node_type": "element", "path": "config", "required": False}
],
example='''
<function_calls>
<invoke name="test_mcp_server_connection">
<parameter name="qualified_name">exa</parameter>
<parameter name="config">{"exaApiKey": "user-api-key"}</parameter>
</invoke>
</function_calls>
'''
)
async def test_mcp_server_connection(
self,
qualified_name: str,
config: Optional[Dict[str, Any]] = None
) -> ToolResult:
"""Test connectivity to an MCP server with provided configuration.
Args:
qualified_name: The qualified name of the MCP server
config: Configuration object with API keys and settings
Returns:
ToolResult with connection test results
"""
try:
# Import MCP components
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
import base64
import os
# Check if Smithery API key is available
smithery_api_key = os.getenv("SMITHERY_API_KEY")
if not smithery_api_key:
return self.fail_response("SMITHERY_API_KEY environment variable is not set")
# Create server URL with provided config
config_json = json.dumps(config or {})
config_b64 = base64.b64encode(config_json.encode()).decode()
server_url = f"https://server.smithery.ai/{qualified_name}/mcp?config={config_b64}&api_key={smithery_api_key}"
# Test connection
async with streamablehttp_client(server_url) as (read_stream, write_stream, _):
async with ClientSession(read_stream, write_stream) as session:
# Initialize the connection
await session.initialize()
# List available tools to verify connection
tools_result = await session.list_tools()
tools = tools_result.tools if hasattr(tools_result, 'tools') else tools_result
tool_names = [tool.name for tool in tools]
return self.success_response({
"message": f"Successfully connected to {qualified_name}",
"qualified_name": qualified_name,
"connection_status": "success",
"available_tools": tool_names,
"total_tools": len(tool_names)
})
except Exception as e:
return self.fail_response(f"Failed to connect to {qualified_name}: {str(e)}")
app/utils/files_utils.py
+10 -10
View File
@@ -40,10 +40,10 @@ EXCLUDED_EXT = {
def should_exclude_file(rel_path: str) -> bool:
"""Check if a file should be excluded based on path, name, or extension
Args:
rel_path: Relative path of the file to check
Returns:
True if the file should be excluded, False otherwise
"""
@@ -62,30 +62,30 @@ def should_exclude_file(rel_path: str) -> bool:
if ext.lower() in EXCLUDED_EXT:
return True
return False
return False
def clean_path(path: str, workspace_path: str = "/workspace") -> str:
"""Clean and normalize a path to be relative to the workspace
Args:
path: The path to clean
workspace_path: The base workspace path to remove (default: "/workspace")
Returns:
The cleaned path, relative to the workspace
"""
# Remove any leading slash
path = path.lstrip('/')
# Remove workspace prefix if present
if path.startswith(workspace_path.lstrip('/')):
path = path[len(workspace_path.lstrip('/')):]
# Remove workspace/ prefix if present
if path.startswith('workspace/'):
path = path[9:]
# Remove any remaining leading slash
path = path.lstrip('/')
return path
return path