app/engine.py

"""Generic task runner: a tool-using chat loop with scoring and LangFuse tracing.

Each task runs against a backend (config.Backend) resolved from its `backend`
name — so different demos can target different endpoints (Anthropic-direct, an
Oumi proxy, ...) in the same app. Two wire protocols per the backend's provider:
- "openai": OpenAI Chat Completions via the langfuse.openai wrapper (calls are
  auto-traced as generations).
- "anthropic": Anthropic Messages (for Oumi deployments that proxy an Anthropic
  model). Tasks still define tools in OpenAI format; we translate at call time
  and record each call as a generation manually.

Either path returns (final_text, transcript, messages); everything downstream —
parsing, scoring, tracing attributes, the return shape — is provider-agnostic.
`messages` is always the complete OpenAI-format trajectory (system → final
answer), including on the anthropic path, so the trajectory artifact the UI
copies/downloads is one standard chat record whatever backend produced it.
"""

import json

from langfuse import get_client, observe, propagate_attributes
from langfuse.openai import OpenAI

from app import config
from app.config import Backend
from app.tasks.base import Task

# Clients are cached per backend identity so repeated calls reuse connections.
_clients: dict[tuple, object] = {}


def _openai(be: Backend):
    key = ("openai", be.base_url, be.api_key)
    if key not in _clients:
        _clients[key] = OpenAI(base_url=be.base_url, api_key=be.api_key)
    return _clients[key]


def _anthropic(be: Backend):
    key = ("anthropic", be.base_url, be.api_key)
    if key not in _clients:
        import anthropic

        _clients[key] = anthropic.Anthropic(base_url=be.base_url, api_key=be.api_key)
    return _clients[key]


def _split_system(messages: list[dict]) -> tuple[str, list[dict]]:
    """Separate the system text from the conversation (Anthropic takes it apart)."""
    system = "\n\n".join(m["content"] for m in messages if m["role"] == "system")
    convo = [m for m in messages if m["role"] != "system"]
    return system, convo


def _run_openai(task: Task, text: str, be: Backend) -> tuple[str, list[dict], list[dict]]:
    messages = task.build_messages(text)
    transcript: list[dict] = []
    for _ in range(config.MAX_TOOL_ITERATIONS):
        response = _openai(be).chat.completions.create(
            model=be.model,
            max_tokens=4096,
            messages=messages,
            **({"tools": task.tools} if task.tools else {}),
        )
        message = response.choices[0].message
        if not message.tool_calls:
            final_text = message.content or ""
            messages.append({"role": "assistant", "content": final_text})
            return final_text, transcript, messages
        messages.append({
            "role": "assistant",
            "content": message.content,
            "tool_calls": [tc.model_dump() for tc in message.tool_calls],
        })
        for tc in message.tool_calls:
            try:
                args = json.loads(tc.function.arguments)
            except json.JSONDecodeError:
                args = {}
            result = task.execute_tool(tc.function.name, args)
            transcript.append({"tool": tc.function.name, "args": args, "result": result})
            messages.append({"role": "tool", "tool_call_id": tc.id, "content": result})
    return "(stopped: tool iteration limit reached)", transcript, messages


def _to_anthropic_tools(tools: list[dict]) -> list[dict]:
    """OpenAI function-tool schema -> Anthropic tool schema."""
    return [
        {
            "name": t["function"]["name"],
            "description": t["function"].get("description", ""),
            "input_schema": t["function"]["parameters"],
        }
        for t in tools
    ]


def _runner_for(be: Backend):
    """Pick the wire-protocol loop for a backend."""
    return _run_anthropic if be.provider == "anthropic" else _run_openai


def _run_anthropic(task: Task, text: str, be: Backend) -> tuple[str, list[dict], list[dict]]:
    # Two views of the same conversation: `messages` is what we send (Anthropic
    # native, system split out); `oai` is the OpenAI-format trajectory we return.
    oai = task.build_messages(text)
    system, messages = _split_system(oai)
    tools = _to_anthropic_tools(task.tools)
    transcript: list[dict] = []
    langfuse = get_client()
    for _ in range(config.MAX_TOOL_ITERATIONS):
        with langfuse.start_as_current_observation(
            name="anthropic-messages", as_type="generation",
            model=be.model, input=messages,
        ) as gen:
            resp = _anthropic(be).messages.create(
                model=be.model,
                max_tokens=4096,
                system=system,
                messages=messages,
                **({"tools": tools} if tools else {}),
            )
            gen.update(
                output=[b.model_dump() for b in resp.content],
                usage_details={
                    "input": resp.usage.input_tokens,
                    "output": resp.usage.output_tokens,
                },
            )
        said = "".join(b.text for b in resp.content if b.type == "text")
        if resp.stop_reason != "tool_use":
            oai.append({"role": "assistant", "content": said})
            return said, transcript, oai
        messages.append({"role": "assistant", "content": resp.content})
        calls = [b for b in resp.content if b.type == "tool_use"]
        oai.append({
            "role": "assistant",
            "content": said or None,
            "tool_calls": [
                {
                    "id": b.id,
                    "type": "function",
                    "function": {"name": b.name, "arguments": json.dumps(b.input)},
                }
                for b in calls
            ],
        })
        results = []
        for b in calls:
            result = task.execute_tool(b.name, b.input)
            transcript.append({"tool": b.name, "args": b.input, "result": result})
            results.append({"type": "tool_result", "tool_use_id": b.id, "content": result})
            oai.append({"role": "tool", "tool_call_id": b.id, "content": result})
        messages.append({"role": "user", "content": results})
    return "(stopped: tool iteration limit reached)", transcript, oai


@observe(name="contimp-app-run", capture_input=False, capture_output=False)
def run_task(
    task: Task,
    text: str,
    input_id: str | None,
    user_id: str,
    session_id: str,
    source: str = "human",
) -> dict:
    truth = task.lookup_truth(input_id) if input_id else None
    be = config.backend(task.backend)
    langfuse = get_client()

    with propagate_attributes(
        trace_name=f"{task.id}",
        user_id=user_id or "anonymous",
        session_id=session_id,
        tags=[task.id, source],
        metadata={
            "task_id": task.id,
            "app_version": config.APP_VERSION,
            "backend": be.name,
            "model": be.model,
            "input_id": input_id or "",
            "ground_truth": json.dumps(truth) if truth else "",
        },
    ):
        final_text, transcript, messages = _runner_for(be)(task, text, be)

        parsed = task.parse_output(final_text)
        scores = task.score(truth, parsed) if truth else {}
        for name, value in scores.items():
            langfuse.score_current_trace(name=name, value=value)
        if task.tools:
            langfuse.score_current_trace(
                name="tool_iterations", value=float(len(transcript))
            )

        langfuse.update_current_span(
            input=text, output=final_text, metadata={"ground_truth": truth}
        )
        trace_id = langfuse.get_current_trace_id()

    return {
        "task_id": task.id,
        "output": task.present(parsed, truth),
        "raw_output": final_text,
        "transcript": transcript,
        "messages": messages,  # complete OpenAI-format trajectory (system → final answer)
        "truth": truth,
        "scores": scores,
        "trace_id": trace_id,
    }


def record_feedback(trace_id: str, value: int | None, comment: str | None = None) -> None:
    """Thumbs (numeric) when value is given; free-text note (TEXT score) otherwise."""
    if value is not None:
        get_client().create_score(
            trace_id=trace_id,
            name="thumbs",
            value=float(value),
            data_type="NUMERIC",
            comment=comment,
        )
    elif comment:
        get_client().create_score(
            trace_id=trace_id,
            name="feedback_note",
            value=comment,
            data_type="TEXT",
        )