58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
|
"""Ollama embedding client -- module 5, phase 4, plan step 0 (docs/kb/modules/05-faza4-plan.md,
|
||
|
|
§3, decision 1). Moved 1:1 out of `jobs/documents_ingest/chunk_embed.py` (`embed_chunk`,
|
||
|
|
`_vector_literal`, `DEFAULT_MODEL`, `DEFAULT_OLLAMA_URL`) so both a venv-based host job
|
||
|
|
(`documents-ingest`) and a long-lived Docker service (`kb-query`) can depend on the same tested
|
||
|
|
client without the service image pulling in `jobs/`'s `anthropic` dependency and CLI scripts.
|
||
|
|
|
||
|
|
`check_ollama_health` is new here (not moved) -- a generic `/api/tags` probe kb-query's
|
||
|
|
health-check + circuit breaker (plan §2 decision 2) builds on, kept in this module so the HTTP
|
||
|
|
probing logic lives in exactly one place.
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
|
||
|
|
import aiohttp
|
||
|
|
|
||
|
|
DEFAULT_OLLAMA_URL = "http://localhost:11434"
|
||
|
|
DEFAULT_MODEL = "bge-m3"
|
||
|
|
|
||
|
|
|
||
|
|
def _vector_literal(embedding: list[float]) -> str:
|
||
|
|
"""Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'."""
|
||
|
|
return "[" + ",".join(repr(v) for v in embedding) + "]"
|
||
|
|
|
||
|
|
|
||
|
|
async def embed_chunk(
|
||
|
|
session: aiohttp.ClientSession, base_url: str, model: str, text: str
|
||
|
|
) -> tuple[list[float], float]:
|
||
|
|
"""POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds).
|
||
|
|
No built-in timeout/retry -- inherits whatever `aiohttp.ClientSession(timeout=...)` the
|
||
|
|
caller constructed; `raise_for_status()` propagates `aiohttp.ClientError` when Ollama is
|
||
|
|
unreachable."""
|
||
|
|
t0 = time.monotonic()
|
||
|
|
async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp:
|
||
|
|
resp.raise_for_status()
|
||
|
|
data = await resp.json()
|
||
|
|
elapsed = time.monotonic() - t0
|
||
|
|
|
||
|
|
embedding = data.get("embedding")
|
||
|
|
if not embedding:
|
||
|
|
raise ValueError(f"ollama response missing 'embedding': {data!r}")
|
||
|
|
return embedding, elapsed
|
||
|
|
|
||
|
|
|
||
|
|
async def check_ollama_health(
|
||
|
|
session: aiohttp.ClientSession, base_url: str, timeout_s: float
|
||
|
|
) -> bool:
|
||
|
|
"""Probe `GET {base_url}/api/tags`. True iff it answers within `timeout_s` with a non-error
|
||
|
|
status; any exception (timeout, connection refused, 4xx/5xx) is treated as `down`, never
|
||
|
|
raised -- callers use this for liveness decisions, not error propagation."""
|
||
|
|
try:
|
||
|
|
async with session.get(
|
||
|
|
f"{base_url}/api/tags", timeout=aiohttp.ClientTimeout(total=timeout_s)
|
||
|
|
) as resp:
|
||
|
|
return resp.status < 400
|
||
|
|
except (aiohttp.ClientError, TimeoutError):
|
||
|
|
return False
|