|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
import requests |
|
|
import os |
|
|
|
|
|
from ..needs_auth.OpenaiAPI import OpenaiAPI |
|
|
from ...requests import StreamSession, raise_for_status |
|
|
from ...providers.response import Usage, Reasoning |
|
|
from ...tools.run_tools import AuthManager |
|
|
from ...typing import AsyncResult, Messages |
|
|
|
|
|
class Ollama(OpenaiAPI): |
|
|
label = "Ollama 🦙" |
|
|
url = "https://ollama.com" |
|
|
login_url = "https://ollama.com/settings/keys" |
|
|
api_endpoint = "https://ollama.com/api/chat" |
|
|
needs_auth = False |
|
|
working = True |
|
|
active_by_default = True |
|
|
local_models: list[str] = [] |
|
|
model_aliases = { |
|
|
"gpt-oss-120b": "gpt-oss:120b", |
|
|
"gpt-oss-20b": "gpt-oss:20b" |
|
|
} |
|
|
|
|
|
@classmethod |
|
|
def get_models(cls, api_key: str = None, api_base: str = None, **kwargs): |
|
|
if not cls.models: |
|
|
cls.models = [] |
|
|
if not api_key: |
|
|
api_key = AuthManager.load_api_key(cls) |
|
|
if api_key: |
|
|
models = requests.get("https://ollama.com/api/tags", {"headers": {"Authorization": f"Bearer {api_key}"}}).json()["models"] |
|
|
if models: |
|
|
cls.live += 1 |
|
|
cls.models = [model["name"] for model in models] |
|
|
if api_base is None: |
|
|
host = os.getenv("OLLAMA_HOST", "127.0.0.1") |
|
|
port = os.getenv("OLLAMA_PORT", "11434") |
|
|
url = f"http://{host}:{port}/api/tags" |
|
|
else: |
|
|
url = api_base.replace("/v1", "/api/tags") |
|
|
try: |
|
|
models = requests.get(url).json()["models"] |
|
|
except requests.exceptions.RequestException as e: |
|
|
return cls.models |
|
|
if cls.live == 0 and models: |
|
|
cls.live += 1 |
|
|
cls.local_models = [model["name"] for model in models] |
|
|
cls.models = cls.models + cls.local_models |
|
|
cls.default_model = next(iter(cls.models), None) |
|
|
return cls.models |
|
|
|
|
|
@classmethod |
|
|
async def create_async_generator( |
|
|
cls, |
|
|
model: str, |
|
|
messages: Messages, |
|
|
api_key: str = None, |
|
|
api_base: str = None, |
|
|
proxy: str = None, |
|
|
**kwargs |
|
|
) -> AsyncResult: |
|
|
if api_base is None: |
|
|
host = os.getenv("OLLAMA_HOST", "localhost") |
|
|
port = os.getenv("OLLAMA_PORT", "11434") |
|
|
api_base: str = f"http://{host}:{port}/v1" |
|
|
if model in cls.local_models or not api_key: |
|
|
async for chunk in super().create_async_generator( |
|
|
model, messages, api_base=api_base, proxy=proxy, **kwargs |
|
|
): |
|
|
yield chunk |
|
|
else: |
|
|
async with StreamSession(headers={"Authorization": f"Bearer {api_key}"}, proxy=proxy) as session: |
|
|
async with session.post(cls.api_endpoint, json={ |
|
|
"model": model, |
|
|
"messages": messages, |
|
|
}) as response: |
|
|
await raise_for_status(response) |
|
|
last_data = {} |
|
|
async for chunk in response.iter_lines(): |
|
|
data = json.loads(chunk) |
|
|
last_data = data |
|
|
thinking = data.get("message", {}).get("thinking", "") |
|
|
if thinking: |
|
|
yield Reasoning(thinking) |
|
|
content = data.get("message", {}).get("content", "") |
|
|
if content: |
|
|
yield content |
|
|
yield Usage( |
|
|
prompt_tokens=last_data.get("prompt_eval_count", 0), |
|
|
completion_tokens=last_data.get("eval_count", 0), |
|
|
total_tokens=last_data.get("prompt_eval_count", 0) + last_data.get("eval_count", 0), |
|
|
) |
