Open WebUI Github Models Manifold V2

github_models_manifold_v2.py

"""
title: GitHub Models Manifold V2
author: wendyliga
author_url: https://github.com/wendyliga
project_url: https://gist.github.com/wendyliga/a1247718c105a45d86e84a1cae60c20e
version: 1.0
license: Apache License 2.0

description: A manifold pipeline for interacting with github models. modified from https://openwebui.com/f/jscheah/github_models_manifold to adapt with new endpoint.
"""

from typing import List, Union, Generator, Iterator
from pydantic import BaseModel

import requests


class Pipe:
    class Valves(BaseModel):
        GITHUB_PAT: str = ""

        # GITHUB_MODELS_BASE_URL: str = "https://models.inference.ai.azure.com"
        GITHUB_MODELS_BASE_URL: str = "https://models.github.ai"
        GITHUB_MODELS_URL: str = GITHUB_MODELS_BASE_URL + "/catalog"
        GITHUB_MODELS_CHAT_COMPLETION: str = GITHUB_MODELS_BASE_URL + "/inference"

    def __init__(self):
        self.id = "github_models"
        self.type = "manifold"
        self.name = "GitHub Models: "

        self.valves = self.Valves()

        self.pipelines = self.get_github_models()

    def get_github_models(self):
        if self.valves.GITHUB_PAT:
            try:
                headers = {
                    "Authorization": f"Bearer {self.valves.GITHUB_PAT}",
                    "Content-Type": "application/json",
                    "Accept": "application/vnd.github+json",
                    "X-GitHub-Api-Version": "2022-11-28",
                }

                r = requests.get(
                    f"{self.valves.GITHUB_MODELS_URL}/models", headers=headers
                )

                models = r.json()
                return [
                    {
                        "id": model["id"],
                        "name": model["name"],
                        "description": (model["summary"] if "summary" in model else ""),
                    }
                    for model in models
                ]

            except Exception as e:

                print(f"Error: {e}")
                return [
                    {
                        "id": "error",
                        "name": "Could not fetch models from GitHub Models, please update the PAT in the valves.",
                    },
                ]
        else:
            return []

    def pipes(self) -> List[dict]:
        return self.get_github_models()

    def pipe(self, body: dict) -> Union[str, Generator, Iterator]:
        # This is where you can add your custom pipelines like RAG.
        print(f"pipe:{__name__}")

        headers = {
            "Authorization": f"Bearer {self.valves.GITHUB_PAT}",
            "Content-Type": "application/json",
            "Accept": "application/vnd.github+json",
            "X-GitHub-Api-Version": "2022-11-28",
        }

        allowed_params = {
            "messages",
            "temperature",
            "top_p",
            "stream",
            "stop",
            "model",
            "max_tokens",
            "stream_options",
        }
        # Remap the model name to the model id
        body["model"] = ".".join(body["model"].split(".")[1:])

        filtered_body = {k: v for k, v in body.items() if k in allowed_params}
        # log fields that were filtered out as a single line
        if len(body) != len(filtered_body):
            print(
                f"Dropped params: {', '.join(set(body.keys()) - set(filtered_body.keys()))}"
            )

        try:
            r = requests.post(
                url=f"{self.valves.GITHUB_MODELS_CHAT_COMPLETION}/chat/completions",
                json=filtered_body,
                headers=headers,
                stream=True,
            )
            r.raise_for_status()

            if body["stream"]:
                return r.iter_lines()
            else:
                return r.json()
        except Exception as e:
            return f"Error: {e} {r.text}"