agent.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import json
import uuid
from pathlib import Path
from typing import Any, Callable
 
import httpx
 
_SYSTEM_PROMPT_PATH = Path(__file__).parent / "prompts" / "system.txt"
 
 
def _load_system_prompt(state: dict[str, Any]) -> str:
    base = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8")
 
    parts: list[str] = []
    model_id = state.get("activeModelId") or state.get("model_id")
    if model_id:
        parts.append(f"Model: {model_id}")
 
    session_mode = state.get("session_mode")
    active_model = state.get("activeModelId")
    if active_model:
        try:
            from aquin.session_mode import mode_for_model_id
            session_mode = mode_for_model_id(str(active_model))
        except ValueError:
            pass
    if session_mode == "llm" or active_model:
        parts.append("Mode: LLM")
    else:
        parts.append("Mode: (no model loaded — ask user to run aquin load)")
 
    last_prompt = state.get("lastPrompt")
    if last_prompt:
        parts.append(f'Last inspected prompt: "{last_prompt}"')
 
    top_features = state.get("lastTopFeatures") or []
    if top_features:
        lines = [f"  - F{f['feature_idx']} \"{f['label']}\" (act={f['activation']:.2f})" for f in top_features[:5]]
        parts.append("Top features from last run:\n" + "\n".join(lines))
 
    memory = state.get("memory") or {}
    if memory:
        mem_lines = [f"  {k}: {v}" for k, v in list(memory.items())[:10]]
        parts.append("Session memory:\n" + "\n".join(mem_lines))
 
    if parts:
        base += "\n\n[Current context]\n" + "\n".join(parts)
 
    return base
 
 
def _build_tools_schema(state: dict[str, Any]) -> list[dict[str, Any]]:
    """Return OpenAI-format tool definitions for the active session mode."""
    try:
        from aquin.session_mode import get_filtered_tool_schemas, mode_for_active_model, mode_for_model_id
 
        model_id = state.get("activeModelId")
        if model_id:
            try:
                mode = mode_for_model_id(str(model_id))
            except ValueError:
                mode = mode_for_active_model()
        else:
            mode = mode_for_active_model()
        if mode is None:
            return []
        return get_filtered_tool_schemas(mode)
    except Exception:
        return []
 
 
def run_agent_turn(
    session_id: str,
    user_message: str,
    state: dict[str, Any],
    api_key: str,
    base_url: str,
    on_token: Callable[[str], None] | None = None,
    on_tool_start: Callable[[str, dict], None] | None = None,
    on_tool_result: Callable[[str, Any, bool], None] | None = None,
) -> None:
    """
    Run one user turn through the orchestrator loop.
 
    Streams text as message.append events and dispatches tool calls via the
    registry, pushing tool.start + tool.result for each. Loops until the model
    stops requesting tool calls.
    """
    from aquin.engine.sync import push_events
    from aquin.engine.tools import dispatch as dispatch_tool  # import via __init__ to trigger stubs registration
 
    system_prompt = _load_system_prompt(state)
    tools_schema = _build_tools_schema(state)
 
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message},
    ]
 
    agent_url = f"{base_url.rstrip('/')}/api/sync/openai/chat"
    assistant_msg_id = str(uuid.uuid4())
 
    while True:
        # ── Stream the assistant response via Aquin proxy ────────────────────
        assistant_text = ""
        pending_tool_calls: list[dict[str, Any]] = []
 
        request_body: dict[str, Any] = {
            "model": "gpt-4o-mini",
            "messages": messages,
        }
        if tools_schema:
            request_body["tools"] = tools_schema
            request_body["tool_choice"] = "auto"
 
        with httpx.stream(
            "POST",
            agent_url,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json=request_body,
            timeout=120,
            verify=False,
        ) as resp:
            resp.raise_for_status()
 
            # Accumulate streamed chunks
            tool_call_accum: dict[int, dict[str, Any]] = {}
 
            for line in resp.iter_lines():
                if not line.startswith("data: "):
                    continue
                raw = line[len("data: "):]
                if raw == "[DONE]":
                    break
                try:
                    chunk = json.loads(raw)
                except json.JSONDecodeError:
                    continue
 
                delta = chunk.get("choices", [{}])[0].get("delta", {})
 
                # Text content — buffer locally, push to web only on flush
                text_piece = delta.get("content") or ""
                if text_piece:
                    assistant_text += text_piece
                    if on_token:
                        on_token(text_piece)
                    else:
                        print(text_piece, end="", flush=True)
 
                # Tool call deltas
                for tc_delta in delta.get("tool_calls") or []:
                    idx = tc_delta.get("index", 0)
                    if idx not in tool_call_accum:
                        tool_call_accum[idx] = {
                            "id": tc_delta.get("id", ""),
                            "type": "function",
                            "function": {"name": "", "arguments": ""},
                        }
                    acc = tool_call_accum[idx]
                    if tc_delta.get("id"):
                        acc["id"] = tc_delta["id"]
                    fn = tc_delta.get("function", {})
                    if fn.get("name"):
                        acc["function"]["name"] += fn["name"]
                    if fn.get("arguments"):
                        acc["function"]["arguments"] += fn["arguments"]
 
            pending_tool_calls = list(tool_call_accum.values())
 
        # ── Push completed assistant message to sync ──────────────────────────
        if assistant_text:
            push_events(
                session_id=session_id,
                events=[{
                    "type": "message.append",
                    "payload": {
                        "message": {
                            "id": assistant_msg_id,
                            "role": "assistant",
                            "content": assistant_text,
                            "toolCalls": [],
                        }
                    },
                }],
                api_key=api_key,
                base_url=base_url,
            )
 
        # ── No tool calls — turn is done ─────────────────────────────────────
        if not pending_tool_calls:
            if assistant_text:
                messages.append({"role": "assistant", "content": assistant_text})
            break
 
        # ── Add assistant message with tool_calls to history ─────────────────
        messages.append({
            "role": "assistant",
            "content": assistant_text or None,
            "tool_calls": pending_tool_calls,
        })
 
        # ── Dispatch each tool call ───────────────────────────────────────────
        for tc in pending_tool_calls:
            tool_name = tc["function"]["name"]
            tool_call_id = tc["id"]
 
            try:
                args = json.loads(tc["function"]["arguments"] or "{}")
            except json.JSONDecodeError:
                args = {}
 
            if on_tool_start:
                on_tool_start(tool_name, args)
 
            # Push tool.start
            push_events(
                session_id=session_id,
                events=[{
                    "type": "tool.start",
                    "payload": {
                        "tool_call_id": tool_call_id,
                        "tool_name": tool_name,
                        "args": args,
                    },
                }],
                api_key=api_key,
                base_url=base_url,
            )
 
            # Dispatch to registry
            result: Any = None
            try:
                from aquin.session_mode import is_tool_allowed, mode_for_active_model, mode_for_model_id
 
                model_id = state.get("activeModelId")
                if model_id:
                    try:
                        mode = mode_for_model_id(str(model_id))
                    except ValueError:
                        mode = mode_for_active_model()
                else:
                    mode = mode_for_active_model()
                if mode is None:
                    raise NotImplementedError("Load a model first: aquin load --model <id>")
                if not is_tool_allowed(tool_name, mode):
                    raise NotImplementedError(
                        f"Tool '{tool_name}' is not available in {mode} session mode"
                    )
 
                result = dispatch_tool(tool_name, args, {
                    "session_id": session_id,
                    "api_key": api_key,
                    "base_url": base_url,
                    "state": state,
                    "tool_call_id": tool_call_id,
                })
                card = result.get("card") if isinstance(result, dict) else None
                if card is None and isinstance(result, dict):
                    from aquin.compute import card_mapper
                    card = card_mapper.to_card(tool_name, result)
                if isinstance(result, dict):
                    result_content = result.get("content") or {
                        k: v for k, v in result.items() if k not in ("card", "capture")
                    }
                else:
                    result_content = result
                if isinstance(result_content, dict):
                    from aquin.compute.sync_slim import slim_tool_result_for_sync
                    result_content = slim_tool_result_for_sync(tool_name, result_content, card)
            except NotImplementedError as exc:
                result_content = {"error": str(exc)}
                card = None
                result = None
            except Exception as exc:
                result_content = {"error": f"Tool '{tool_name}' failed: {exc}"}
                card = None
                result = None
 
            if on_tool_result:
                is_error = isinstance(result_content, dict) and "error" in result_content
                on_tool_result(tool_name, result_content, is_error)
 
            # Push tool.result
            tool_result_payload: dict[str, Any] = {
                "tool_call_id": tool_call_id,
                "result": result_content,
            }
            if card is not None:
                tool_result_payload["card"] = card
 
            result_events: list[dict] = [{"type": "tool.result", "payload": tool_result_payload}]
            if isinstance(result, dict) and result.get("capture"):
                result_events.append({"type": "capture.ready", "payload": result["capture"]})
 
            push_events(
                session_id=session_id,
                events=result_events,
                api_key=api_key,
                base_url=base_url,
                timeout=90.0 if tool_name == "ensure_umap_loaded" else 15.0,
            )
 
            # Append tool result to message history for next loop iteration
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call_id,
                "content": json.dumps(result_content) if not isinstance(result_content, str) else result_content,
            })
 
        # Re-assign a fresh message id for any subsequent assistant turn
        assistant_msg_id = str(uuid.uuid4())