41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
|
|
from fastapi import FastAPI
|
|
import requests, os
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import PointStruct
|
|
import hashlib
|
|
import json
|
|
|
|
app = FastAPI()
|
|
|
|
QDRANT_URL = os.getenv("QDRANT_URL")
|
|
OLLAMA_API = os.getenv("OLLAMA_API")
|
|
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "chat-memory")
|
|
|
|
client = QdrantClient(url=QDRANT_URL)
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "qdrant": QDRANT_URL, "ollama": OLLAMA_API}
|
|
|
|
def embed(text):
|
|
r = requests.post(f"{OLLAMA_API}/api/embeddings", json={"model":"nomic-embed-text","prompt":text})
|
|
return r.json()["embedding"]
|
|
|
|
@app.post("/store")
|
|
def store(item: dict):
|
|
text = item["text"]
|
|
metadata = item.get("metadata", {})
|
|
vec = embed(text)
|
|
pid = hashlib.sha256(text.encode()).hexdigest()
|
|
client.upsert(collection_name=COLLECTION_NAME, points=[PointStruct(id=pid, vector=vec, payload={"text": text, **metadata})])
|
|
return {"stored": True}
|
|
|
|
@app.post("/search")
|
|
def search(query: dict):
|
|
q = query["text"]
|
|
top_k = query.get("top_k", 5)
|
|
vec = embed(q)
|
|
result = client.search(collection_name=COLLECTION_NAME, query_vector=vec, limit=top_k)
|
|
return [{"score": r.score, "text": r.payload["text"]} for r in result]
|