agent.py

"""
agent.py - LLM ReAct agent brain for the Agentic Wallet SDK.
Uses OpenRouter (OpenAI-compatible) function calling.
All actions validated by SafetyGuard before execution.
"""
import asyncio, json, os
from typing import Optional, Dict, Any, List
from openai import AsyncOpenAI
from wallet  import SolanaWallet, DEVNET_RPC, MAINNET_RPC
from jupiter import JupiterClient
from safety  import SafetyGuard, SafetyConfig
from logger  import AgentLogger

OR_BASE = "https://openrouter.ai/api/v1"
DEFAULT_MODEL = "anthropic/claude-3.5-haiku"

TOOLS: List[Dict] = [
    {"type": "function", "function": {
        "name": "get_balance",
        "description": "Get current SOL balance and full token portfolio.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "transfer_sol",
        "description": "Send SOL to a recipient. Safety limits enforced.",
        "parameters": {"type": "object",
            "properties": {
                "recipient":  {"type": "string",  "description": "Recipient Solana address (base58)"},
                "amount_sol": {"type": "number",  "description": "Amount in SOL"}},
            "required": ["recipient", "amount_sol"]}}},
    {"type": "function", "function": {
        "name": "swap_tokens",
        "description": "Swap SOL to another token via Jupiter DEX aggregator.",
        "parameters": {"type": "object",
            "properties": {
                "output_token":  {"type": "string",  "description": "Token symbol (USDC, BONK, JUP, RAY, MSOL)"},
                "sol_amount":    {"type": "number",  "description": "SOL amount to swap"},
                "slippage_bps":  {"type": "integer", "description": "Slippage in bps (default 50)", "default": 50}},
            "required": ["output_token", "sol_amount"]}}},
    {"type": "function", "function": {
        "name": "get_token_price",
        "description": "Get USD price of a token via Jupiter Price API.",
        "parameters": {"type": "object",
            "properties": {"token": {"type": "string", "description": "Token symbol (SOL, USDC, BONK...)"}},
            "required": ["token"]}}},
    {"type": "function", "function": {
        "name": "get_transaction_history",
        "description": "Get recent transaction signatures for the wallet.",
        "parameters": {"type": "object",
            "properties": {"limit": {"type": "integer", "description": "Number of txns (max 20)", "default": 5}},
            "required": []}}},
    {"type": "function", "function": {
        "name": "get_safety_status",
        "description": "Check safety guard status: limits, daily spend, rate limits, emergency stop.",
        "parameters": {"type": "object", "properties": {}, "required": []}}},
    {"type": "function", "function": {
        "name": "request_devnet_airdrop",
        "description": "Request devnet SOL airdrop for testing (devnet only).",
        "parameters": {"type": "object",
            "properties": {"amount_sol": {"type": "number", "description": "SOL to request (max 2.0)", "default": 1.0}},
            "required": []}}},
    {"type": "function", "function": {
        "name": "done",
        "description": "Signal task completion with a summary of actions taken.",
        "parameters": {"type": "object",
            "properties": {"summary": {"type": "string", "description": "Human-readable completion summary"}},
            "required": ["summary"]}}},
]


class AgenticWallet:
    """
    AI-controlled Solana wallet.
    run(instruction) executes a full ReAct loop:
      Reason -> Select tool -> Safety check -> Execute -> Observe -> Repeat -> Done
    """
    def __init__(self, wallet: SolanaWallet, safety: SafetyGuard,
                 logger: AgentLogger, openrouter_api_key: str,
                 model: str = DEFAULT_MODEL, max_iterations: int = 12):
        self.wallet         = wallet
        self.safety         = safety
        self.log            = logger
        self.model          = model
        self.max_iterations = max_iterations
        self._llm = AsyncOpenAI(api_key=openrouter_api_key, base_url=OR_BASE)
        self._jupiter: Optional[JupiterClient] = None

    async def _get_jup(self) -> JupiterClient:
        if not self._jupiter:
            c = await self.wallet._client_()
            self._jupiter = JupiterClient(self.wallet.keypair, c)
        return self._jupiter

    def _system_prompt(self) -> str:
        s = self.safety.get_status()
        return (
            f"You are an autonomous Solana wallet agent controlling a real wallet.\n"
            f"All transactions are IRREVERSIBLE. Think carefully before acting.\n\n"
            f"Wallet:   {self.wallet.get_public_key()}\n"
            f"Network:  {self.wallet.network}\n"
            f"Explorer: {self.wallet.get_explorer_url()}\n\n"
            f"Safety limits (enforced independently, you CANNOT override them):\n"
            f"  max single tx:    {s['max_single_tx_sol']} SOL\n"
            f"  max single swap:  {s['max_single_swap_sol']} SOL\n"
            f"  daily remaining:  {s['daily_remaining_sol']} SOL\n"
            f"  txns this hour:   {s['txns_last_hour']}/{s['max_txns_per_hour']}\n"
            f"  emergency stop:   {s['emergency_stop']}\n"
            f"  transfers:        {s['allow_transfers']}\n"
            f"  swaps:            {s['allow_swaps']}\n\n"
            "Guidelines:\n"
            "- Use get_balance first if you need current balance info\n"
            "- Never attempt actions you know will be blocked\n"
            "- Always call done() when finished with a full summary\n"
            "- On devnet, airdrops are free - use them freely for testing\n"
        )

    async def _exec(self, name: str, args: Dict) -> str:
        self.log.log_action(name, json.dumps(args)[:100])
        try:
            if name == "get_balance":
                result = await self.wallet.get_portfolio()
            elif name == "transfer_sol":
                to, amt = args["recipient"], float(args["amount_sol"])
                bal = await self.wallet.get_sol_balance()
                chk = self.safety.check_transfer(to, amt, bal)
                self.log.log_safety_check(f"transfer {amt} SOL", chk.decision.value, chk.reason)
                if not chk.allowed: return json.dumps({"error": f"SAFETY DENIED: {chk.reason}"})
                result = await self.wallet.transfer_sol(to, amt)
                self.safety.record_executed(amt)
                self.log.log_transaction("transfer", str(self.wallet.pubkey), to, amt, "SOL",
                                         result.get("signature"), True)
            elif name == "swap_tokens":
                out  = args["output_token"]
                amt  = float(args["sol_amount"])
                slip = int(args.get("slippage_bps", 50))
                bal  = await self.wallet.get_sol_balance()
                chk  = self.safety.check_swap("SOL", out, amt, bal)
                self.log.log_safety_check(f"swap {amt} SOL>{out}", chk.decision.value, chk.reason)
                if not chk.allowed: return json.dumps({"error": f"SAFETY DENIED: {chk.reason}"})
                jup    = await self._get_jup()
                result = await jup.swap_sol_to_token(out, amt, slip)
                self.safety.record_executed(amt)
                self.log.log_transaction("swap", str(self.wallet.pubkey), None,
                                         amt, f"SOL>{out}", result.get("signature"), True)
            elif name == "get_token_price":
                jup    = await self._get_jup()
                result = await jup.get_token_price(args["token"])
            elif name == "get_transaction_history":
                result = await self.wallet.get_transaction_history(limit=min(int(args.get("limit",5)),20))
            elif name == "get_safety_status":
                result = self.safety.get_status()
            elif name == "request_devnet_airdrop":
                amt    = min(float(args.get("amount_sol", 1.0)), 2.0)
                sig    = await self.wallet.request_airdrop(amt)
                result = {"success": True, "amount_sol": amt, "signature": sig}
                self.log.log_transaction("airdrop", "devnet-faucet",
                                         str(self.wallet.pubkey), amt, "SOL", sig, True)
            elif name == "done":
                result = {"done": True, "summary": args.get("summary", "")}
            else:
                result = {"error": f"Unknown tool: {name}"}
        except Exception as e:
            self.log.error(f"Tool {name} raised", error=str(e))
            result = {"error": str(e)}
        return json.dumps(result, default=str)

    async def run(self, instruction: str) -> Dict[str, Any]:
        self.log.info(f"Task: {instruction}")
        messages = [
            {"role": "system", "content": self._system_prompt()},
            {"role": "user",   "content": instruction},
        ]
        for iteration in range(1, self.max_iterations + 1):
            self.log.info(f"ReAct iteration {iteration}/{self.max_iterations}")
            resp = await self._llm.chat.completions.create(
                model=self.model, messages=messages,
                tools=TOOLS, tool_choice="auto", max_tokens=1024)
            msg = resp.choices[0].message
            messages.append(msg.model_dump(exclude_unset=True))
            if msg.tool_calls:
                for tc in msg.tool_calls:
                    fn   = tc.function.name
                    args = json.loads(tc.function.arguments)
                    self.log.log_llm_decision(instruction[:60], fn, json.dumps(args)[:150])
                    res  = await self._exec(fn, args)
                    messages.append({"role": "tool", "tool_call_id": tc.id, "content": res})
                    if fn == "done":
                        summary = args.get("summary", "Task complete")
                        self.log.info(f"Done after {iteration} iterations")
                        return {"success": True, "summary": summary,
                                "iterations": iteration, "session": self.log.get_session_summary()}
            else:
                return {"success": True, "summary": msg.content or "Task complete",
                        "iterations": iteration, "session": self.log.get_session_summary()}
        return {"success": False, "summary": f"Max iterations ({self.max_iterations}) reached",
                "iterations": self.max_iterations}

    async def close(self):
        await self.wallet.close()
        if self._jupiter: await self._jupiter.close()


async def create_agent(
        keypair_path: Optional[str] = None,
        network: str = "devnet",
        openrouter_api_key: Optional[str] = None,
        safety_config: Optional[SafetyConfig] = None,
        model: str = DEFAULT_MODEL,
        verbose: bool = True) -> AgenticWallet:
    key = openrouter_api_key or os.getenv("OPENROUTER_API_KEY", "")
    if not key: raise ValueError("OPENROUTER_API_KEY is required")
    rpc = DEVNET_RPC if network == "devnet" else MAINNET_RPC
    w   = SolanaWallet(keypair_path=keypair_path, rpc_url=rpc, network=network)
    s   = SafetyGuard(safety_config or SafetyConfig.devnet_testing())
    lg  = AgentLogger(verbose=verbose)
    return AgenticWallet(w, s, lg, key, model)