118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
import base64
|
||
import mimetypes
|
||
import os
|
||
import time
|
||
from typing import Any, Dict, List
|
||
|
||
import requests
|
||
|
||
|
||
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://192.168.1.250:11434")
|
||
DEFAULT_MODEL = os.getenv("VLM_MODEL", "qwen2.5vl:32b")
|
||
|
||
|
||
def _img_to_data_uri(path: str) -> str:
|
||
mime, _ = mimetypes.guess_type(path)
|
||
if not mime:
|
||
ext = os.path.splitext(path)[1].lower()
|
||
if ext in {".jpg", ".jpeg"}:
|
||
mime = "image/jpeg"
|
||
elif ext == ".png":
|
||
mime = "image/png"
|
||
else:
|
||
mime = "application/octet-stream"
|
||
with open(path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode("utf-8")
|
||
return f"data:{mime};base64,{b64}"
|
||
|
||
|
||
def chat_vision(
|
||
prompt: str,
|
||
image_paths: List[str],
|
||
*,
|
||
temperature: float = 0.1,
|
||
options: Dict[str, Any] | None = None, # kept for API compatibility, not used in /v1 schema
|
||
model: str = DEFAULT_MODEL,
|
||
timeout_sec: int = 120,
|
||
) -> str:
|
||
# Logging: request summary before heavy work
|
||
try:
|
||
img_sizes = []
|
||
for p in image_paths:
|
||
try:
|
||
img_sizes.append(os.path.getsize(p))
|
||
except Exception:
|
||
img_sizes.append(-1)
|
||
print(
|
||
f"[Ollama] Prepare chat url={OLLAMA_URL}/v1/chat/completions model={model} temp={temperature} "
|
||
f"images={len(image_paths)} sizes={img_sizes} timeout={timeout_sec}s"
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
# OpenAI-compatible vision message format: content is a list of parts (text + image_url)
|
||
content_parts: List[Dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||
for p in image_paths:
|
||
content_parts.append({
|
||
"type": "image_url",
|
||
"image_url": {"url": _img_to_data_uri(p)}
|
||
})
|
||
|
||
payload: Dict[str, Any] = {
|
||
"model": model,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": content_parts,
|
||
}
|
||
],
|
||
"temperature": temperature,
|
||
"stream": False,
|
||
}
|
||
|
||
t0 = time.time()
|
||
try:
|
||
resp = requests.post(f"{OLLAMA_URL}/v1/chat/completions", json=payload, timeout=timeout_sec)
|
||
dt = time.time() - t0
|
||
print(f"[Ollama] Response status={resp.status_code} elapsed={dt:.2f}s")
|
||
if resp.status_code >= 400:
|
||
body = resp.text
|
||
snippet = (body[:400]).replace("\n", " ")
|
||
print(f"[Ollama][ErrorBody] {snippet}...")
|
||
resp.raise_for_status()
|
||
|
||
# Parse OpenAI-compatible response
|
||
# Expected: { choices: [ { message: { content: "..." } } ], ... }
|
||
data = resp.json()
|
||
if not isinstance(data, dict) or "choices" not in data or not data["choices"]:
|
||
snippet = (resp.text[:400]).replace("\n", " ")
|
||
print(f"[Ollama][Error] Unexpected response shape: {snippet}")
|
||
raise RuntimeError("Unexpected response from /v1/chat/completions")
|
||
|
||
choice0 = data["choices"][0]
|
||
msg = choice0.get("message") or {}
|
||
text = msg.get("content") or ""
|
||
if not text:
|
||
# Some servers might put text in 'choices[0].text' (rare for chat) – try fallback
|
||
text = choice0.get("text") or ""
|
||
# _preview = (text[:120]).replace("\n", " ")
|
||
result = text.replace("\n", " ")
|
||
# print(f"[Ollama] Text length={len(text)} preview={_preview}...")
|
||
print(f"[Ollama] Text length={len(text)} preview={result}")
|
||
if not text:
|
||
raise RuntimeError("Empty content in choices[0].message.content")
|
||
return text
|
||
except Exception as e:
|
||
dt = time.time() - t0
|
||
# Print body to help debugging when JSON parsing fails, etc.
|
||
try:
|
||
body = resp.text # type: ignore[name-defined]
|
||
snippet = (body[:400]).replace("\n", " ")
|
||
print(f"[Ollama][ErrorBody] {snippet}...")
|
||
except Exception:
|
||
pass
|
||
print(f"[Ollama][Error] during request elapsed={dt:.2f}s: {e}")
|
||
raise
|
||
|
||
|