from __future__ import annotations import uuid from dataclasses import dataclass from typing import Any, Awaitable, Callable from traderai.memory import MemoryStore from traderai.scheduler import WakeScheduler from traderai.uex_client import UEXClient ToolHandler = Callable[..., Awaitable[dict[str, Any]]] @dataclass class PendingAction: id: str label: str endpoint: str payload: dict[str, Any] class ToolRegistry: def __init__( self, uex: UEXClient, require_write_approval: bool = True, memory: MemoryStore | None = None, scheduler: WakeScheduler | None = None, ) -> None: self.uex = uex self.require_write_approval = require_write_approval self.memory = memory self.scheduler = scheduler self.pending_actions: dict[str, PendingAction] = {} self.handlers: dict[str, ToolHandler] = { "search_marketplace_listings": self.search_marketplace_listings, "get_marketplace_listing": self.get_marketplace_listing, "list_marketplace_negotiations": self.list_marketplace_negotiations, "get_negotiation_messages": self.get_negotiation_messages, "draft_negotiation_message": self.draft_negotiation_message, "draft_marketplace_listing": self.draft_marketplace_listing, "remember_user_fact": self.remember_user_fact, "recall_memory": self.recall_memory, "schedule_wake_job": self.schedule_wake_job, "list_wake_jobs": self.list_wake_jobs, "check_uex_notifications": self.check_uex_notifications, } @property def schemas(self) -> list[dict[str, Any]]: return [ { "type": "function", "function": { "name": "search_marketplace_listings", "description": "Search active/current UEX marketplace listings only. Prices are in-game aUEC/UEC credits, not real-world dollars. Do not use this as historical sale or completed-sale information. UEX returns up to 100 active listings; filters are applied locally.", "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "Text to search in title, description, location, advertiser, or slug."}, "operation": {"type": "string", "enum": ["buy", "sell"]}, "type": {"type": "string", "enum": ["item", "service", "contract"]}, "username": {"type": "string", "description": "Advertiser IGN."}, "location": {"type": "string"}, "min_price": {"type": "number"}, "max_price": {"type": "number"}, "limit": {"type": "integer", "minimum": 1, "maximum": 25}, }, }, }, }, { "type": "function", "function": { "name": "get_marketplace_listing", "description": "Fetch a specific UEX marketplace listing by id or slug.", "parameters": { "type": "object", "properties": { "id": {"type": "integer"}, "slug": {"type": "string"}, }, }, }, }, { "type": "function", "function": { "name": "list_marketplace_negotiations", "description": "List authenticated marketplace negotiations for the configured UEX user.", "parameters": { "type": "object", "properties": { "id": {"type": "integer"}, "id_listing": {"type": "integer"}, "hash": {"type": "string"}, }, }, }, }, { "type": "function", "function": { "name": "get_negotiation_messages", "description": "Fetch authenticated messages from a marketplace negotiation by hash or id_negotiation.", "parameters": { "type": "object", "properties": { "hash": {"type": "string"}, "id_negotiation": {"type": "integer"}, }, }, }, }, { "type": "function", "function": { "name": "draft_negotiation_message", "description": "Draft a message or offer to a UEX negotiation. This creates a pending action that must be approved before sending.", "parameters": { "type": "object", "required": ["message"], "properties": { "message": {"type": "string"}, "hash": {"type": "string"}, "id_negotiation": {"type": "integer"}, "is_production": {"type": "integer", "enum": [0, 1], "default": 0}, }, }, }, }, { "type": "function", "function": { "name": "draft_marketplace_listing", "description": "Draft a new UEX marketplace listing. Listing prices are in-game aUEC/UEC credits, not real-world dollars. This creates a pending action that must be approved before posting.", "parameters": { "type": "object", "required": ["id_category", "operation", "type", "unit", "title", "description", "price", "currency", "language"], "properties": { "id_item": {"type": "integer"}, "id_star_system": {"type": "integer"}, "id_organization": {"type": "integer"}, "id_category": {"type": "integer"}, "operation": {"type": "string", "enum": ["buy", "sell"]}, "type": {"type": "string", "enum": ["item", "service", "contract"]}, "unit": {"type": "string"}, "title": {"type": "string"}, "description": {"type": "string"}, "price": {"type": "number"}, "currency": {"type": "string", "enum": ["UEC", "WIF"]}, "language": {"type": "string", "default": "en_US"}, "location": {"type": "string"}, "source": {"type": "string"}, "availability": {"type": "string"}, "in_stock": {"type": "integer"}, "hours_expiration": {"type": "integer"}, "is_hidden": {"type": "integer", "enum": [0, 1]}, "is_production": {"type": "integer", "enum": [0, 1], "default": 0}, }, }, }, }, { "type": "function", "function": { "name": "remember_user_fact", "description": "Persist a durable user preference, identity detail, trading rule, or long-term note for future chats.", "parameters": { "type": "object", "required": ["content"], "properties": { "content": {"type": "string"}, "kind": {"type": "string", "enum": ["user", "preference", "trading", "project", "note"], "default": "note"}, "importance": {"type": "integer", "minimum": 1, "maximum": 5, "default": 3}, }, }, }, }, { "type": "function", "function": { "name": "recall_memory", "description": "Search long-term memory for relevant prior facts, preferences, and chat context.", "parameters": { "type": "object", "properties": { "query": {"type": "string"}, "limit": {"type": "integer", "minimum": 1, "maximum": 10, "default": 6}, }, }, }, }, { "type": "function", "function": { "name": "schedule_wake_job", "description": "Create a scheduled wake-up job for the assistant. Use either run_at for one-time jobs or cron for recurring jobs.", "parameters": { "type": "object", "required": ["prompt"], "properties": { "prompt": {"type": "string", "description": "What the AI should consider or do when it wakes."}, "run_at": {"type": "string", "description": "ISO datetime for a one-time wake job, such as 2026-05-05T20:30:00-04:00."}, "cron": {"type": "string", "description": "Five-field cron expression for recurring jobs, such as 0 9 * * *."}, }, }, }, }, { "type": "function", "function": { "name": "list_wake_jobs", "description": "List currently enabled scheduled assistant wake jobs.", "parameters": {"type": "object", "properties": {}}, }, }, { "type": "function", "function": { "name": "check_uex_notifications", "description": "Check authenticated UEX user notifications and return unread pending notifications.", "parameters": {"type": "object", "properties": {}}, }, }, ] async def execute(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: handler = self.handlers.get(name) if not handler: return {"error": f"Unknown tool: {name}"} try: return await handler(**arguments) except Exception as exc: return {"error": str(exc)} async def approve(self, action_id: str) -> dict[str, Any]: action = self.pending_actions.pop(action_id, None) if not action: return {"error": f"Pending action not found: {action_id}"} return await self.uex.post(action.endpoint, action.payload, authenticated=True) async def decline(self, action_id: str) -> dict[str, Any]: action = self.pending_actions.pop(action_id, None) if not action: return {"error": f"Pending action not found: {action_id}"} return { "declined": True, "pending_action": { "id": action.id, "label": action.label, "endpoint": action.endpoint, "payload": action.payload, }, } async def search_marketplace_listings( self, query: str | None = None, operation: str | None = None, type: str | None = None, username: str | None = None, location: str | None = None, min_price: float | None = None, max_price: float | None = None, limit: int = 10, ) -> dict[str, Any]: response = await self.uex.get("marketplace_listings", {"username": username}) listings = response.get("data") or [] filtered = [] q = (query or "").casefold() loc = (location or "").casefold() for listing in listings: if operation and listing.get("operation") != operation: continue if type and listing.get("type") != type: continue if min_price is not None and float(listing.get("price") or 0) < min_price: continue if max_price is not None and float(listing.get("price") or 0) > max_price: continue if loc and loc not in str(listing.get("location") or "").casefold(): continue haystack = " ".join(str(listing.get(k) or "") for k in ["title", "description", "location", "user_username", "slug"]).casefold() if q and q not in haystack: continue filtered.append(self._summarize_listing(listing)) if len(filtered) >= max(1, min(limit, 25)): break return {"count": len(filtered), "listings": filtered} async def get_marketplace_listing(self, id: int | None = None, slug: str | None = None) -> dict[str, Any]: response = await self.uex.get("marketplace_listings", {"id": id, "slug": slug}) return {"listing": response.get("data")} async def list_marketplace_negotiations( self, id: int | None = None, id_listing: int | None = None, hash: str | None = None, ) -> dict[str, Any]: return await self.uex.get("marketplace_negotiations", {"id": id, "id_listing": id_listing, "hash": hash}, authenticated=True) async def get_negotiation_messages(self, hash: str | None = None, id_negotiation: int | None = None) -> dict[str, Any]: return await self.uex.get("marketplace_negotiations_messages", {"hash": hash, "id_negotiation": id_negotiation}, authenticated=True) async def draft_negotiation_message( self, message: str, hash: str | None = None, id_negotiation: int | None = None, is_production: int = 0, ) -> dict[str, Any]: payload = {"message": message, "hash": hash, "id_negotiation": id_negotiation, "is_production": is_production} return self._pending("Send negotiation message", "marketplace_negotiations_messages", payload) async def draft_marketplace_listing(self, **payload: Any) -> dict[str, Any]: return self._pending("Post marketplace listing", "marketplace_advertise", payload) async def remember_user_fact(self, content: str, kind: str = "note", importance: int = 3) -> dict[str, Any]: if self.memory is None: return {"error": "Memory store is not configured."} return {"memory": self.memory.remember(kind, content, importance)} async def recall_memory(self, query: str = "", limit: int = 6) -> dict[str, Any]: if self.memory is None: return {"error": "Memory store is not configured."} return {"memories": self.memory.recall(query, max(1, min(limit, 10)))} async def schedule_wake_job( self, prompt: str, run_at: str | None = None, cron: str | None = None, ) -> dict[str, Any]: if self.scheduler is None: return {"error": "Scheduler is not configured."} if bool(run_at) == bool(cron): return {"error": "Provide exactly one of run_at or cron."} if run_at: return {"scheduled_job": self.scheduler.schedule_date(run_at, prompt)} return {"scheduled_job": self.scheduler.schedule_cron(cron or "", prompt)} async def list_wake_jobs(self) -> dict[str, Any]: if self.scheduler is None: return {"error": "Scheduler is not configured."} return {"scheduled_jobs": self.scheduler.list_jobs()} async def check_uex_notifications(self) -> dict[str, Any]: response = await self.uex.get_user_notifications() notifications = response.get("notifications") or [] pending = [item for item in notifications if not item.get("date_read")] return {"count": len(pending), "notifications": pending} def _pending(self, label: str, endpoint: str, payload: dict[str, Any]) -> dict[str, Any]: action_id = str(uuid.uuid4()) payload = {key: value for key, value in payload.items() if value is not None} self.pending_actions[action_id] = PendingAction(action_id, label, endpoint, payload) return { "pending_action": { "id": action_id, "label": label, "endpoint": endpoint, "payload": payload, "approval_required": self.require_write_approval, } } @staticmethod def _summarize_listing(listing: dict[str, Any]) -> dict[str, Any]: return { "id": listing.get("id"), "slug": listing.get("slug"), "title": listing.get("title"), "operation": listing.get("operation"), "type": listing.get("type"), "price": listing.get("price"), "currency": listing.get("currency"), "unit": listing.get("unit"), "location": listing.get("location"), "availability": listing.get("availability"), "in_stock": listing.get("in_stock"), "advertiser": listing.get("user_username"), "expires_at": listing.get("date_expiration"), }