Initial commit — kb-app RAG server

2026-05-10 23:07:43 +00:00
commit 89331c1fa5
13 changed files with 1055 additions and 0 deletions
@@ -0,0 +1,46 @@
+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]