知微
fa45d8aa5f
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
Privoxy对node122:18003返回500，直连正常
266 lines
8.4 KiB
Python
266 lines
8.4 KiB
Python
from typing import Any, List, Optional, Union
|
|
|
|
import httpx
|
|
|
|
from litellm import verbose_logger
|
|
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|
|
|
|
|
class OllamaError(BaseLLMException):
|
|
def __init__(
|
|
self, status_code: int, message: str, headers: Union[dict, httpx.Headers]
|
|
):
|
|
super().__init__(status_code=status_code, message=message, headers=headers)
|
|
|
|
|
|
def _convert_image(image):
|
|
"""
|
|
Convert image to base64 encoded image if not already in base64 format
|
|
|
|
If image is already in base64 format AND is a jpeg/png, return it
|
|
|
|
If image is not JPEG/PNG, convert it to JPEG base64 format
|
|
"""
|
|
import base64
|
|
import io
|
|
|
|
try:
|
|
from PIL import Image
|
|
except Exception:
|
|
raise Exception(
|
|
"ollama image conversion failed please run `pip install Pillow`"
|
|
)
|
|
|
|
orig = image
|
|
if image.startswith("data:"):
|
|
image = image.split(",")[-1]
|
|
try:
|
|
image_data = Image.open(io.BytesIO(base64.b64decode(image)))
|
|
if image_data.format in ["JPEG", "PNG"]:
|
|
return image
|
|
except Exception:
|
|
return orig
|
|
jpeg_image = io.BytesIO()
|
|
image_data.convert("RGB").save(jpeg_image, "JPEG")
|
|
jpeg_image.seek(0)
|
|
return base64.b64encode(jpeg_image.getvalue()).decode("utf-8")
|
|
|
|
|
|
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
|
|
|
|
|
|
class OllamaModelInfo(BaseLLMModelInfo):
|
|
"""
|
|
Dynamic model listing for Ollama server.
|
|
Fetches /api/models and /api/tags, then for each tag also /api/models?tag=...
|
|
Returns the union of all model names.
|
|
"""
|
|
|
|
@staticmethod
|
|
def get_api_key(api_key=None) -> Optional[str]:
|
|
"""Get API key from environment variables or litellm configuration"""
|
|
import os
|
|
|
|
import litellm
|
|
from litellm.secret_managers.main import get_secret_str
|
|
|
|
return (
|
|
api_key
|
|
or os.environ.get("OLLAMA_API_KEY")
|
|
or litellm.api_key
|
|
or litellm.openai_key
|
|
or get_secret_str("OLLAMA_API_KEY")
|
|
)
|
|
|
|
@staticmethod
|
|
def get_api_base(api_base: Optional[str] = None) -> str:
|
|
from litellm.secret_managers.main import get_secret_str
|
|
|
|
# env var OLLAMA_API_BASE or default
|
|
return api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434"
|
|
|
|
@classmethod
|
|
def get_server_api_base(cls, api_base: Optional[str] = None) -> str:
|
|
api_base = cls.get_api_base(api_base).rstrip("/")
|
|
for suffix in (
|
|
"/api/generate",
|
|
"/api/chat",
|
|
"/api/embed",
|
|
"/api/embeddings",
|
|
"/api/show",
|
|
"/api/tags",
|
|
):
|
|
if api_base.endswith(suffix):
|
|
return api_base[: -len(suffix)]
|
|
return api_base
|
|
|
|
def get_models(self, api_key=None, api_base: Optional[str] = None) -> List[str]:
|
|
"""
|
|
List all models available on the Ollama server via /api/tags endpoint.
|
|
"""
|
|
|
|
passed_api_base = api_base
|
|
base = self.get_server_api_base(api_base)
|
|
api_key = (
|
|
self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
|
)
|
|
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
|
|
|
names: set[str] = set()
|
|
try:
|
|
resp = httpx.get(f"{base}/api/tags", headers=headers)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
# Expecting a dict with a 'models' list
|
|
models_list = []
|
|
if (
|
|
isinstance(data, dict)
|
|
and "models" in data
|
|
and isinstance(data["models"], list)
|
|
):
|
|
models_list = data["models"]
|
|
elif isinstance(data, list):
|
|
models_list = data
|
|
# Extract model names
|
|
for entry in models_list:
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
nm = entry.get("name") or entry.get("model")
|
|
if isinstance(nm, str):
|
|
names.add(nm if nm.startswith("ollama/") else f"ollama/{nm}")
|
|
except Exception as e:
|
|
verbose_logger.warning(f"Error retrieving ollama tag endpoint: {e}")
|
|
# If tags endpoint fails, fall back to static list
|
|
try:
|
|
from litellm import models_by_provider
|
|
|
|
static = models_by_provider.get("ollama", []) or []
|
|
return [f"ollama/{m}" for m in static]
|
|
except Exception as e1:
|
|
verbose_logger.warning(
|
|
f"Error retrieving static ollama models as fallback: {e1}"
|
|
)
|
|
return []
|
|
# assemble full model names
|
|
result = sorted(names)
|
|
return result
|
|
|
|
@staticmethod
|
|
def _strip_ollama_model_prefix(model: str) -> str:
|
|
if model.startswith("ollama/") or model.startswith("ollama_chat/"):
|
|
return model.split("/", 1)[1]
|
|
return model
|
|
|
|
@staticmethod
|
|
def _is_static_ollama_model(model: str) -> bool:
|
|
from litellm import model_cost
|
|
|
|
stripped_model = OllamaModelInfo._strip_ollama_model_prefix(model)
|
|
potential_model_names = {
|
|
model,
|
|
stripped_model,
|
|
"ollama/" + stripped_model,
|
|
"ollama_chat/" + stripped_model,
|
|
}
|
|
model_cost_keys = {key.lower() for key in model_cost}
|
|
return any(name.lower() in model_cost_keys for name in potential_model_names)
|
|
|
|
@staticmethod
|
|
def _supports_function_calling(ollama_model_info: dict) -> bool:
|
|
_template: str = str(ollama_model_info.get("template", "") or "")
|
|
return "tools" in _template.lower()
|
|
|
|
@staticmethod
|
|
def _get_max_tokens(ollama_model_info: dict) -> Optional[int]:
|
|
_model_info: dict = ollama_model_info.get("model_info", {})
|
|
|
|
for key, value in _model_info.items():
|
|
if "context_length" in key:
|
|
return value
|
|
return None
|
|
|
|
def get_runtime_model_info(
|
|
self,
|
|
model: str,
|
|
api_base: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
) -> dict[str, Any]:
|
|
from litellm import module_level_client
|
|
|
|
model = self._strip_ollama_model_prefix(model)
|
|
passed_api_base = api_base
|
|
api_base = self.get_server_api_base(api_base)
|
|
api_key = (
|
|
self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
|
)
|
|
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
|
|
|
try:
|
|
response = module_level_client.post(
|
|
url=f"{api_base}/api/show",
|
|
json={"name": model},
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
except Exception:
|
|
verbose_logger.debug("OllamaError: Could not get model info.")
|
|
return {
|
|
"key": model,
|
|
"litellm_provider": "ollama",
|
|
"mode": "chat",
|
|
"input_cost_per_token": 0.0,
|
|
"output_cost_per_token": 0.0,
|
|
"max_tokens": None,
|
|
"max_input_tokens": None,
|
|
"max_output_tokens": None,
|
|
}
|
|
|
|
model_info = response.json()
|
|
max_tokens = self._get_max_tokens(model_info)
|
|
|
|
return {
|
|
"key": model,
|
|
"litellm_provider": "ollama",
|
|
"mode": "chat",
|
|
"supports_function_calling": self._supports_function_calling(model_info),
|
|
"input_cost_per_token": 0.0,
|
|
"output_cost_per_token": 0.0,
|
|
"max_tokens": max_tokens,
|
|
"max_input_tokens": max_tokens,
|
|
"max_output_tokens": max_tokens,
|
|
}
|
|
|
|
def get_model_info(
|
|
self,
|
|
model: str,
|
|
api_base: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
) -> Optional[dict[str, Any]]:
|
|
if self._is_static_ollama_model(model):
|
|
return None
|
|
return self.get_runtime_model_info(
|
|
model=model, api_base=api_base, api_key=api_key
|
|
)
|
|
|
|
def validate_environment(
|
|
self,
|
|
headers: dict,
|
|
model: str,
|
|
messages: list,
|
|
optional_params: dict,
|
|
litellm_params: dict,
|
|
api_key=None,
|
|
api_base=None,
|
|
) -> dict:
|
|
"""
|
|
No-op environment validation for Ollama.
|
|
"""
|
|
return {}
|
|
|
|
@staticmethod
|
|
def get_base_model(model: str) -> str:
|
|
"""
|
|
Return the base model name for Ollama (no-op).
|
|
"""
|
|
return model
|