Chat Handlers¶

from typing import List
from neurosurfer.server import NeurosurferApp
from neurosurfer.server.schemas import ChatCompletionRequest
from neurosurfer.server.runtime import RequestContext
from neurosurfer.config import config

app = NeurosurferApp()

# Optional: RAG orchestrator/service may be initialized in startup hooks
RAG = None
LLM = None

@app.on_startup
def boot():
    global LLM, RAG
    # Load your model(s), embedder, and RAG orchestrator here.
    # Add the model to app.model_registry for client discovery.
    ...

@app.chat()
def handle_chat(request: ChatCompletionRequest, ctx: RequestContext):
    # 1) Extract prompts
    user_msgs = [m["content"] for m in request.messages if m["role"] == "user"]
    system_msgs = [m["content"] for m in request.messages if m["role"] == "system"]
    user_prompt = user_msgs[-1] if user_msgs else ""
    system_prompt = system_msgs[-1] if system_msgs else config.model.system_prompt

    # 2) Optional: compact recent chat history
    chat_history = request.messages[-10:-1]

    # 3) Optional: RAG augmentation per thread
    if RAG and request.thread_id is not None:
        rag_res = RAG.apply(
            actor_id=(getattr(ctx, "meta", {}) or {}).get("actor_id", 0),
            thread_id=request.thread_id,
            user_query=user_prompt,
            files=[f.model_dump() for f in (request.files or [])]
        )
        user_prompt = rag_res.augmented_query  # falls back to original if not used

    # 4) Model call (stream or not)
    kwargs = {
        "temperature": request.temperature or config.model.temperature,
        "max_new_tokens": request.max_tokens or config.model.max_new_tokens,
        "stream": request.stream,
    }

    result = LLM.ask(
        user_prompt=user_prompt,
        system_prompt=system_prompt,
        chat_history=chat_history,
        **kwargs
    )

    # 5) Return directly for non-streaming
    if not request.stream:
        return result

    # 6) For streaming, wrap yielded tokens into OpenAI delta chunks if needed
    def to_stream(gen):
        for chunk in gen:
            # If your model already yields OpenAI-shaped deltas, just yield them.
            # Otherwise adapt each token to {"choices":[{"delta":{"content": token}}]}
            yield chunk
    return to_stream(result)

# Assume 'toolkit' and 'LLM' (and optionally 'EMBEDDER', 'RAG') were initialized in startup.
@app.chat()
def handler(request, ctx):
    tool_name = getattr(request, "tool", None)
    tool_input = getattr(request, "tool_input", {}) or {}

    if tool_name:
        tool = toolkit.registry.get(tool_name)
        if not tool:
            return {"error": {"code": "unknown_tool", "message": f"Tool '{tool_name}' not registered"}}

        try:
            data = tool.spec.check_inputs(tool_input)
        except Exception as e:
            return {"error": {"code": "invalid_tool_input", "message": str(e)}}

        out = tool(**data, llm=LLM, embedder=EMBEDDER, rag=RAG, stream=request.stream)
        obs = out.observation

        if request.stream and hasattr(obs, "__iter__") and not isinstance(obs, (str, bytes)):
            for token in obs:
                yield {"choices": [{"delta": {"content": str(token)}}]}
        else:
            return {"choices": [{"message": {"role": "assistant", "content": str(obs)}}]}

    # …otherwise fall back to your standard chat flow (RAG + LLM).
    ...

import json

@app.chat()
def handler(request, ctx):
    user_last = next((m["content"] for m in reversed(request.messages) if m["role"] == "user"), "")
    history = request.messages[-10:-1]

    menu = toolkit.get_tools_description()
    system_prompt = (
        "You are Neurosurfer. Decide if a tool is needed.\n\n"
        "TOOLS:\n" + menu + "\n\n"
        "INSTRUCTIONS:\n"
        "- If you can answer directly, output ONLY the final answer text.\n"
        "- If you need a tool, output ONLY this JSON (no backticks, no extra text):\n"
        '  {\"call_tool\": {\"name\": \"<tool_name>\", \"args\": { ... }}}\n'
        "- Use exact parameter names/types from the tool spec."
    )

    # First pass (non-streaming) to inspect if tool JSON is returned
    first = LLM.ask(
        system_prompt=system_prompt,
        user_prompt=user_last,
        chat_history=history,
        temperature=request.temperature or 0.7,
        max_new_tokens=request.max_tokens or 1024,
        stream=False,
    )

    text = first.choices[0].message.content or ""
    try:
        call = json.loads(text).get("call_tool")
    except Exception:
        call = None

    # No tool requested → return (or optionally re-ask in streaming mode)
    if not call:
        if request.stream:
            return LLM.ask(
                system_prompt=system_prompt,
                user_prompt=user_last,
                chat_history=history,
                temperature=request.temperature or 0.7,
                max_new_tokens=request.max_tokens or 1024,
                stream=True,
            )
        return first

    # Tool requested → validate and run
    name, args = call.get("name"), (call.get("args") or {})
    tool = toolkit.registry.get(name)
    if not tool:
        return {"error": {"code": "unknown_tool", "message": f"Tool '{name}' is not registered"}}

    try:
        data = tool.spec.check_inputs(args)
    except Exception as e:
        return {"error": {"code": "invalid_tool_input", "message": str(e)}}

    out = tool(**data, llm=LLM, embedder=EMBEDDER, rag=RAG, stream=request.stream)
    obs = out.observation

    if request.stream and hasattr(obs, "__iter__") and not isinstance(obs, (str, bytes)):
        for token in obs:
            yield {"choices": [{"delta": {"content": str(token)}}]}
    else:
        return {"choices": [{"message": {"role": "assistant", "content": str(obs)}}]}