""" title: Ollama Api Facade Metadata author: Gregor Biswanger author_url: https://github.com/GregorBiswanger/OllamaApiFacade version: 1.0.0 date: 2024-11-28 license: Apache 2.0 description: A Pipe that enhances .NET Ollama API Facade requests by including additional metadata fields such as id, chat_id, and session_id for advanced context handling. """ from pydantic import BaseModel, Field from typing import Union, AsyncGenerator import aiohttp import json import requests class Pipe: class Valves(BaseModel): OLLAMA_API_BASE_URL: str = Field( default="http://host.docker.internal:11434", description="Base URL for accessing the Ollama API Facade.", ) def __init__(self): self.type = "manifold" self.valves = self.Valves() def pipes(self): try: headers = {"Content-Type": "application/json"} response = requests.get( f"{self.valves.OLLAMA_API_BASE_URL}/api/tags", headers=headers ) response.raise_for_status() models = response.json().get("models", []) print(f"Debug: Available models: {models}") return [{"id": model["model"], "name": model["name"]} for model in models] except Exception as e: print(f"Error fetching models: {e}") return [{"id": "error", "name": "Failed to fetch models."}] async def pipe( self, body: dict, __user__: dict, __metadata__: dict ) -> AsyncGenerator[str, None]: headers = {"Content-Type": "application/json"} model_name = body["model"].split(".")[-1] transformed_body = { "stream": True, "model": model_name, "messages": body.get("messages", []), "options": body.get("options", {}), "session_id": __metadata__.get("session_id"), "chat_id": __metadata__.get("chat_id"), "id": __user__.get("id"), "max_tokens": body.get("max_tokens", 256), } print(f"Request to Ollama: {json.dumps(transformed_body, indent=2)}") try: async with aiohttp.ClientSession() as session: async with session.post( url=f"{self.valves.OLLAMA_API_BASE_URL}/api/chat", json=transformed_body, headers=headers, ) as response: response.raise_for_status() print("Debug: Starting stream processing...") async for line in response.content: line = line.decode("utf-8").strip() if line: try: parsed_line = json.loads(line) print(f"Debug: Received data: {parsed_line}") if ( "message" in parsed_line and "content" in parsed_line["message"] ): yield parsed_line["message"]["content"] except json.JSONDecodeError as e: print(f"Error parsing line: {e}") continue except aiohttp.ClientError as e: print(f"Error during request: {e}") yield json.dumps({"error": f"Request error: {e}"})