localRAG/realtime_chat.py

"""
Text-only LangChain chat model for OpenAI's Realtime API (e.g. gpt-realtime-2.1-mini).

Realtime models are not served on Chat Completions or the Responses API, only on
wss://api.openai.com/v1/realtime. This wrapper opens one WebSocket session per call, replays the
messages as conversation items, asks for a text response and returns it as an AIMessage (with
tool_calls if the model calls a function). That way LocalRAG.query_with_history can use it like
ChatOpenAI or ChatOllama, including bind_tools().
"""
import json
import uuid

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.utils.function_calling import convert_to_openai_tool
from websockets.sync.client import connect

REALTIME_URL = "wss://api.openai.com/v1/realtime?model={model}"


def _text(content):
    """Plain text of a LangChain message content (string or list of blocks)."""
    if isinstance(content, list):
        return "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content)
    return content or ""


class ChatOpenAIRealtime(BaseChatModel):
    model: str
    api_key: str
    timeout: float = 120.0

    @property
    def _llm_type(self):
        return "openai-realtime"

    def bind_tools(self, tools, **kwargs):
        """Tools in the Realtime format: {"type": "function", "name", "description", "parameters"}."""
        specs = []
        for tool in tools:
            fn = convert_to_openai_tool(tool)["function"]
            specs.append({"type": "function", "name": fn["name"],
                          "description": fn.get("description", ""), "parameters": fn["parameters"]})
        return self.bind(tools=specs, **kwargs)

    @staticmethod
    def _items(messages):
        """System prompt and conversation items for the Realtime session."""
        instructions, items = [], []
        for m in messages:
            if isinstance(m, SystemMessage):
                instructions.append(_text(m.content))
            elif isinstance(m, HumanMessage):
                items.append({"type": "message", "role": "user",
                              "content": [{"type": "input_text", "text": _text(m.content)}]})
            elif isinstance(m, AIMessage):
                if _text(m.content).strip():
                    items.append({"type": "message", "role": "assistant",
                                  "content": [{"type": "output_text", "text": _text(m.content)}]})
                for call in m.tool_calls:
                    items.append({"type": "function_call", "call_id": call["id"], "name": call["name"],
                                  "arguments": json.dumps(call["args"], ensure_ascii=False)})
            elif isinstance(m, ToolMessage):
                items.append({"type": "function_call_output", "call_id": m.tool_call_id,
                              "output": _text(m.content)})
        return "\n\n".join(instructions), items

    def _generate(self, messages, stop=None, run_manager=None, tools=None, **kwargs):
        instructions, items = self._items(messages)
        session = {"type": "realtime", "output_modalities": ["text"], "instructions": instructions}
        if tools:
            session.update(tools=tools, tool_choice="auto")

        with connect(REALTIME_URL.format(model=self.model),
                     additional_headers={"Authorization": f"Bearer {self.api_key}"},
                     open_timeout=30, max_size=None) as ws:
            def send(event):
                ws.send(json.dumps(event, ensure_ascii=False))

            send({"type": "session.update", "session": session})
            for item in items:
                send({"type": "conversation.item.create", "item": item})
            send({"type": "response.create"})

            while True:
                event = json.loads(ws.recv(timeout=self.timeout))
                if event["type"] == "error":
                    raise RuntimeError(f"Realtime API error: {event.get('error')}")
                if event["type"] == "response.done":
                    response = event["response"]
                    break

        if response.get("status") not in (None, "completed"):
            raise RuntimeError(f"Realtime response {response.get('status')}: "
                               f"{response.get('status_details')}")
        text, tool_calls = [], []
        for item in response.get("output", []):
            if item.get("type") == "message":
                text += [part.get("text", "") for part in item.get("content", [])
                         if part.get("type") in ("output_text", "text")]
            elif item.get("type") == "function_call":
                tool_calls.append({"id": item.get("call_id") or f"call_{uuid.uuid4().hex[:8]}",
                                   "name": item["name"], "args": json.loads(item.get("arguments") or "{}")})
        message = AIMessage(content="".join(text), tool_calls=tool_calls,
                            response_metadata={"model": self.model, "usage": response.get("usage")})
        return ChatResult(generations=[ChatGeneration(message=message)])