VoiceAgent/ollama_llm.py

"""A Pipecat processor that puts Ollama / OpenCode / Ollama Cloud in the LLM slot.

Drives Ollama API (local or remote/cloud) with streaming responses so text-to-speech
downstream starts speaking immediately as tokens arrive.
"""

import asyncio
import json
import os
import re
from pathlib import Path
import aiohttp
from loguru import logger

from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    Frame,
    InterruptionFrame,
    LLMContextFrame,
    LLMFullResponseEndFrame,
    LLMFullResponseStartFrame,
    LLMTextFrame,
    StartFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor

_NOISE_TRANSCRIPTS = {
    "",
    ".",
    "thank you.",
    "thanks for watching!",
    "you",
    "bye.",
    "okay.",
    "[blank_audio]",
    "[silence]",
}


def _clean_spoken_text(text: str) -> str:
    """Clean text for speech output and truncate fake turn generations."""
    if not text:
        return ""
    # Truncate if model hallucinates fake turn markers
    for marker in ("User:", "Human:", "Assistant:", "\nUser", "\nHuman", "\nAssistant"):
        if marker in text:
            text = text.split(marker)[0]
    # Remove markdown code blocks
    text = re.sub(r"```[\s\S]*?```", "", text)
    # Remove inline code ticks
    text = re.sub(r"`[^`]*`", "", text)
    # Remove markdown syntax characters
    text = re.sub(r"[\#\*\_\~]", "", text)
    # Flatten newlines into clear speech
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    return " ".join(lines).strip()


def get_default_ollama_host() -> str:
    host = os.environ.get("OLLAMA_HOST") or os.environ.get("OLLAMA_URL") or "http://localhost:11434"
    if not host.startswith("http://") and not host.startswith("https://"):
        host = f"http://{host}"
    return host.rstrip("/")


async def probe_ollama(host: str | None = None, model: str = "gemma4:31b") -> tuple[bool, str]:
    """Check if Ollama server and requested model are reachable."""
    host = host or get_default_ollama_host()
    url = f"{host}/api/tags"
    try:
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=4.0)) as session:
            async with session.get(url) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    models = [m.get("name", "") for m in data.get("models", [])]
                    match = any(m.startswith(model.split(":")[0]) for m in models)
                    if match or not models:
                        return True, f"Ollama host {host} active with model {model}"
                    return True, f"Ollama host {host} active (models: {', '.join(models[:4])})"
                return False, f"Ollama returned HTTP {resp.status}"
    except Exception as e:
        return False, f"Cannot connect to Ollama at {host}: {e}"


class OllamaLLM(FrameProcessor):
    """Runs user turns through Ollama / OpenCode API."""

    def __init__(
        self,
        *,
        model: str = "gemma4:31b",
        host: str | None = None,
        system_prompt: str | None = None,
        observer=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._model = model
        self._host = host or get_default_ollama_host()
        self._system_prompt = system_prompt or "You are a helpful spoken voice assistant. Keep answers brief and conversational."
        self._on_reply = observer
        self._turn_task: asyncio.Task | None = None
        self._history: list[dict[str, str]] = []

    async def process_frame(self, frame: Frame, direction: FrameDirection):
        await super().process_frame(frame, direction)

        if isinstance(frame, StartFrame):
            await self.push_frame(frame, direction)
            logger.info(f"Ollama LLM engine initialized: host={self._host}, model={self._model}")
        elif isinstance(frame, (EndFrame, CancelFrame)):
            await self._cancel_turn()
            await self.push_frame(frame, direction)
        elif isinstance(frame, InterruptionFrame):
            await self._cancel_turn()
            await self.push_frame(frame, direction)
        elif isinstance(frame, LLMContextFrame):
            text = self._latest_user_text(frame.context)
            await self._maybe_start_turn(text)
        else:
            await self.push_frame(frame, direction)

    def _latest_user_text(self, context) -> str:
        if not context or not hasattr(context, "messages"):
            return ""
        for msg in reversed(context.messages):
            if isinstance(msg, dict) and msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, str):
                    return content
                elif isinstance(content, list):
                    text_parts = [c.get("text", "") for c in content if isinstance(c, dict)]
                    return " ".join(text_parts)
        return ""

    async def _maybe_start_turn(self, text: str):
        utterance = text.strip()
        if utterance.lower() in _NOISE_TRANSCRIPTS or len(utterance) < 2:
            logger.debug(f"Ignoring noise transcript: {utterance!r}")
            return

        await self._cancel_turn()
        logger.info(f"You: {utterance}")
        self._turn_task = self.create_task(self._run_turn(utterance))

    async def _cancel_turn(self):
        if not self._turn_task:
            return
        task, self._turn_task = self._turn_task, None
        await self.cancel_task(task)

    async def _run_turn(self, utterance: str):
        self._history.append({"role": "user", "content": utterance})
        recent_history = self._history[-8:]
        messages = [{"role": "system", "content": self._system_prompt}] + recent_history

        await self.push_frame(LLMFullResponseStartFrame())
        chunks: list[str] = []

        url = f"{self._host}/api/chat"
        payload = {
            "model": self._model,
            "messages": messages,
            "stream": True,
            "options": {
                "temperature": 0.7,
            },
        }

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as resp:
                    if resp.status != 200:
                        err_text = await resp.text()
                        logger.error(f"Ollama API HTTP {resp.status}: {err_text}")
                        err_msg = f"Sorry, Ollama API returned error {resp.status}."
                        chunks.append(err_msg)
                        await self.push_frame(LLMTextFrame(err_msg))
                    else:
                        async for line in resp.content:
                            line_str = line.decode("utf-8").strip()
                            if not line_str:
                                continue
                            try:
                                data = json.loads(line_str)
                                msg = data.get("message", {})
                                chunk = msg.get("content", "")
                                if chunk:
                                    chunks.append(chunk)
                                    cleaned_chunk = _clean_spoken_text(chunk)
                                    if cleaned_chunk:
                                        await self.push_frame(LLMTextFrame(cleaned_chunk))
                            except json.JSONDecodeError:
                                pass
        except asyncio.CancelledError:
            logger.info("Ollama turn cancelled mid-response.")
            raise
        except Exception as e:
            logger.error(f"Ollama LLM connection error: {e}")
            err_msg = "Sorry, I ran into an error connecting to Ollama."
            chunks.append(err_msg)
            await self.push_frame(LLMTextFrame(err_msg))
        finally:
            await self.push_frame(LLMFullResponseEndFrame())

        full_reply = "".join(chunks).strip()
        cleaned_reply = _clean_spoken_text(full_reply)
        if cleaned_reply:
            self._history.append({"role": "assistant", "content": cleaned_reply})
            logger.info(f"Ollama LLM ({self._model}): {cleaned_reply}")
            if self._on_reply:
                self._on_reply(cleaned_reply)