Julian Carlson
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
import requests
|
|
from .config import CONFIG
|
|
|
|
|
|
def embed_one(text):
|
|
"""Embed a single text. Handles both old and new Ollama embedding APIs."""
|
|
try:
|
|
resp = requests.post(
|
|
f"{CONFIG.ollama_url}/api/embed",
|
|
json={"model": CONFIG.embed_model, "input": text},
|
|
timeout=120,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
if "embeddings" in data:
|
|
return data["embeddings"][0]
|
|
if "embedding" in data:
|
|
return data["embedding"]
|
|
except Exception:
|
|
resp = requests.post(
|
|
f"{CONFIG.ollama_url}/api/embeddings",
|
|
json={"model": CONFIG.embed_model, "prompt": text},
|
|
timeout=120,
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()["embedding"]
|
|
|
|
raise RuntimeError(f"Unexpected embedding response: {data}")
|
|
|
|
|
|
def embed_batch(texts):
|
|
"""Embed multiple texts. Use new API native batching when possible."""
|
|
try:
|
|
resp = requests.post(
|
|
f"{CONFIG.ollama_url}/api/embed",
|
|
json={"model": CONFIG.embed_model, "input": texts},
|
|
timeout=300,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
if "embeddings" in data:
|
|
return data["embeddings"]
|
|
except Exception:
|
|
pass
|
|
|
|
return [embed_one(t) for t in texts]
|