Open WebUI

"""
title: AgentOS
description: Chat with an agent that runs on AgentOS. Open WebUI stays the chat UI; AgentOS runs, governs and logs the agent.
version: 0.1.0
"""

import json

import httpx
from pydantic import BaseModel, Field


def _text(content) -> str:
    """Open WebUI message content: a string, or a list of parts (text, images) for multimodal messages."""
    if isinstance(content, str):
        return content
    return "".join(part.get("text", "") for part in content or [] if isinstance(part, dict))


def _problem(response: httpx.Response) -> str:
    try:
        error = response.json()
        return f"AgentOS: {error.get('message') or error.get('code')} (HTTP {response.status_code})"
    except ValueError:
        return f"AgentOS: HTTP {response.status_code}"


class Pipe:
    class Valves(BaseModel):
        AGENTOS_URL: str = Field(default="https://agentos.example.com", description="Your AgentOS address")
        AGENTOS_API_KEY: str = Field(default="", description="A service account's API key (aos_…), member role")
        ENVIRONMENT: str = Field(default="default", description="The environment the agent is deployed to")
        AGENT: str = Field(default="webui-chat", description="The agent's name")
        CAPABILITY: str = Field(default="chat", description="The capability to run for each message")
        STREAM: bool = Field(default=True, description="Show the answer as the agent writes it (passthrough)")

    def __init__(self):
        self.valves = self.Valves()
        self.sessions: dict[str, str] = {}   # Open WebUI chat id → AgentOS session id

    async def pipe(self, body: dict, __user__: dict, __chat_id__: str | None = None, __task__: str | None = None):
        if __task__:
            return   # titles, tags and follow-up suggestions: don't start an agent run for them

        v = self.valves
        base = f"{v.AGENTOS_URL.rstrip('/')}/v0/envs/{v.ENVIRONMENT}"
        messages = [{"role": m["role"], "content": _text(m.get("content"))}
                    for m in body.get("messages", []) if m.get("role") in ("user", "assistant")]
        url = f"{base}/agents/{v.AGENT}/capabilities/{v.CAPABILITY}/runs"

        async with httpx.AsyncClient(headers={"Authorization": f"Bearer {v.AGENTOS_API_KEY}"},
                                     timeout=httpx.Timeout(10, read=330)) as client:
            run = {"input": {"messages": messages}}
            if __chat_id__:
                session_id = await self._session(client, base, __chat_id__, __user__)
                if session_id is None:
                    yield "AgentOS: couldn't open a session for this chat."
                    return
                run["session_id"] = session_id

            if body.get("stream") and v.STREAM:
                # The agent's own stream, as it writes it (plain text from the example agent).
                async with client.stream("POST", url, params={"passthrough": "true"}, json=run) as response:
                    if response.status_code >= 400:
                        await response.aread()
                        yield _problem(response)
                        return
                    if response.headers.get("content-type", "").startswith("application/json"):
                        await response.aread()   # the agent answered without streaming, or the run is held
                        yield self._answer(response.json())
                        return
                    async for text in response.aiter_text():
                        yield text
            else:
                response = await client.post(url, params={"wait": "true"}, json=run)
                yield _problem(response) if response.status_code >= 400 else self._answer(response.json())

    async def _session(self, client: httpx.AsyncClient, base: str, chat_id: str, user: dict) -> str | None:
        """One AgentOS session per Open WebUI chat, naming the Open WebUI user as the session's end user."""
        if chat_id not in self.sessions:
            response = await client.post(f"{base}/sessions", json={"user": user["id"]} if user.get("id") else {})
            if response.status_code >= 400:
                return None
            self.sessions[chat_id] = response.json()["session_id"]
        return self.sessions[chat_id]

    @staticmethod
    def _answer(run: dict) -> str:
        if run.get("status") == "completed":
            output = run.get("output") or {}
            return output.get("answer") if isinstance(output, dict) and "answer" in output else json.dumps(output)
        if run.get("status") == "held":
            return f"Waiting for a person's approval in AgentOS (run {run['run_id']})."
        error = run.get("error") or {}
        return f"AgentOS: the run {run.get('status')} ({error.get('code', 'no code')}: {error.get('message', '')})"