src/api/server.py

"""
Frox AI — Morph API Server

The connective tissue between the Morph model and every client:
frox-chat, frox-code (or any Tauri/Electron desktop app), myai-cli,
frox-mobile, and anything else you build. All of them talk to this one
HTTP surface — none of them import Python or know anything about
MorphInferenceEngine directly.

Contract: OpenAI-compatible (/v1/chat/completions, /v1/models),
including tool/function calling (`tools` in the request, `tool_calls`
in the response — bridged to Morph's native <|tool_call|> format),
plus Frox-specific extensions: `session_id` for persistent KV-cache
reuse, /v1/tools/execute for direct tool access, and a "conductor"
pseudo-model that orchestrates across the whole Morph family instead
of answering from a single tier (see orchestration/conductor.py).

Run with:
    python scripts/serve.py --family classic --model ./path/to/checkpoint
    python scripts/serve.py --family nano                    # untrained, for testing wiring
"""
from __future__ import annotations

import json
import os
import re
import time
import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse

from api.schemas import (
    ChatCompletionChunk, ChatCompletionChunkChoice, ChatCompletionChunkDelta,
    ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse,
    ChatMessage, ErrorDetail, ErrorResponse, HealthResponse,
    ModelInfo, ModelListResponse, ToolCall, ToolCallFunction,
    ToolExecuteRequest, Usage,
)
from utils.common import load_family_config, FAMILY_TIERS, get_device, describe_device


# ── App state ──────────────────────────────────────────────────────

class ServerState:
    engine: Optional[object] = None      # MorphInferenceEngine, set at startup
    family: str = "classic"
    model_name: str = "Morph Classic"
    memory_store: Optional[object] = None
    knowledge_base: Optional[object] = None
    search_client: Optional[object] = None
    conductor: Optional[object] = None   # MorphConductor, set at startup if enabled


state = ServerState()


def _remote_worker_call_fn(base_url: str):
    """
    Build a call_fn for a WorkerSpec that hits a *separate* serve.py
    instance (a different tier, running as its own process — e.g. on
    another Kaggle/Colab session, or a second GPU). Only used if you
    actually have the hardware to run more than one tier at once;
    otherwise Conductor falls back to self-orchestration on whichever
    single engine this process loaded.
    """
    import requests

    def _call(system_prompt: str, user_content: str) -> str:
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": user_content})
        resp = requests.post(
            f"{base_url.rstrip('/')}/chat/completions",
            json={"model": "classic", "messages": messages, "stream": False, "max_tokens": 1024},
            timeout=120,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"] or ""

    return _call


def _load_conductor():
    """
    Build the Conductor's worker pool.

    For each Morph tier name, check for MORPH_WORKER_<NAME>_URL — if
    set, that worker calls a separately-running serve.py instance over
    HTTP (true multi-tier orchestration, needs the hardware to run
    more than one model at once). If not set, fall back to
    self-orchestration: the one engine this process already loaded,
    given a different role-prompt per worker name. Self-orchestration
    is the realistic default for a single Kaggle/Colab GPU — it still
    gets you task decomposition, specialization-by-prompt, and the
    critic/refine loop, just not genuinely different model weights
    per worker.
    """
    from orchestration.conductor import MorphConductor, WorkerSpec, DEFAULT_ROLE_PROMPTS

    if os.environ.get("MORPH_CONDUCTOR_ENABLED", "true").lower() == "false":
        return

    worker_names = ["nano", "mini", "classic", "pro", "code", "critic"]
    workers = []
    for name in worker_names:
        remote_url = os.environ.get(f"MORPH_WORKER_{name.upper()}_URL")
        if remote_url:
            workers.append(WorkerSpec(
                name=name, description=_worker_description(name),
                call_fn=_remote_worker_call_fn(remote_url),
                role_prompt=DEFAULT_ROLE_PROMPTS.get(name),
            ))
        elif name != state.family:   # don't add a self-orchestration worker identical to the base engine
            workers.append(WorkerSpec(
                name=name, description=_worker_description(name),
                engine=state.engine, role_prompt=DEFAULT_ROLE_PROMPTS.get(name),
            ))

    mode = os.environ.get("MORPH_CONDUCTOR_MODE", "auto")
    state.conductor = MorphConductor(
        orchestrator=state.engine, workers=workers, default_mode=mode,
    )
    any_remote = any(os.environ.get(f"MORPH_WORKER_{n.upper()}_URL") for n in worker_names)
    print(f"✓ Conductor ready ({len(workers)} workers, "
          f"{'remote tiers configured' if any_remote else 'self-orchestration mode'}, default_mode={mode})")


def _worker_description(name: str) -> str:
    return {
        "nano":    "Fastest tier — best for simple, quick questions.",
        "mini":    "Fast general-purpose tier for everyday questions.",
        "classic": "Balanced tier — general chat, tool use, most tasks.",
        "pro":     "Highest-quality tier — hard multi-step reasoning, verification.",
        "code":    "Code-specialized — programming, debugging, repo-level tasks.",
        "critic":  "Reviews another worker's output for correctness before it's returned.",
    }.get(name, name)


def _load_engine():
    """
    Load the model once at startup. Reads configuration from
    environment variables so `scripts/serve.py` (or a container's env)
    controls what gets loaded without editing this file.
    """
    from multimodal.fusion.morph_multimodal import MorphMultimodalModel
    from tokenizer.morph_tokenizer import build_morph_tokenizer
    from inference.engine.morph_engine import MorphInferenceEngine

    family = os.environ.get("MORPH_FAMILY", "classic")
    checkpoint_path = os.environ.get("MORPH_CHECKPOINT")
    quantization = os.environ.get("MORPH_QUANTIZATION") or None

    state.family = family
    config, module = load_family_config(family)
    state.model_name = getattr(module, "MODEL_NAME", family.title())

    if checkpoint_path:
        print(f"Loading {state.model_name} from checkpoint: {checkpoint_path}")
        engine = MorphInferenceEngine.from_pretrained(
            checkpoint_path, quantization=quantization,
        )
    else:
        print(f"⚠ No MORPH_CHECKPOINT set — building an UNTRAINED {state.model_name}. "
              f"Responses will be gibberish. Set MORPH_CHECKPOINT to a real checkpoint dir.")
        tokenizer = build_morph_tokenizer()
        model = MorphMultimodalModel(config)
        device = get_device()
        engine = MorphInferenceEngine(model=model, tokenizer=tokenizer, config=config, device=device)

    state.engine = engine

    # Tool dependencies (memory/knowledge-base are self-contained local stores)
    try:
        import tools  # registers every @tool
        from tools.memory import MemoryStore
        from tools.knowledge_base import KnowledgeBase
        state.memory_store = MemoryStore(path=os.environ.get("MORPH_MEMORY_PATH", "./data/memories.json"))
        state.knowledge_base = KnowledgeBase()
    except ImportError as e:
        print(f"⚠ Tools unavailable ({e}) — /v1/tools/execute will return errors")

    _load_conductor()


@asynccontextmanager
async def lifespan(app: FastAPI):
    _load_engine()
    yield
    # (no teardown needed — process exit frees everything)


app = FastAPI(title="Frox AI — Morph API", version="1.1.0", lifespan=lifespan)

# Permissive CORS for local dev (Tauri apps, localhost web dev servers).
# Restrict allow_origins to your actual client origins in production.
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("MORPH_CORS_ORIGINS", "*").split(","),
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Auth ────────────────────────────────────────────────────────────
# This server had no auth check at all: any client that could reach the
# HTTP port could run inference, regardless of whether it sent a key.
# That's fine on localhost, but the moment this is tunneled out (ngrok,
# a VPS reverse proxy, etc.) the URL is effectively public — anyone who
# finds it can run your GPU for free. If MORPH_API_KEY is set, require a
# matching `Authorization: Bearer <key>` header on every request except
# /health (so uptime checks / ngrok's own health probe still work).
_MORPH_API_KEY = os.environ.get("MORPH_API_KEY", "")
if not _MORPH_API_KEY:
    print("⚠️  MORPH_API_KEY is not set — this server accepts requests from anyone who can "
          "reach it. Fine for pure localhost testing; set MORPH_API_KEY before exposing it "
          "via ngrok or a public URL.")


@app.middleware("http")
async def require_api_key(request: Request, call_next):
    if _MORPH_API_KEY and request.url.path != "/health":
        auth = request.headers.get("authorization", "")
        token = auth[7:] if auth.lower().startswith("bearer ") else ""
        if token != _MORPH_API_KEY:
            return JSONResponse(
                status_code=401,
                content=ErrorResponse(error=ErrorDetail(
                    message="Invalid or missing API key.", type="invalid_request_error",
                )).model_dump(),
            )
    return await call_next(request)


# ── Error handling (OpenAI-shaped errors) ──────────────────────────

@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content=ErrorResponse(error=ErrorDetail(
            message=str(exc), type="server_error",
        )).model_dump(),
    )


# ── Health ────────────────────────────────────────────────────────

@app.get("/health", response_model=HealthResponse)
async def health():
    if state.engine is None:
        return HealthResponse(status="loading")
    return HealthResponse(
        status="ok", model=state.model_name, device=str(state.engine.device),
    )


# ── Models ────────────────────────────────────────────────────────

@app.get("/v1/models", response_model=ModelListResponse)
async def list_models():
    """
    Returns the model(s) this server instance can serve: the loaded
    base tier, plus "conductor" if orchestration is enabled — clients
    that list models (Open WebUI's dropdown, for instance) can offer
    both without any code change on their end.
    """
    models = [ModelInfo(id=state.family)]
    if state.conductor is not None:
        models.append(ModelInfo(id="conductor"))
    return ModelListResponse(data=models)


# ── Chat completions ─────────────────────────────────────────────

def _messages_to_dicts(messages: list[ChatMessage]) -> list[dict]:
    out = []
    for m in messages:
        content = m.content or ""
        if m.role == "tool":
            # Wrap in Morph's dedicated tool-result tags rather than passing
            # raw text under a generic "tool" role header — matches what the
            # tokenizer/training data actually teach the model to expect.
            content = f"<|tool_result|>{content}<|/tool_result|>"
        out.append({"role": m.role, "content": content})
    return out


def _extract_system_and_last_user(messages: list[ChatMessage]) -> tuple[Optional[str], str]:
    system = next((m.content for m in messages if m.role == "system"), None)
    last_user = next((m.content for m in reversed(messages) if m.role == "user"), "")
    return system, (last_user or "")


# ── Tool-calling bridge ──────────────────────────────────────────
#
# frox-code (and any OpenAI-style client) sends `tools` as JSON function
# schemas and expects `tool_calls` back on the response message. Morph
# itself was designed around a simpler native format — a
# <|tool_call|>{"name":...,"args":...}<|/tool_call|> block in the raw
# generated text (see tokenizer/morph_tokenizer.py and
# MorphInferenceEngine.parse_tool_calls()). This section bridges the two:
# describe the OpenAI tool schemas to the model as a system-prompt
# addendum instructing it to answer in Morph's native format, then parse
# that format back out of the generated text into real `tool_calls`.
#
# Note: this makes the WIRING correct. Whether the model reliably
# produces well-formed tool calls also depends on it having been
# fine-tuned on examples that do so — the training pipeline already
# reserves the right special tokens and SFT format for this (see
# training/pipeline/trainer.py), but wiring and training are separate
# concerns. An untrained/base checkpoint won't use this reliably.

_TOOL_CALL_OPEN = "<|tool_call|>"
_TOOL_CALL_CLOSE = "<|/tool_call|>"
_TOOL_CALL_BLOCK_RE = re.compile(
    re.escape(_TOOL_CALL_OPEN) + r"(.*?)" + re.escape(_TOOL_CALL_CLOSE), re.DOTALL,
)


def _build_tools_system_addendum(tools: list[dict]) -> str:
    lines = [
        "You have access to the following tools. When you need to use one, "
        "respond with exactly this format and nothing else in that turn:",
        f'{_TOOL_CALL_OPEN}{{"name": "<tool_name>", "args": {{<arguments as JSON>}}}}{_TOOL_CALL_CLOSE}',
        "You can emit more than one such block to call multiple tools in one turn. "
        "Wait for each tool's result before continuing your response.",
        "",
        "Available tools:",
    ]
    for t in tools:
        fn = t.get("function", t)   # tolerate either {"type":"function","function":{...}} or a flat dict
        name = fn.get("name", "unknown")
        desc = fn.get("description", "")
        params = fn.get("parameters", {})
        lines.append(f"- {name}: {desc}\n  parameters: {json.dumps(params)}")
    return "\n".join(lines)


def _inject_tools_system_message(messages: list[dict], tools: list[dict]) -> list[dict]:
    addendum = _build_tools_system_addendum(tools)
    messages = list(messages)
    if messages and messages[0]["role"] == "system":
        messages[0] = {**messages[0], "content": messages[0]["content"] + "\n\n" + addendum}
    else:
        messages = [{"role": "system", "content": addendum}] + messages
    return messages


def _strip_tool_call_tags(text: str) -> str:
    return _TOOL_CALL_BLOCK_RE.sub("", text).strip()


def _to_openai_tool_calls(parsed: list[dict]) -> list[ToolCall]:
    return [
        ToolCall(
            id=f"call_{uuid.uuid4().hex[:24]}",
            function=ToolCallFunction(
                name=p.get("name", ""),
                arguments=json.dumps(p.get("args", {})),
            ),
        )
        for p in parsed
    ]


def _safe_flush_length(buffer: str, tag: str) -> int:
    """
    How many characters of `buffer` are safe to emit immediately — i.e.
    everything except a trailing suffix that could still be the start of
    `tag` once more text streams in. Without this, a tag split across two
    streamed chunks (e.g. one chunk ending in a lone "<") would leak that
    fragment to the client before we know whether it's actually a tag.
    """
    max_check = min(len(tag) - 1, len(buffer))
    for length in range(max_check, 0, -1):
        if tag.startswith(buffer[-length:]):
            return len(buffer) - length
    return len(buffer)


@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    if state.engine is None:
        raise HTTPException(status_code=503, detail="Model still loading")

    if request.model in ("conductor", "auto"):
        if state.conductor is None:
            raise HTTPException(
                status_code=503,
                detail="Conductor isn't enabled on this server (MORPH_CONDUCTOR_ENABLED=false)",
            )
        return await _conductor_chat_completion(request)

    if request.model not in FAMILY_TIERS and request.model != state.family:
        # Not a hard error — this server only ever serves the one loaded
        # family, so just note the mismatch and continue serving it.
        pass

    if request.stream:
        return StreamingResponse(
            _stream_chat_completion(request),
            media_type="text/event-stream",
        )
    return await _full_chat_completion(request)


async def _conductor_chat_completion(request: ChatCompletionRequest):
    """
    Dispatch through Morph Conductor. Conductor's run()/run_recursive()
    are synchronous and internally make several blocking GPU calls (one
    per workflow step) — always thread-offloaded, same reasoning as
    every other blocking call in this file.

    Streaming note: Conductor doesn't have a true token-by-token
    streaming mode (each step is a complete generate() call, not a
    generator) — a streamed request still gets the full orchestrated
    result, just delivered as a sequence of word-sized chunks so the
    client still sees progressive output rather than one long pause.
    """
    import asyncio

    conductor = state.conductor
    messages = _messages_to_dicts(request.messages)
    use_recursive = request.session_id == "conductor-recursive"   # opt-in via a sentinel session_id

    if use_recursive:
        result = await asyncio.to_thread(conductor.run_recursive, messages)
    else:
        result = await asyncio.to_thread(conductor.run, messages)

    if not request.stream:
        prompt_text = " ".join(m.content or "" for m in request.messages)
        prompt_tokens = len(state.engine.tokenizer.encode(prompt_text, add_special_tokens=False))
        completion_tokens = len(state.engine.tokenizer.encode(result.text, add_special_tokens=False))
        return ChatCompletionResponse(
            model="conductor",
            choices=[ChatCompletionChoice(
                message=ChatMessage(role="assistant", content=result.text),
                finish_reason="stop",
            )],
            usage=Usage(
                prompt_tokens=prompt_tokens, completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            ),
            frox_trace=result.trace,
        )

    async def _fake_stream():
        completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"

        def _chunk(content=None, role=None, finish_reason=None) -> str:
            chunk = ChatCompletionChunk(
                id=completion_id, model="conductor",
                choices=[ChatCompletionChunkChoice(
                    delta=ChatCompletionChunkDelta(role=role, content=content),
                    finish_reason=finish_reason,
                )],
            )
            return f"data: {chunk.model_dump_json()}\n\n"

        yield _chunk(role="assistant")
        words = result.text.split(" ")
        for i, word in enumerate(words):
            yield _chunk(content=word + (" " if i < len(words) - 1 else ""))
        yield _chunk(finish_reason="stop")
        yield "data: [DONE]\n\n"

    return StreamingResponse(_fake_stream(), media_type="text/event-stream")


async def _full_chat_completion(request: ChatCompletionRequest) -> ChatCompletionResponse:
    import asyncio
    engine = state.engine

    messages = _messages_to_dicts(request.messages)
    if request.tools:
        messages = _inject_tools_system_message(messages, request.tools)

    # engine.chat()/generate() are synchronous and GPU-bound. Calling them
    # directly here would block the asyncio event loop for the entire
    # generation — freezing every other connected client's request until
    # this one finishes. asyncio.to_thread() runs it in a worker thread
    # instead, keeping the event loop free to serve concurrent requests.
    if request.session_id and not request.tools:
        # Tools + session_id together aren't supported: chat() only takes
        # this turn's new user message, not the full list, so a
        # tools-addendum built from THIS request can't be reliably baked
        # into an already-cached system prompt from an earlier turn.
        # Falling back to non-session generation whenever tools are used
        # keeps behavior correct rather than silently wrong.
        system, last_user = _extract_system_and_last_user(request.messages)
        response_text = await asyncio.to_thread(
            engine.chat, request.session_id, last_user,
            system_prompt=system, max_new_tokens=request.max_tokens,
            temperature=request.temperature,
        )
    else:
        response_text = await asyncio.to_thread(
            engine.generate, messages, max_new_tokens=request.max_tokens,
            temperature=request.temperature, top_p=request.top_p,
        )

    tool_calls = None
    finish_reason = "stop"
    if request.tools:
        parsed = engine.parse_tool_calls(response_text)
        if parsed:
            tool_calls = _to_openai_tool_calls(parsed)
            finish_reason = "tool_calls"
            response_text = _strip_tool_call_tags(response_text)

    prompt_text = " ".join(m.content or "" for m in request.messages)
    prompt_tokens = len(engine.tokenizer.encode(prompt_text, add_special_tokens=False))
    completion_tokens = len(engine.tokenizer.encode(response_text, add_special_tokens=False))

    return ChatCompletionResponse(
        model=state.family,
        choices=[ChatCompletionChoice(
            message=ChatMessage(
                role="assistant",
                content=response_text if response_text else None,
                tool_calls=tool_calls,
            ),
            finish_reason=finish_reason,
        )],
        usage=Usage(
            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        ),
    )


async def _stream_chat_completion(request: ChatCompletionRequest) -> AsyncGenerator[str, None]:
    import asyncio
    import threading

    engine = state.engine
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"

    def _chunk(content: Optional[str] = None, role: Optional[str] = None,
              finish_reason: Optional[str] = None, tool_calls=None) -> str:
        chunk = ChatCompletionChunk(
            id=completion_id, model=state.family,
            choices=[ChatCompletionChunkChoice(
                delta=ChatCompletionChunkDelta(role=role, content=content, tool_calls=tool_calls),
                finish_reason=finish_reason,
            )],
        )
        return f"data: {chunk.model_dump_json()}\n\n"

    yield _chunk(role="assistant")

    messages = _messages_to_dicts(request.messages)
    if request.tools:
        messages = _inject_tools_system_message(messages, request.tools)

    # engine.chat()/generate_stream() are synchronous generators doing
    # GPU-bound work between each yielded token. Draining them directly
    # in this async generator would block the event loop for the whole
    # response, freezing every other connected client in the meantime.
    # Instead: run the blocking work in a background thread, and bridge
    # each token back to this async generator through an asyncio.Queue
    # via call_soon_threadsafe — the standard pattern for wrapping a
    # sync producer with an async consumer without blocking the loop.
    loop = asyncio.get_event_loop()
    out_queue: asyncio.Queue = asyncio.Queue()
    _DONE = object()

    def _worker():
        try:
            if request.session_id and not request.tools:
                system, last_user = _extract_system_and_last_user(request.messages)
                engine.chat(
                    request.session_id, last_user, system_prompt=system,
                    max_new_tokens=request.max_tokens, temperature=request.temperature,
                    stream_callback=lambda delta: loop.call_soon_threadsafe(
                        out_queue.put_nowait, delta
                    ),
                )
            else:
                for delta in engine.generate_stream(
                    messages, max_new_tokens=request.max_tokens,
                    temperature=request.temperature, top_p=request.top_p,
                ):
                    loop.call_soon_threadsafe(out_queue.put_nowait, delta)
        except Exception as e:
            loop.call_soon_threadsafe(out_queue.put_nowait, RuntimeError(str(e)))
        finally:
            loop.call_soon_threadsafe(out_queue.put_nowait, _DONE)

    threading.Thread(target=_worker, daemon=True).start()

    # Buffer text so a <|tool_call|>...<|/tool_call|> block is never
    # partially shown to the client as raw tag markup. Text is only ever
    # withheld for (a) the characters currently inside an open tool-call
    # block, or (b) a trailing suffix that could still turn into the
    # opening tag once more text arrives (see _safe_flush_length) — never
    # withheld indefinitely, and never more than len(tag)-1 characters
    # in case (b).
    buffer = ""
    in_tool_call = False
    full_text_parts: list[str] = []
    saw_error: Optional[RuntimeError] = None

    while True:
        item = await out_queue.get()
        if item is _DONE:
            break
        if isinstance(item, RuntimeError):
            saw_error = item
            break

        buffer += item
        full_text_parts.append(item)

        while True:
            if not in_tool_call:
                idx = buffer.find(_TOOL_CALL_OPEN)
                if idx == -1:
                    safe_len = _safe_flush_length(buffer, _TOOL_CALL_OPEN)
                    if safe_len > 0:
                        yield _chunk(content=buffer[:safe_len])
                        buffer = buffer[safe_len:]
                    break
                if idx > 0:
                    yield _chunk(content=buffer[:idx])
                buffer = buffer[idx:]
                in_tool_call = True
                # loop again — the close tag might already be in this buffer too
            else:
                idx = buffer.find(_TOOL_CALL_CLOSE)
                if idx == -1:
                    break   # still waiting for the rest of the tool call
                buffer = buffer[idx + len(_TOOL_CALL_CLOSE):]
                in_tool_call = False
                # loop again in case there's more real text after this

    if saw_error:
        yield _chunk(content=f"\n\n[error: {saw_error}]")
        yield _chunk(finish_reason="stop")
        yield "data: [DONE]\n\n"
        return

    if request.tools:
        full_text = "".join(full_text_parts)
        parsed = engine.parse_tool_calls(full_text)
        if parsed:
            yield _chunk(tool_calls=_to_openai_tool_calls(parsed))
            yield _chunk(finish_reason="tool_calls")
            yield "data: [DONE]\n\n"
            return

    yield _chunk(finish_reason="stop")
    yield "data: [DONE]\n\n"


# ── Tool execution ────────────────────────────────────────────────

@app.post("/v1/tools/execute")
async def execute_tool(request: ToolExecuteRequest):
    try:
        from tools.registry import registry, ToolContext
    except ImportError:
        raise HTTPException(status_code=503, detail="Tools package unavailable")

    ctx = ToolContext(
        engine=state.engine, memory_store=state.memory_store,
        knowledge_base=state.knowledge_base, user_id=request.user_id,
        session_id=request.session_id, plan=request.plan,
        extra={"search_client": state.search_client},
    )
    result = await registry.execute(request.tool, request.args, ctx)
    return result.to_dict()


@app.get("/v1/tools")
async def list_tools(plan: str = "free"):
    try:
        from tools.registry import registry
    except ImportError:
        return {"tools": []}
    return {"tools": registry.list_tools(plan=plan)}