openwebui/functions/gemma4_audio_pipe.py

"""
title: Gemma 4 Omni (Audio)
author: gemma4-stack
version: 0.1.0
required_open_webui_version: 0.5.0
description: >
  Sends an attached audio clip to Gemma 4 as a NATIVE multimodal input
  (OpenAI `input_audio` content part) via llama-swap/llama-server — the model
  "hears" the audio instead of transcribing it with Whisper. Resamples to
  16 kHz mono first (Gemma 4's reliable envelope, clips <= ~30 s).

HOW TO USE
  Admin → Functions → "+" → paste this file → Save → enable it.
  In a new chat pick the model "Gemma 4 · Omni (audio)", attach a short audio
  clip, type your question (Persian or any language), send.

FRAGILE BIT (read me)
  Open WebUI hands a Pipe only FILE REFERENCES, not bytes. _resolve_local_path()
  below fetches the real file from Open WebUI's store. That store API moves
  between versions — if audio isn't found, adjust _resolve_local_path() to your
  installed Open WebUI version first (see the strategies inside).
"""

import os
import json
import base64
import glob
import shutil
import subprocess
import tempfile
from typing import List, Optional

import requests
from pydantic import BaseModel, Field

AUDIO_EXTS = (".wav", ".mp3", ".flac", ".ogg", ".m4a", ".webm", ".aac", ".opus")
DATA_DIR = os.environ.get("DATA_DIR", "/app/backend/data")
UPLOADS_DIR = os.path.join(DATA_DIR, "uploads")


class Pipe:
    class Valves(BaseModel):
        LLAMASWAP_URL: str = Field(
            default="http://llama-swap:8080/v1",
            description="OpenAI-compatible base URL of llama-swap.",
        )
        MODEL: str = Field(
            default="gemma-e4b",
            description="llama-swap model key (audio-capable: gemma-e4b/gemma-12b).",
        )
        API_KEY: str = Field(default="sk-local", description="Bearer key for llama-swap.")
        TEMPERATURE: float = Field(default=1.0)
        TOP_K: int = Field(default=64)
        TOP_P: float = Field(default=0.95)
        MAX_TOKENS: int = Field(default=512)
        TARGET_SR: int = Field(default=16000, description="Resample rate (Gemma wants 16 kHz mono).")
        DEFAULT_PROMPT: str = Field(
            default="این فایل صوتی را دقیق بنویس و در صورت نیاز توضیح بده.",
            description="Used when the user attaches audio without typing a question.",
        )

    def __init__(self):
        self.valves = self.Valves()

    def pipes(self):
        return [{"id": "gemma4-audio", "name": "Gemma 4 · Omni (audio)"}]

    # ------------------------------------------------------------------ helpers
    def _collect_file_refs(self, body, __files__, __metadata__) -> List[dict]:
        refs = []
        if __files__:
            refs += __files__
        if isinstance(__metadata__, dict):
            refs += __metadata__.get("files", []) or []
        meta = (body or {}).get("metadata", {}) or {}
        refs += meta.get("files", []) or []
        refs += (body or {}).get("files", []) or []
        return refs

    def _ref_id_and_name(self, ref: dict):
        # Open WebUI nests the actual record under "file" in some versions.
        inner = ref.get("file", ref) if isinstance(ref, dict) else {}
        fid = ref.get("id") or inner.get("id")
        name = (
            ref.get("name")
            or inner.get("filename")
            or (inner.get("meta") or {}).get("name")
            or ""
        )
        ctype = (inner.get("meta") or {}).get("content_type", "") or ref.get("type", "")
        return fid, name, ctype

    def _resolve_local_path(self, fid, name) -> Optional[str]:
        """Turn a file reference into a real on-disk path. Version-sensitive."""
        # Strategy 1: official Files model.
        try:
            from open_webui.models.files import Files  # type: ignore

            rec = Files.get_file_by_id(fid)
            if rec is not None:
                p = getattr(rec, "path", None) or (getattr(rec, "meta", {}) or {}).get("path")
                if p:
                    if not os.path.isabs(p):
                        p = os.path.join(DATA_DIR, p)
                    if os.path.exists(p):
                        return p
        except Exception:
            pass
        # Strategy 2: storage provider abstraction.
        try:
            from open_webui.storage.provider import Storage  # type: ignore

            p = Storage.get_file(f"uploads/{fid}")  # may raise / vary
            if p and os.path.exists(p):
                return p
        except Exception:
            pass
        # Strategy 3: scan the uploads dir for <id> or <name>.
        for pattern in (f"*{fid}*", f"*{name}*"):
            if not pattern.strip("*"):
                continue
            hits = glob.glob(os.path.join(UPLOADS_DIR, pattern))
            hits = [h for h in hits if os.path.isfile(h)]
            if hits:
                return max(hits, key=os.path.getmtime)
        return None

    def _to_16k_mono_wav(self, src: str) -> (str, str):
        """Return (path, format). Resample via ffmpeg if available, else pass through."""
        ffmpeg = shutil.which("ffmpeg")
        if ffmpeg:
            out = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
            try:
                subprocess.run(
                    [ffmpeg, "-y", "-i", src, "-ar", str(self.valves.TARGET_SR),
                     "-ac", "1", "-f", "wav", out],
                    check=True, capture_output=True,
                )
                return out, "wav"
            except Exception:
                pass
        ext = os.path.splitext(src)[1].lower().lstrip(".") or "wav"
        return src, ("wav" if ext not in ("mp3", "flac", "wav") else ext)

    def _latest_user_text(self, body) -> str:
        for msg in reversed((body or {}).get("messages", [])):
            if msg.get("role") == "user":
                c = msg.get("content")
                if isinstance(c, str):
                    return c.strip()
                if isinstance(c, list):
                    parts = [p.get("text", "") for p in c if p.get("type") == "text"]
                    return " ".join(t for t in parts if t).strip()
        return ""

    # --------------------------------------------------------------------- main
    def pipe(self, body: dict, __user__=None, __request__=None,
             __files__=None, __metadata__=None):
        refs = self._collect_file_refs(body, __files__, __metadata__)
        audio_paths = []
        for ref in refs:
            fid, name, ctype = self._ref_id_and_name(ref)
            is_audio = ctype.startswith("audio") or name.lower().endswith(AUDIO_EXTS)
            if not is_audio:
                continue
            local = self._resolve_local_path(fid, name)
            if local:
                audio_paths.append(local)

        if not audio_paths:
            return (
                "⚠️ No audio found. Attach a short clip (≤ ~30 s) and ask your "
                "question. If you *did* attach audio, the file-store lookup needs "
                "adapting to your Open WebUI version — see _resolve_local_path()."
            )

        text = self._latest_user_text(body) or self.valves.DEFAULT_PROMPT
        content = [{"type": "text", "text": text}]
        for p in audio_paths:
            wav, fmt = self._to_16k_mono_wav(p)
            with open(wav, "rb") as f:
                content.append({
                    "type": "input_audio",
                    "input_audio": {
                        "data": base64.b64encode(f.read()).decode("ascii"),
                        "format": fmt,
                    },
                })

        payload = {
            "model": self.valves.MODEL,
            "messages": [{"role": "user", "content": content}],
            "temperature": self.valves.TEMPERATURE,
            "top_k": self.valves.TOP_K,
            "top_p": self.valves.TOP_P,
            "max_tokens": self.valves.MAX_TOKENS,
            "stream": True,
        }
        headers = {"Authorization": f"Bearer {self.valves.API_KEY}"}
        url = self.valves.LLAMASWAP_URL.rstrip("/") + "/chat/completions"

        def gen():
            with requests.post(url, json=payload, headers=headers,
                               stream=True, timeout=600) as r:
                r.raise_for_status()
                for line in r.iter_lines(decode_unicode=True):
                    if not line or not line.startswith("data:"):
                        continue
                    data = line[len("data:"):].strip()
                    if data == "[DONE]":
                        break
                    try:
                        delta = json.loads(data)["choices"][0]["delta"]
                        piece = delta.get("content")
                        if piece:
                            yield piece
                    except Exception:
                        continue

        return gen()