File size: 1,861 Bytes
5da4770 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import structlog
from typing import Optional
from agentpress.tool import Tool
from agentpress.thread_manager import ThreadManager
from utils.logger import logger
class AgentBuilderBaseTool(Tool):
def __init__(self, thread_manager: ThreadManager, db_connection, agent_id: str):
super().__init__()
self.thread_manager = thread_manager
self.db = db_connection
self.agent_id = agent_id
async def _get_current_account_id(self) -> str:
try:
context_vars = structlog.contextvars.get_contextvars()
thread_id = context_vars.get('thread_id')
if not thread_id:
raise ValueError("No thread_id available from execution context")
client = await self.db.client
thread_result = await client.table('threads').select('account_id').eq('thread_id', thread_id).limit(1).execute()
if not thread_result.data:
raise ValueError(f"Could not find thread with ID: {thread_id}")
account_id = thread_result.data[0]['account_id']
if not account_id:
raise ValueError("Thread has no associated account_id")
return account_id
except Exception as e:
logger.error(f"Error getting current account_id: {e}")
raise
async def _get_agent_data(self) -> Optional[dict]:
try:
client = await self.db.client
result = await client.table('agents').select('*').eq('agent_id', self.agent_id).execute()
if not result.data:
return None
return result.data[0]
except Exception as e:
logger.error(f"Error getting agent data: {e}")
return None |