import json from typing import List, Union, Generator, Iterator from pydantic import BaseModel import requests import re from datetime import timedelta class Pipe: class Valves(BaseModel): GITHUB_PAT: str = "" GITHUB_MODELS_BASE_URL: str = "https://models.inference.ai.azure.com" GITHUB_MARKETPLACE_BASE_URL: str = "https://github.com/marketplace" def __init__(self): self.id = "github_models" self.type = "manifold" self.name = "GitHub: " 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}", "Accept": "application/json", } models = [] page = 1 while True: r = requests.get( f"{self.valves.GITHUB_MARKETPLACE_BASE_URL}?page={page}&type=models", headers=headers, ) response = r.json() models.extend(response.get("results", [])) if page >= response.get("totalPages", 1): break page += 1 return [ { "id": ( model["original_name"] if model["original_name"] else model["name"] ), "name": ( model["friendly_name"] if "friendly_name" in model else model["name"] ), "description": (model["summary"] if "summary" in model else ""), } for model in models if model["task"] == "chat-completion" ] except Exception as e: print(f"Error: {e}") return [] def pipes(self) -> List[dict]: return self.get_github_models() def pipe(self, body: dict) -> Union[str, Generator, Iterator]: print(f"pipe:{__name__}") headers = { "Authorization": f"Bearer {self.valves.GITHUB_PAT}", "Content-Type": "application/json", } allowed_params = { "messages", "temperature", "top_p", "stop", "model", "max_tokens", "stream_options", "stream", # Include stream in allowed params } # 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} # Handle streaming for o1 models differently is_o1_model = "o1" in body["model"] should_stream = filtered_body.get("stream", False) # Don't remove stream parameters for o1 models, just set them appropriately if is_o1_model: filtered_body["stream"] = False # 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_BASE_URL}/chat/completions", json=filtered_body, headers=headers, stream=should_stream and not is_o1_model, # Only stream for non-o1 models ) if not r.status_code == 200: error_data = r.json() error_info = error_data.get("error", {}) message = error_info.get("message", "") try: message = json.loads(message).get("message", message) except Exception: pass details = error_info.get("details", "") # Extract wait time in seconds from the details match = re.search(r'Please wait (\d+) seconds before retrying', details) if not match: match = re.search(r'Please wait (\d+) seconds before retrying', message) if match: wait_seconds = int(match.group(1)) delta = timedelta(seconds=wait_seconds) readable_time = str(delta) message = f"Rate limit exceeded. Please wait {readable_time} before retrying." return f"Error: {message}" r.raise_for_status() if should_stream and not is_o1_model: return r.iter_lines() else: response_json = r.json() if response_json.get("choices"): return response_json["choices"][0]["message"]["content"] return response_json except Exception as e: return f"Error: {e} {r.text}"