eliminate sandox redundant information
This commit is contained in:
white-rm
@@ -1,298 +0,0 @@
|
||||
"""
|
||||
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 app.services.supabase import DBConnection
|
||||
from app.services.llm import make_llm_api_call
|
||||
from app.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
|
||||
@@ -1,787 +0,0 @@
|
||||
"""
|
||||
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 app.services.llm import make_llm_api_call
|
||||
from app.agentpress.tool import Tool
|
||||
from app.agentpress.tool_registry import ToolRegistry
|
||||
from app.agentpress.context_manager import ContextManager
|
||||
from app.agentpress.response_processor import (
|
||||
ResponseProcessor,
|
||||
ProcessorConfig
|
||||
)
|
||||
from app.services.supabase import DBConnection
|
||||
from app.utils.logger import logger
|
||||
from langfuse.client import StatefulGenerationClient, StatefulTraceClient
|
||||
from app.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()
|
||||
@@ -7,7 +7,6 @@ This directory contains the agent sandbox implementation - a Docker-based virtua
|
||||
The sandbox provides a complete containerized Linux environment with:
|
||||
- Chrome browser for web interactions
|
||||
- VNC server for accessing the Web User
|
||||
- Web server for serving content (port 8080) -> loading html files from the /workspace directory
|
||||
- Full file system access
|
||||
- Full sudo access
|
||||
|
||||
@@ -18,9 +17,8 @@ You can modify the sandbox environment for development or to add new capabilitie
|
||||
1. Edit files in the `docker/` directory
|
||||
2. Build a custom image:
|
||||
```
|
||||
cd backend/sandbox/docker
|
||||
docker compose build
|
||||
docker push kortix/suna:0.1.3
|
||||
cd daytona/docker
|
||||
docker build
|
||||
```
|
||||
3. Test your changes locally using docker-compose
|
||||
|
||||
@@ -36,10 +34,6 @@ To use your custom sandbox image:
|
||||
|
||||
When publishing a new version of the sandbox:
|
||||
|
||||
1. Update the version number in `docker-compose.yml` (e.g., from `0.1.2` to `0.1.3`)
|
||||
2. Build the new image: `docker compose build`
|
||||
3. Push the new version: `docker push kortix/suna:0.1.3`
|
||||
4. Update all references to the image version in:
|
||||
- `backend/utils/config.py`
|
||||
- Daytona images
|
||||
- Any other services using this image
|
||||
@@ -1 +0,0 @@
|
||||
# Sandbox
|
||||
@@ -6,7 +6,7 @@ services:
|
||||
dockerfile: ${DOCKERFILE:-Dockerfile}
|
||||
args:
|
||||
TARGETPLATFORM: ${TARGETPLATFORM:-linux/amd64}
|
||||
image: kortix/suna:0.1.3
|
||||
image: daytona_sanbox:1.0
|
||||
ports:
|
||||
- "6080:6080" # noVNC web interface
|
||||
- "5901:5901" # VNC port
|
||||
|
||||
@@ -24,7 +24,6 @@ class ThreadMessage:
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert the message to a dictionary for API calls"""
|
||||
return {
|
||||
"thread_id": self.thread_id,
|
||||
"type": self.type,
|
||||
"content": self.content,
|
||||
"is_llm_message": self.is_llm_message,
|
||||
|
||||
@@ -7,662 +7,10 @@ import logging
|
||||
from typing import Optional, Dict,Literal
|
||||
import os
|
||||
from pydantic import Field
|
||||
# from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
|
||||
from app.daytona.tool_base import SandboxToolsBase, Sandbox
|
||||
from app.tool.base import ToolResult
|
||||
|
||||
|
||||
KEYBOARD_KEYS = [
|
||||
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
|
||||
'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
|
||||
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
|
||||
'enter', 'esc', 'backspace', 'tab', 'space', 'delete',
|
||||
'ctrl', 'alt', 'shift', 'win',
|
||||
'up', 'down', 'left', 'right',
|
||||
'f1', 'f2', 'f3', 'f4', 'f5', 'f6', 'f7', 'f8', 'f9', 'f10', 'f11', 'f12',
|
||||
'ctrl+c', 'ctrl+v', 'ctrl+x', 'ctrl+z', 'ctrl+a', 'ctrl+s',
|
||||
'alt+tab', 'alt+f4', 'ctrl+alt+delete'
|
||||
]
|
||||
|
||||
# class ComputerUseTool(SandboxToolsBase):
|
||||
# """Computer automation tool for controlling the sandbox browser and GUI."""
|
||||
|
||||
# def __init__(self, sandbox: Sandbox):
|
||||
# """Initialize automation tool with sandbox connection."""
|
||||
# super().__init__(sandbox)
|
||||
# self.session = None
|
||||
# self.mouse_x = 0 # Track current mouse position
|
||||
# self.mouse_y = 0
|
||||
# # Get automation service URL using port 8000
|
||||
# self.api_base_url = self.sandbox.get_preview_link(8000)
|
||||
# logging.info(f"Initialized Computer Use Tool with API URL: {self.api_base_url}")
|
||||
|
||||
# async def _get_session(self) -> aiohttp.ClientSession:
|
||||
# """Get or create aiohttp session for API requests."""
|
||||
# if self.session is None or self.session.closed:
|
||||
# self.session = aiohttp.ClientSession()
|
||||
# return self.session
|
||||
|
||||
# async def _api_request(self, method: str, endpoint: str, data: Optional[Dict] = None) -> Dict:
|
||||
# """Send request to automation service API."""
|
||||
# try:
|
||||
# session = await self._get_session()
|
||||
# url = f"{self.api_base_url}/api{endpoint}"
|
||||
|
||||
# logging.debug(f"API request: {method} {url} {data}")
|
||||
|
||||
# if method.upper() == "GET":
|
||||
# async with session.get(url) as response:
|
||||
# result = await response.json()
|
||||
# else: # POST
|
||||
# async with session.post(url, json=data) as response:
|
||||
# result = await response.json()
|
||||
|
||||
# logging.debug(f"API response: {result}")
|
||||
# return result
|
||||
|
||||
# except Exception as e:
|
||||
# logging.error(f"API request failed: {str(e)}")
|
||||
# return {"success": False, "error": str(e)}
|
||||
|
||||
# async def cleanup(self):
|
||||
# """Clean up resources."""
|
||||
# if self.session and not self.session.closed:
|
||||
# await self.session.close()
|
||||
# self.session = None
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "move_to",
|
||||
# "description": "Move cursor to specified position",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "x": {
|
||||
# "type": "number",
|
||||
# "description": "X coordinate"
|
||||
# },
|
||||
# "y": {
|
||||
# "type": "number",
|
||||
# "description": "Y coordinate"
|
||||
# }
|
||||
# },
|
||||
# "required": ["x", "y"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="move-to",
|
||||
# mappings=[
|
||||
# {"param_name": "x", "node_type": "attribute", "path": "."},
|
||||
# {"param_name": "y", "node_type": "attribute", "path": "."}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="move_to">
|
||||
# <parameter name="x">100</parameter>
|
||||
# <parameter name="y">200</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def move_to(self, x: float, y: float) -> ToolResult:
|
||||
# """Move cursor to specified position."""
|
||||
# try:
|
||||
# x_int = int(round(float(x)))
|
||||
# y_int = int(round(float(y)))
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/move", {
|
||||
# "x": x_int,
|
||||
# "y": y_int
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# self.mouse_x = x_int
|
||||
# self.mouse_y = y_int
|
||||
# return ToolResult(success=True, output=f"Moved to ({x_int}, {y_int})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to move: {result.get('error', 'Unknown error')}")
|
||||
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to move: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "click",
|
||||
# "description": "Click at current or specified position",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "button": {
|
||||
# "type": "string",
|
||||
# "description": "Mouse button to click",
|
||||
# "enum": ["left", "right", "middle"],
|
||||
# "default": "left"
|
||||
# },
|
||||
# "x": {
|
||||
# "type": "number",
|
||||
# "description": "Optional X coordinate"
|
||||
# },
|
||||
# "y": {
|
||||
# "type": "number",
|
||||
# "description": "Optional Y coordinate"
|
||||
# },
|
||||
# "num_clicks": {
|
||||
# "type": "integer",
|
||||
# "description": "Number of clicks",
|
||||
# "enum": [1, 2, 3],
|
||||
# "default": 1
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="click",
|
||||
# mappings=[
|
||||
# {"param_name": "x", "node_type": "attribute", "path": "x"},
|
||||
# {"param_name": "y", "node_type": "attribute", "path": "y"},
|
||||
# {"param_name": "button", "node_type": "attribute", "path": "button"},
|
||||
# {"param_name": "num_clicks", "node_type": "attribute", "path": "num_clicks"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="click">
|
||||
# <parameter name="x">100</parameter>
|
||||
# <parameter name="y">200</parameter>
|
||||
# <parameter name="button">left</parameter>
|
||||
# <parameter name="num_clicks">1</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def click(self, x: Optional[float] = None, y: Optional[float] = None,
|
||||
# button: str = "left", num_clicks: int = 1) -> ToolResult:
|
||||
# """Click at current or specified position."""
|
||||
# try:
|
||||
# x_val = x if x is not None else self.mouse_x
|
||||
# y_val = y if y is not None else self.mouse_y
|
||||
|
||||
# x_int = int(round(float(x_val)))
|
||||
# y_int = int(round(float(y_val)))
|
||||
# num_clicks = int(num_clicks)
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/click", {
|
||||
# "x": x_int,
|
||||
# "y": y_int,
|
||||
# "clicks": num_clicks,
|
||||
# "button": button.lower()
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# self.mouse_x = x_int
|
||||
# self.mouse_y = y_int
|
||||
# return ToolResult(success=True,
|
||||
# output=f"{num_clicks} {button} click(s) performed at ({x_int}, {y_int})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to click: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to click: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "scroll",
|
||||
# "description": "Scroll the mouse wheel at current position",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "amount": {
|
||||
# "type": "integer",
|
||||
# "description": "Scroll amount (positive for up, negative for down)",
|
||||
# "minimum": -10,
|
||||
# "maximum": 10
|
||||
# }
|
||||
# },
|
||||
# "required": ["amount"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="scroll",
|
||||
# mappings=[
|
||||
# {"param_name": "amount", "node_type": "attribute", "path": "amount"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="scroll">
|
||||
# <parameter name="amount">-3</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def scroll(self, amount: int) -> ToolResult:
|
||||
# """
|
||||
# Scroll the mouse wheel at current position.
|
||||
# Positive values scroll up, negative values scroll down.
|
||||
# """
|
||||
# try:
|
||||
# amount = int(float(amount))
|
||||
# amount = max(-10, min(10, amount))
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/scroll", {
|
||||
# "clicks": amount,
|
||||
# "x": self.mouse_x,
|
||||
# "y": self.mouse_y
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# direction = "up" if amount > 0 else "down"
|
||||
# steps = abs(amount)
|
||||
# return ToolResult(success=True,
|
||||
# output=f"Scrolled {direction} {steps} step(s) at position ({self.mouse_x}, {self.mouse_y})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to scroll: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to scroll: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "typing",
|
||||
# "description": "Type specified text",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "text": {
|
||||
# "type": "string",
|
||||
# "description": "Text to type"
|
||||
# }
|
||||
# },
|
||||
# "required": ["text"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="typing",
|
||||
# mappings=[
|
||||
# {"param_name": "text", "node_type": "content", "path": "text"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="typing">
|
||||
# <parameter name="text">Hello World!</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def typing(self, text: str) -> ToolResult:
|
||||
# """Type specified text."""
|
||||
# try:
|
||||
# text = str(text)
|
||||
|
||||
# result = await self._api_request("POST", "/automation/keyboard/write", {
|
||||
# "message": text,
|
||||
# "interval": 0.01
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# return ToolResult(success=True, output=f"Typed: {text}")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to type: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to type: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "press",
|
||||
# "description": "Press and release a key",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "key": {
|
||||
# "type": "string",
|
||||
# "description": "Key to press",
|
||||
# "enum": KEYBOARD_KEYS
|
||||
# }
|
||||
# },
|
||||
# "required": ["key"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="press",
|
||||
# mappings=[
|
||||
# {"param_name": "key", "node_type": "attribute", "path": "key"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="press">
|
||||
# <parameter name="key">enter</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def press(self, key: str) -> ToolResult:
|
||||
# """Press and release a key."""
|
||||
# try:
|
||||
# key = str(key).lower()
|
||||
|
||||
# result = await self._api_request("POST", "/automation/keyboard/press", {
|
||||
# "keys": key,
|
||||
# "presses": 1
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# return ToolResult(success=True, output=f"Pressed key: {key}")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to press key: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to press key: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "wait",
|
||||
# "description": "Wait for specified duration",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "duration": {
|
||||
# "type": "number",
|
||||
# "description": "Duration in seconds",
|
||||
# "default": 0.5
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="wait",
|
||||
# mappings=[
|
||||
# {"param_name": "duration", "node_type": "attribute", "path": "duration"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="wait">
|
||||
# <parameter name="duration">1.5</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def wait(self, duration: float = 0.5) -> ToolResult:
|
||||
# """Wait for specified duration."""
|
||||
# try:
|
||||
# duration = float(duration)
|
||||
# duration = max(0, min(10, duration))
|
||||
# await asyncio.sleep(duration)
|
||||
# return ToolResult(success=True, output=f"Waited {duration} seconds")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to wait: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "mouse_down",
|
||||
# "description": "Press a mouse button",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "button": {
|
||||
# "type": "string",
|
||||
# "description": "Mouse button to press",
|
||||
# "enum": ["left", "right", "middle"],
|
||||
# "default": "left"
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="mouse-down",
|
||||
# mappings=[
|
||||
# {"param_name": "button", "node_type": "attribute", "path": "button"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="mouse_down">
|
||||
# <parameter name="button">left</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def mouse_down(self, button: str = "left", x: Optional[float] = None, y: Optional[float] = None) -> ToolResult:
|
||||
# """Press a mouse button at current or specified position."""
|
||||
# try:
|
||||
# x_val = x if x is not None else self.mouse_x
|
||||
# y_val = y if y is not None else self.mouse_y
|
||||
|
||||
# x_int = int(round(float(x_val)))
|
||||
# y_int = int(round(float(y_val)))
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/down", {
|
||||
# "x": x_int,
|
||||
# "y": y_int,
|
||||
# "button": button.lower()
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# self.mouse_x = x_int
|
||||
# self.mouse_y = y_int
|
||||
# return ToolResult(success=True, output=f"{button} button pressed at ({x_int}, {y_int})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to press button: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to press button: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "mouse_up",
|
||||
# "description": "Release a mouse button",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "button": {
|
||||
# "type": "string",
|
||||
# "description": "Mouse button to release",
|
||||
# "enum": ["left", "right", "middle"],
|
||||
# "default": "left"
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="mouse-up",
|
||||
# mappings=[
|
||||
# {"param_name": "button", "node_type": "attribute", "path": "button"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="mouse_up">
|
||||
# <parameter name="button">left</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def mouse_up(self, button: str = "left", x: Optional[float] = None, y: Optional[float] = None) -> ToolResult:
|
||||
# """Release a mouse button at current or specified position."""
|
||||
# try:
|
||||
# x_val = x if x is not None else self.mouse_x
|
||||
# y_val = y if y is not None else self.mouse_y
|
||||
|
||||
# x_int = int(round(float(x_val)))
|
||||
# y_int = int(round(float(y_val)))
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/up", {
|
||||
# "x": x_int,
|
||||
# "y": y_int,
|
||||
# "button": button.lower()
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# self.mouse_x = x_int
|
||||
# self.mouse_y = y_int
|
||||
# return ToolResult(success=True, output=f"{button} button released at ({x_int}, {y_int})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to release button: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to release button: {str(e)}")
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "drag_to",
|
||||
# "description": "Drag cursor to specified position",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "x": {
|
||||
# "type": "number",
|
||||
# "description": "Target X coordinate"
|
||||
# },
|
||||
# "y": {
|
||||
# "type": "number",
|
||||
# "description": "Target Y coordinate"
|
||||
# }
|
||||
# },
|
||||
# "required": ["x", "y"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="drag-to",
|
||||
# mappings=[
|
||||
# {"param_name": "x", "node_type": "attribute", "path": "x"},
|
||||
# {"param_name": "y", "node_type": "attribute", "path": "y"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="drag_to">
|
||||
# <parameter name="x">500</parameter>
|
||||
# <parameter name="y">50</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def drag_to(self, x: float, y: float) -> ToolResult:
|
||||
# """Click and drag from current position to target position."""
|
||||
# try:
|
||||
# target_x = int(round(float(x)))
|
||||
# target_y = int(round(float(y)))
|
||||
# start_x = self.mouse_x
|
||||
# start_y = self.mouse_y
|
||||
|
||||
# result = await self._api_request("POST", "/automation/mouse/drag", {
|
||||
# "x": target_x,
|
||||
# "y": target_y,
|
||||
# "duration": 0.3,
|
||||
# "button": "left"
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# self.mouse_x = target_x
|
||||
# self.mouse_y = target_y
|
||||
# return ToolResult(success=True,
|
||||
# output=f"Dragged from ({start_x}, {start_y}) to ({target_x}, {target_y})")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to drag: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to drag: {str(e)}")
|
||||
|
||||
# async def get_screenshot_base64(self) -> Optional[dict]:
|
||||
# """Capture screen and return as base64 encoded image."""
|
||||
# try:
|
||||
# result = await self._api_request("POST", "/automation/screenshot")
|
||||
|
||||
# if "image" in result:
|
||||
# base64_str = result["image"]
|
||||
# timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
# # Save screenshot to file
|
||||
# screenshots_dir = "screenshots"
|
||||
# if not os.path.exists(screenshots_dir):
|
||||
# os.makedirs(screenshots_dir)
|
||||
|
||||
# timestamped_filename = os.path.join(screenshots_dir, f"screenshot_{timestamp}.png")
|
||||
# latest_filename = "latest_screenshot.png"
|
||||
|
||||
# # Decode base64 string and save to file
|
||||
# img_data = base64.b64decode(base64_str)
|
||||
# with open(timestamped_filename, 'wb') as f:
|
||||
# f.write(img_data)
|
||||
|
||||
# # Save a copy as the latest screenshot
|
||||
# with open(latest_filename, 'wb') as f:
|
||||
# f.write(img_data)
|
||||
|
||||
# return {
|
||||
# "content_type": "image/png",
|
||||
# "base64": base64_str,
|
||||
# "timestamp": timestamp,
|
||||
# "filename": timestamped_filename
|
||||
# }
|
||||
# else:
|
||||
# return None
|
||||
|
||||
# except Exception as e:
|
||||
# print(f"[Screenshot] Error during screenshot process: {str(e)}")
|
||||
# return None
|
||||
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "hotkey",
|
||||
# "description": "Press a key combination",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "keys": {
|
||||
# "type": "string",
|
||||
# "description": "Key combination to press",
|
||||
# "enum": KEYBOARD_KEYS
|
||||
# }
|
||||
# },
|
||||
# "required": ["keys"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="hotkey",
|
||||
# mappings=[
|
||||
# {"param_name": "keys", "node_type": "attribute", "path": "keys"}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="hotkey">
|
||||
# <parameter name="keys">ctrl+a</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def hotkey(self, keys: str) -> ToolResult:
|
||||
# """Press a key combination."""
|
||||
# try:
|
||||
# keys = str(keys).lower().strip()
|
||||
# key_sequence = keys.split('+')
|
||||
|
||||
# result = await self._api_request("POST", "/automation/keyboard/hotkey", {
|
||||
# "keys": key_sequence,
|
||||
# "interval": 0.01
|
||||
# })
|
||||
|
||||
# if result.get("success", False):
|
||||
# return ToolResult(success=True, output=f"Pressed key combination: {keys}")
|
||||
# else:
|
||||
# return ToolResult(success=False, output=f"Failed to press keys: {result.get('error', 'Unknown error')}")
|
||||
# except Exception as e:
|
||||
# return ToolResult(success=False, output=f"Failed to press keys: {str(e)}")
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# print("This module should be imported, not run directly.")
|
||||
|
||||
|
||||
|
||||
KEYBOARD_KEYS = [
|
||||
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
|
||||
'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
|
||||
@@ -1030,21 +378,3 @@ class ComputerUseTool(SandboxToolsBase):
|
||||
loop = asyncio.new_event_loop()
|
||||
loop.run_until_complete(self.cleanup())
|
||||
loop.close()
|
||||
# def __del__(self):
|
||||
# """Ensure cleanup when object is destroyed."""
|
||||
# if self.session is not None:
|
||||
# try:
|
||||
# asyncio.run(self.cleanup())
|
||||
# except RuntimeError:
|
||||
# loop = asyncio.new_event_loop()
|
||||
# loop.run_until_complete(self.cleanup())
|
||||
# loop.close()
|
||||
# @classmethod
|
||||
# def create_with_sandbox(cls, sandbox: Sandbox) -> "ComputerUseTool":
|
||||
# """Factory method to create a ComputerUseTool with a sandbox connection."""
|
||||
# tool = cls()
|
||||
# tool.sandbox = sandbox
|
||||
# tool.api_base_url = sandbox.get_preview_link(8000)
|
||||
# logging.info(f"Initialized Computer Use Tool with API URL: {tool.api_base_url}")
|
||||
# return tool
|
||||
|
||||
|
||||
+2
-1070
File diff suppressed because it is too large
Load Diff
@@ -1,147 +0,0 @@
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from agentpress.tool import ToolResult, openapi_schema, xml_schema
|
||||
from daytona.tool_base import SandboxToolsBase
|
||||
from utils.files_utils import clean_path
|
||||
from agentpress.thread_manager import ThreadManager
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
class SandboxDeployTool(SandboxToolsBase):
|
||||
"""Tool for deploying static websites from a Daytona sandbox to Cloudflare Pages."""
|
||||
|
||||
def __init__(self, project_id: str, thread_manager: ThreadManager):
|
||||
super().__init__(project_id, thread_manager)
|
||||
self.workspace_path = "/workspace" # Ensure we're always operating in /workspace
|
||||
self.cloudflare_api_token = os.getenv("CLOUDFLARE_API_TOKEN")
|
||||
|
||||
def clean_path(self, path: str) -> str:
|
||||
"""Clean and normalize a path to be relative to /workspace"""
|
||||
return clean_path(path, self.workspace_path)
|
||||
|
||||
@openapi_schema({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "deploy",
|
||||
"description": "Deploy a static website (HTML+CSS+JS) from a directory in the sandbox to Cloudflare Pages. Only use this tool when permanent deployment to a production environment is needed. The directory path must be relative to /workspace. The website will be deployed to {name}.kortix.cloud.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "Name for the deployment, will be used in the URL as {name}.kortix.cloud"
|
||||
},
|
||||
"directory_path": {
|
||||
"type": "string",
|
||||
"description": "Path to the directory containing the static website files to deploy, relative to /workspace (e.g., 'build')"
|
||||
}
|
||||
},
|
||||
"required": ["name", "directory_path"]
|
||||
}
|
||||
}
|
||||
})
|
||||
@xml_schema(
|
||||
tag_name="deploy",
|
||||
mappings=[
|
||||
{"param_name": "name", "node_type": "attribute", "path": "name"},
|
||||
{"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
|
||||
|
||||
NOTE: If the same name is used, it will redeploy to the same project as before
|
||||
-->
|
||||
|
||||
<function_calls>
|
||||
<invoke name="deploy">
|
||||
<parameter name="name">my-site</parameter>
|
||||
<parameter name="directory_path">website</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
'''
|
||||
)
|
||||
async def deploy(self, name: str, directory_path: str) -> ToolResult:
|
||||
"""
|
||||
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
|
||||
- Failure: Error message if deployment fails
|
||||
"""
|
||||
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)
|
||||
if not dir_info.is_dir:
|
||||
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 &&
|
||||
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",
|
||||
"output": response.result
|
||||
})
|
||||
else:
|
||||
return self.fail_response(f"Deployment failed with exit code {response.exit_code}: {response.result}")
|
||||
except Exception as e:
|
||||
return self.fail_response(f"Error during deployment: {str(e)}")
|
||||
except Exception as e:
|
||||
return self.fail_response(f"Error deploying website: {str(e)}")
|
||||
|
||||
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())
|
||||
|
||||
@@ -1,126 +0,0 @@
|
||||
from agentpress.tool import ToolResult, openapi_schema, xml_schema
|
||||
from daytona.tool_base import SandboxToolsBase
|
||||
from agentpress.thread_manager import ThreadManager
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
class SandboxExposeTool(SandboxToolsBase):
|
||||
"""Tool for exposing and retrieving preview URLs for sandbox ports."""
|
||||
|
||||
def __init__(self, project_id: str, thread_manager: ThreadManager):
|
||||
super().__init__(project_id, thread_manager)
|
||||
|
||||
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({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "expose_port",
|
||||
"description": "Expose a port from the agent's sandbox environment to the public internet and get its preview URL. This is essential for making services running in the sandbox accessible to users, such as web applications, APIs, or other network services. The exposed URL can be shared with users to allow them to interact with the sandbox environment.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"port": {
|
||||
"type": "integer",
|
||||
"description": "The port number to expose. Must be a valid port number between 1 and 65535.",
|
||||
"minimum": 1,
|
||||
"maximum": 65535
|
||||
}
|
||||
},
|
||||
"required": ["port"]
|
||||
}
|
||||
}
|
||||
})
|
||||
@xml_schema(
|
||||
tag_name="expose-port",
|
||||
mappings=[
|
||||
{"param_name": "port", "node_type": "content", "path": "."}
|
||||
],
|
||||
example='''
|
||||
<!-- Example 1: Expose a web server running on port 8000 -->
|
||||
<function_calls>
|
||||
<invoke name="expose_port">
|
||||
<parameter name="port">8000</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
|
||||
<!-- Example 2: Expose an API service running on port 3000 -->
|
||||
<function_calls>
|
||||
<invoke name="expose_port">
|
||||
<parameter name="port">3000</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
|
||||
<!-- Example 3: Expose a development server running on port 5173 -->
|
||||
<function_calls>
|
||||
<invoke name="expose_port">
|
||||
<parameter name="port">5173</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
'''
|
||||
)
|
||||
async def expose_port(self, port: int) -> ToolResult:
|
||||
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.")
|
||||
|
||||
# Wait for sandbox services to be ready (especially important for workflows)
|
||||
services_ready = await self._wait_for_sandbox_services()
|
||||
if not services_ready:
|
||||
return self.fail_response(f"Sandbox services are not fully started yet. Please wait a moment and try again, or ensure a service is running on port {port}.")
|
||||
|
||||
# Check if something is actually listening on the port (for custom ports)
|
||||
if port not in [6080, 8080, 8003]: # Skip check for known sandbox ports
|
||||
try:
|
||||
port_check = self.sandbox.process.exec(f"netstat -tlnp | grep :{port}", timeout=5)
|
||||
if port_check.exit_code != 0:
|
||||
return self.fail_response(f"No service is currently listening on port {port}. Please start a service on this port first.")
|
||||
except Exception:
|
||||
# If we can't check, proceed anyway - the user might be starting a service
|
||||
pass
|
||||
|
||||
# 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:
|
||||
return self.fail_response(f"Error exposing port {port}: {str(e)}")
|
||||
@@ -3,9 +3,7 @@ from typing import Optional, TypeVar
|
||||
from app.tool.base import ToolResult
|
||||
from app.daytona.tool_base import SandboxToolsBase, Sandbox
|
||||
from app.utils.files_utils import should_exclude_file, clean_path
|
||||
# from app.agentpress.thread_manager import ThreadManager
|
||||
from app.utils.logger import logger
|
||||
import os
|
||||
import asyncio
|
||||
|
||||
Context = TypeVar("Context")
|
||||
|
||||
@@ -377,426 +377,3 @@ class SandboxShellTool(SandboxToolsBase):
|
||||
except Exception as e:
|
||||
logger.error(f"Error shell box cleanup action: {e}")
|
||||
pass
|
||||
# from typing import Optional, Dict, Any
|
||||
# import time
|
||||
# from uuid import uuid4
|
||||
# from agentpress.tool import ToolResult, openapi_schema, xml_schema
|
||||
# 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.
|
||||
# Uses sessions for maintaining state between commands and provides comprehensive process management."""
|
||||
#
|
||||
# def __init__(self, project_id: str, thread_manager: ThreadManager):
|
||||
# super().__init__(project_id, thread_manager)
|
||||
# self._sessions: Dict[str, str] = {} # Maps session names to session IDs
|
||||
# self.workspace_path = "/workspace" # Ensure we're always operating in /workspace
|
||||
#
|
||||
# async def _ensure_session(self, session_name: str = "default") -> str:
|
||||
# """Ensure a session exists and return its ID."""
|
||||
# if session_name not in self._sessions:
|
||||
# session_id = str(uuid4())
|
||||
# try:
|
||||
# await self._ensure_sandbox() # Ensure sandbox is initialized
|
||||
# self.sandbox.process.create_session(session_id)
|
||||
# self._sessions[session_name] = session_id
|
||||
# except Exception as e:
|
||||
# raise RuntimeError(f"Failed to create session: {str(e)}")
|
||||
# return self._sessions[session_name]
|
||||
#
|
||||
# async def _cleanup_session(self, session_name: str):
|
||||
# """Clean up a session if it exists."""
|
||||
# if session_name in self._sessions:
|
||||
# try:
|
||||
# await self._ensure_sandbox() # Ensure sandbox is initialized
|
||||
# self.sandbox.process.delete_session(self._sessions[session_name])
|
||||
# del self._sessions[session_name]
|
||||
# except Exception as e:
|
||||
# print(f"Warning: Failed to cleanup session {session_name}: {str(e)}")
|
||||
#
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "execute_command",
|
||||
# "description": "Execute a shell command in the workspace directory. IMPORTANT: Commands are non-blocking by default and run in a tmux session. This is ideal for long-running operations like starting servers or build processes. Uses sessions to maintain state between commands. This tool is essential for running CLI tools, installing packages, and managing system operations.",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "command": {
|
||||
# "type": "string",
|
||||
# "description": "The shell command to execute. Use this for running CLI tools, installing packages, or system operations. Commands can be chained using &&, ||, and | operators."
|
||||
# },
|
||||
# "folder": {
|
||||
# "type": "string",
|
||||
# "description": "Optional relative path to a subdirectory of /workspace where the command should be executed. Example: 'data/pdfs'"
|
||||
# },
|
||||
# "session_name": {
|
||||
# "type": "string",
|
||||
# "description": "Optional name of the tmux session to use. Use named sessions for related commands that need to maintain state. Defaults to a random session name.",
|
||||
# },
|
||||
# "blocking": {
|
||||
# "type": "boolean",
|
||||
# "description": "Whether to wait for the command to complete. Defaults to false for non-blocking execution.",
|
||||
# "default": False
|
||||
# },
|
||||
# "timeout": {
|
||||
# "type": "integer",
|
||||
# "description": "Optional timeout in seconds for blocking commands. Defaults to 60. Ignored for non-blocking commands.",
|
||||
# "default": 60
|
||||
# }
|
||||
# },
|
||||
# "required": ["command"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="execute-command",
|
||||
# mappings=[
|
||||
# {"param_name": "command", "node_type": "content", "path": "."},
|
||||
# {"param_name": "folder", "node_type": "attribute", "path": ".", "required": False},
|
||||
# {"param_name": "session_name", "node_type": "attribute", "path": ".", "required": False},
|
||||
# {"param_name": "blocking", "node_type": "attribute", "path": ".", "required": False},
|
||||
# {"param_name": "timeout", "node_type": "attribute", "path": ".", "required": False}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="execute_command">
|
||||
# <parameter name="command">npm run dev</parameter>
|
||||
# <parameter name="session_name">dev_server</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
#
|
||||
# <!-- Example 2: Running in Specific Directory -->
|
||||
# <function_calls>
|
||||
# <invoke name="execute_command">
|
||||
# <parameter name="command">npm run build</parameter>
|
||||
# <parameter name="folder">frontend</parameter>
|
||||
# <parameter name="session_name">build_process</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
#
|
||||
# <!-- Example 3: Blocking command (wait for completion) -->
|
||||
# <function_calls>
|
||||
# <invoke name="execute_command">
|
||||
# <parameter name="command">npm install</parameter>
|
||||
# <parameter name="blocking">true</parameter>
|
||||
# <parameter name="timeout">300</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def execute_command(
|
||||
# self,
|
||||
# command: str,
|
||||
# folder: Optional[str] = None,
|
||||
# session_name: Optional[str] = None,
|
||||
# blocking: bool = False,
|
||||
# timeout: int = 60
|
||||
# ) -> ToolResult:
|
||||
# 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,
|
||||
# "cwd": cwd,
|
||||
# "completed": True
|
||||
# })
|
||||
# else:
|
||||
# # For non-blocking, just return immediately
|
||||
# return self.success_response({
|
||||
# "session_name": session_name,
|
||||
# "cwd": cwd,
|
||||
# "message": f"Command sent to tmux session '{session_name}'. Use check_command_output to view results.",
|
||||
# "completed": False
|
||||
# })
|
||||
#
|
||||
# except Exception as e:
|
||||
# # Attempt to clean up session in case of error
|
||||
# if session_name:
|
||||
# try:
|
||||
# await self._execute_raw_command(f"tmux kill-session -t {session_name}")
|
||||
# except:
|
||||
# pass
|
||||
# return self.fail_response(f"Error executing command: {str(e)}")
|
||||
#
|
||||
# async def _execute_raw_command(self, command: str) -> Dict[str, Any]:
|
||||
# """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(
|
||||
# command=command,
|
||||
# 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
|
||||
# }
|
||||
#
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "check_command_output",
|
||||
# "description": "Check the output of a previously executed command in a tmux session. Use this to monitor the progress or results of non-blocking commands.",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "session_name": {
|
||||
# "type": "string",
|
||||
# "description": "The name of the tmux session to check."
|
||||
# },
|
||||
# "kill_session": {
|
||||
# "type": "boolean",
|
||||
# "description": "Whether to terminate the tmux session after checking. Set to true when you're done with the command.",
|
||||
# "default": False
|
||||
# }
|
||||
# },
|
||||
# "required": ["session_name"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="check-command-output",
|
||||
# mappings=[
|
||||
# {"param_name": "session_name", "node_type": "attribute", "path": ".", "required": True},
|
||||
# {"param_name": "kill_session", "node_type": "attribute", "path": ".", "required": False}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="check_command_output">
|
||||
# <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">
|
||||
# <parameter name="session_name">build_process</parameter>
|
||||
# <parameter name="kill_session">true</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def check_command_output(
|
||||
# self,
|
||||
# session_name: str,
|
||||
# kill_session: bool = False
|
||||
# ) -> ToolResult:
|
||||
# try:
|
||||
# # Ensure sandbox is initialized
|
||||
# await self._ensure_sandbox()
|
||||
#
|
||||
# # Check if session exists
|
||||
# check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'")
|
||||
# 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)}")
|
||||
#
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "terminate_command",
|
||||
# "description": "Terminate a running command by killing its tmux session.",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {
|
||||
# "session_name": {
|
||||
# "type": "string",
|
||||
# "description": "The name of the tmux session to terminate."
|
||||
# }
|
||||
# },
|
||||
# "required": ["session_name"]
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="terminate-command",
|
||||
# mappings=[
|
||||
# {"param_name": "session_name", "node_type": "attribute", "path": ".", "required": True}
|
||||
# ],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="terminate_command">
|
||||
# <parameter name="session_name">dev_server</parameter>
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def terminate_command(
|
||||
# self,
|
||||
# session_name: str
|
||||
# ) -> ToolResult:
|
||||
# try:
|
||||
# # Ensure sandbox is initialized
|
||||
# await self._ensure_sandbox()
|
||||
#
|
||||
# # Check if session exists
|
||||
# check_result = await self._execute_raw_command(f"tmux has-session -t {session_name} 2>/dev/null || echo 'not_exists'")
|
||||
# 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)}")
|
||||
#
|
||||
# @openapi_schema({
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": "list_commands",
|
||||
# "description": "List all running tmux sessions and their status.",
|
||||
# "parameters": {
|
||||
# "type": "object",
|
||||
# "properties": {}
|
||||
# }
|
||||
# }
|
||||
# })
|
||||
# @xml_schema(
|
||||
# tag_name="list-commands",
|
||||
# mappings=[],
|
||||
# example='''
|
||||
# <function_calls>
|
||||
# <invoke name="list_commands">
|
||||
# </invoke>
|
||||
# </function_calls>
|
||||
# '''
|
||||
# )
|
||||
# async def list_commands(self) -> ToolResult:
|
||||
# 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'):
|
||||
# if line.strip():
|
||||
# parts = line.split(':')
|
||||
# if parts:
|
||||
# session_name = parts[0].strip()
|
||||
# sessions.append(session_name)
|
||||
#
|
||||
# return self.success_response({
|
||||
# "message": f"Found {len(sessions)} active sessions.",
|
||||
# "sessions": sessions
|
||||
# })
|
||||
#
|
||||
# except Exception as e:
|
||||
# return self.fail_response(f"Error listing commands: {str(e)}")
|
||||
#
|
||||
# async def cleanup(self):
|
||||
# """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
|
||||
|
||||
@@ -2,7 +2,6 @@ from pydantic import Field
|
||||
import os
|
||||
import base64
|
||||
import mimetypes
|
||||
import json
|
||||
from typing import Optional
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
|
||||
@@ -1,395 +0,0 @@
|
||||
from tavily import AsyncTavilyClient
|
||||
import httpx
|
||||
from dotenv import load_dotenv
|
||||
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema
|
||||
from utils.config import config
|
||||
from daytona.tool_base import SandboxToolsBase
|
||||
from agentpress.thread_manager import ThreadManager
|
||||
import json
|
||||
import os
|
||||
import datetime
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
# TODO: add subpages, etc... in filters as sometimes its necessary
|
||||
|
||||
class SandboxWebSearchTool(SandboxToolsBase):
|
||||
"""Tool for performing web searches using Tavily API and web scraping using Firecrawl."""
|
||||
|
||||
def __init__(self, project_id: str, thread_manager: ThreadManager):
|
||||
super().__init__(project_id, thread_manager)
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
# Use API keys from config
|
||||
self.tavily_api_key = config.TAVILY_API_KEY
|
||||
self.firecrawl_api_key = config.FIRECRAWL_API_KEY
|
||||
self.firecrawl_url = config.FIRECRAWL_URL
|
||||
|
||||
if not self.tavily_api_key:
|
||||
raise ValueError("TAVILY_API_KEY not found in configuration")
|
||||
if not self.firecrawl_api_key:
|
||||
raise ValueError("FIRECRAWL_API_KEY not found in configuration")
|
||||
|
||||
# Tavily asynchronous search client
|
||||
self.tavily_client = AsyncTavilyClient(api_key=self.tavily_api_key)
|
||||
|
||||
@openapi_schema({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_search",
|
||||
"description": "Search the web for up-to-date information on a specific topic using the Tavily API. This tool allows you to gather real-time information from the internet to answer user queries, research topics, validate facts, and find recent developments. Results include titles, URLs, and publication dates. Use this tool for discovering relevant web pages before potentially crawling them for complete content.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "The search query to find relevant web pages. Be specific and include key terms to improve search accuracy. For best results, use natural language questions or keyword combinations that precisely describe what you're looking for."
|
||||
},
|
||||
"num_results": {
|
||||
"type": "integer",
|
||||
"description": "The number of search results to return. Increase for more comprehensive research or decrease for focused, high-relevance results.",
|
||||
"default": 20
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
}
|
||||
}
|
||||
})
|
||||
@xml_schema(
|
||||
tag_name="web-search",
|
||||
mappings=[
|
||||
{"param_name": "query", "node_type": "attribute", "path": "."},
|
||||
{"param_name": "num_results", "node_type": "attribute", "path": "."}
|
||||
],
|
||||
example='''
|
||||
<function_calls>
|
||||
<invoke name="web_search">
|
||||
<parameter name="query">what is Kortix AI and what are they building?</parameter>
|
||||
<parameter name="num_results">20</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
|
||||
<!-- Another search example -->
|
||||
<function_calls>
|
||||
<invoke name="web_search">
|
||||
<parameter name="query">latest AI research on transformer models</parameter>
|
||||
<parameter name="num_results">20</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
'''
|
||||
)
|
||||
async def web_search(
|
||||
self,
|
||||
query: str,
|
||||
num_results: int = 20
|
||||
) -> ToolResult:
|
||||
"""
|
||||
Search the web using the Tavily API to find relevant and up-to-date information.
|
||||
"""
|
||||
try:
|
||||
# Ensure we have a valid query
|
||||
if not query or not isinstance(query, str):
|
||||
return self.fail_response("A valid search query is required.")
|
||||
|
||||
# Normalize num_results
|
||||
if num_results is None:
|
||||
num_results = 20
|
||||
elif isinstance(num_results, int):
|
||||
num_results = max(1, min(num_results, 50))
|
||||
elif isinstance(num_results, str):
|
||||
try:
|
||||
num_results = max(1, min(int(num_results), 50))
|
||||
except ValueError:
|
||||
num_results = 20
|
||||
else:
|
||||
num_results = 20
|
||||
|
||||
# Execute the search with Tavily
|
||||
logging.info(f"Executing web search for query: '{query}' with {num_results} results")
|
||||
search_response = await self.tavily_client.search(
|
||||
query=query,
|
||||
max_results=num_results,
|
||||
include_images=True,
|
||||
include_answer="advanced",
|
||||
search_depth="advanced",
|
||||
)
|
||||
|
||||
# Check if we have actual results or an answer
|
||||
results = search_response.get('results', [])
|
||||
answer = search_response.get('answer', '')
|
||||
|
||||
# Return the complete Tavily response
|
||||
# This includes the query, answer, results, images and more
|
||||
logging.info(f"Retrieved search results for query: '{query}' with answer and {len(results)} results")
|
||||
|
||||
# Consider search successful if we have either results OR an answer
|
||||
if len(results) > 0 or (answer and answer.strip()):
|
||||
return ToolResult(
|
||||
success=True,
|
||||
output=json.dumps(search_response, ensure_ascii=False)
|
||||
)
|
||||
else:
|
||||
# No results or answer found
|
||||
logging.warning(f"No search results or answer found for query: '{query}'")
|
||||
return ToolResult(
|
||||
success=False,
|
||||
output=json.dumps(search_response, ensure_ascii=False)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
logging.error(f"Error performing web search for '{query}': {error_message}")
|
||||
simplified_message = f"Error performing web search: {error_message[:200]}"
|
||||
if len(error_message) > 200:
|
||||
simplified_message += "..."
|
||||
return self.fail_response(simplified_message)
|
||||
|
||||
@openapi_schema({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "scrape_webpage",
|
||||
"description": "Extract full text content from multiple webpages in a single operation. IMPORTANT: You should ALWAYS collect multiple relevant URLs from web-search results and scrape them all in a single call for efficiency. This tool saves time by processing multiple pages simultaneously rather than one at a time. The extracted text includes the main content of each page without HTML markup.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"urls": {
|
||||
"type": "string",
|
||||
"description": "Multiple URLs to scrape, separated by commas. You should ALWAYS include several URLs when possible for efficiency. Example: 'https://example.com/page1,https://example.com/page2,https://example.com/page3'"
|
||||
}
|
||||
},
|
||||
"required": ["urls"]
|
||||
}
|
||||
}
|
||||
})
|
||||
@xml_schema(
|
||||
tag_name="scrape-webpage",
|
||||
mappings=[
|
||||
{"param_name": "urls", "node_type": "attribute", "path": "."}
|
||||
],
|
||||
example='''
|
||||
<function_calls>
|
||||
<invoke name="scrape_webpage">
|
||||
<parameter name="urls">https://www.kortix.ai/,https://github.com/kortix-ai/suna</parameter>
|
||||
</invoke>
|
||||
</function_calls>
|
||||
'''
|
||||
)
|
||||
async def scrape_webpage(
|
||||
self,
|
||||
urls: str
|
||||
) -> ToolResult:
|
||||
"""
|
||||
Retrieve the complete text content of multiple webpages in a single efficient operation.
|
||||
|
||||
ALWAYS collect multiple relevant URLs from search results and scrape them all at once
|
||||
rather than making separate calls for each URL. This is much more efficient.
|
||||
|
||||
Parameters:
|
||||
- urls: Multiple URLs to scrape, separated by commas
|
||||
"""
|
||||
try:
|
||||
logging.info(f"Starting to scrape webpages: {urls}")
|
||||
|
||||
# Ensure sandbox is initialized
|
||||
await self._ensure_sandbox()
|
||||
|
||||
# Parse the URLs parameter
|
||||
if not urls:
|
||||
logging.warning("Scrape attempt with empty URLs")
|
||||
return self.fail_response("Valid URLs are required.")
|
||||
|
||||
# Split the URLs string into a list
|
||||
url_list = [url.strip() for url in urls.split(',') if url.strip()]
|
||||
|
||||
if not url_list:
|
||||
logging.warning("No valid URLs found in the input")
|
||||
return self.fail_response("No valid URLs provided.")
|
||||
|
||||
if len(url_list) == 1:
|
||||
logging.warning("Only a single URL provided - for efficiency you should scrape multiple URLs at once")
|
||||
|
||||
logging.info(f"Processing {len(url_list)} URLs: {url_list}")
|
||||
|
||||
# Process each URL and collect results
|
||||
results = []
|
||||
for url in url_list:
|
||||
try:
|
||||
# Add protocol if missing
|
||||
if not (url.startswith('http://') or url.startswith('https://')):
|
||||
url = 'https://' + url
|
||||
logging.info(f"Added https:// protocol to URL: {url}")
|
||||
|
||||
# Scrape this URL
|
||||
result = await self._scrape_single_url(url)
|
||||
results.append(result)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing URL {url}: {str(e)}")
|
||||
results.append({
|
||||
"url": url,
|
||||
"success": False,
|
||||
"error": str(e)
|
||||
})
|
||||
|
||||
# Summarize results
|
||||
successful = sum(1 for r in results if r.get("success", False))
|
||||
failed = len(results) - successful
|
||||
|
||||
# Create success/failure message
|
||||
if successful == len(results):
|
||||
message = f"Successfully scraped all {len(results)} URLs. Results saved to:"
|
||||
for r in results:
|
||||
if r.get("file_path"):
|
||||
message += f"\n- {r.get('file_path')}"
|
||||
elif successful > 0:
|
||||
message = f"Scraped {successful} URLs successfully and {failed} failed. Results saved to:"
|
||||
for r in results:
|
||||
if r.get("success", False) and r.get("file_path"):
|
||||
message += f"\n- {r.get('file_path')}"
|
||||
message += "\n\nFailed URLs:"
|
||||
for r in results:
|
||||
if not r.get("success", False):
|
||||
message += f"\n- {r.get('url')}: {r.get('error', 'Unknown error')}"
|
||||
else:
|
||||
error_details = "; ".join([f"{r.get('url')}: {r.get('error', 'Unknown error')}" for r in results])
|
||||
return self.fail_response(f"Failed to scrape all {len(results)} URLs. Errors: {error_details}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
output=message
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
logging.error(f"Error in scrape_webpage: {error_message}")
|
||||
return self.fail_response(f"Error processing scrape request: {error_message[:200]}")
|
||||
|
||||
async def _scrape_single_url(self, url: str) -> dict:
|
||||
"""
|
||||
Helper function to scrape a single URL and return the result information.
|
||||
"""
|
||||
logging.info(f"Scraping single URL: {url}")
|
||||
|
||||
try:
|
||||
# ---------- Firecrawl scrape endpoint ----------
|
||||
logging.info(f"Sending request to Firecrawl for URL: {url}")
|
||||
async with httpx.AsyncClient() as client:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.firecrawl_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"url": url,
|
||||
"formats": ["markdown"]
|
||||
}
|
||||
|
||||
# Use longer timeout and retry logic for more reliability
|
||||
max_retries = 3
|
||||
timeout_seconds = 120
|
||||
retry_count = 0
|
||||
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
logging.info(f"Sending request to Firecrawl (attempt {retry_count + 1}/{max_retries})")
|
||||
response = await client.post(
|
||||
f"{self.firecrawl_url}/v1/scrape",
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
logging.info(f"Successfully received response from Firecrawl for {url}")
|
||||
break
|
||||
except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.ReadError) as timeout_err:
|
||||
retry_count += 1
|
||||
logging.warning(f"Request timed out (attempt {retry_count}/{max_retries}): {str(timeout_err)}")
|
||||
if retry_count >= max_retries:
|
||||
raise Exception(f"Request timed out after {max_retries} attempts with {timeout_seconds}s timeout")
|
||||
# Exponential backoff
|
||||
logging.info(f"Waiting {2 ** retry_count}s before retry")
|
||||
await asyncio.sleep(2 ** retry_count)
|
||||
except Exception as e:
|
||||
# Don't retry on non-timeout errors
|
||||
logging.error(f"Error during scraping: {str(e)}")
|
||||
raise e
|
||||
|
||||
# Format the response
|
||||
title = data.get("data", {}).get("metadata", {}).get("title", "")
|
||||
markdown_content = data.get("data", {}).get("markdown", "")
|
||||
logging.info(f"Extracted content from {url}: title='{title}', content length={len(markdown_content)}")
|
||||
|
||||
formatted_result = {
|
||||
"title": title,
|
||||
"url": url,
|
||||
"text": markdown_content
|
||||
}
|
||||
|
||||
# Add metadata if available
|
||||
if "metadata" in data.get("data", {}):
|
||||
formatted_result["metadata"] = data["data"]["metadata"]
|
||||
logging.info(f"Added metadata: {data['data']['metadata'].keys()}")
|
||||
|
||||
# Create a simple filename from the URL domain and date
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
# Extract domain from URL for the filename
|
||||
from urllib.parse import urlparse
|
||||
parsed_url = urlparse(url)
|
||||
domain = parsed_url.netloc.replace("www.", "")
|
||||
|
||||
# Clean up domain for filename
|
||||
domain = "".join([c if c.isalnum() else "_" for c in domain])
|
||||
safe_filename = f"{timestamp}_{domain}.json"
|
||||
|
||||
logging.info(f"Generated filename: {safe_filename}")
|
||||
|
||||
# Save results to a file in the /workspace/scrape directory
|
||||
scrape_dir = f"{self.workspace_path}/scrape"
|
||||
self.sandbox.fs.create_folder(scrape_dir, "755")
|
||||
|
||||
results_file_path = f"{scrape_dir}/{safe_filename}"
|
||||
json_content = json.dumps(formatted_result, ensure_ascii=False, indent=2)
|
||||
logging.info(f"Saving content to file: {results_file_path}, size: {len(json_content)} bytes")
|
||||
|
||||
self.sandbox.fs.upload_file(
|
||||
json_content.encode(),
|
||||
results_file_path,
|
||||
)
|
||||
|
||||
return {
|
||||
"url": url,
|
||||
"success": True,
|
||||
"title": title,
|
||||
"file_path": results_file_path,
|
||||
"content_length": len(markdown_content)
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
logging.error(f"Error scraping URL '{url}': {error_message}")
|
||||
|
||||
# Create an error result
|
||||
return {
|
||||
"url": url,
|
||||
"success": False,
|
||||
"error": error_message
|
||||
}
|
||||
|
||||
if __name__ == "__main__":
|
||||
async def test_web_search():
|
||||
"""Test function for the web search tool"""
|
||||
# This test function is not compatible with the sandbox version
|
||||
print("Test function needs to be updated for sandbox version")
|
||||
|
||||
async def test_scrape_webpage():
|
||||
"""Test function for the webpage scrape tool"""
|
||||
# This test function is not compatible with the sandbox version
|
||||
print("Test function needs to be updated for sandbox version")
|
||||
|
||||
async def run_tests():
|
||||
"""Run all test functions"""
|
||||
await test_web_search()
|
||||
await test_scrape_webpage()
|
||||
|
||||
asyncio.run(run_tests())
|
||||
Reference in New Issue
Block a user