101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
"""
|
|
Search API for ChromaDB log collections.
|
|
Called by the Open WebUI Tool to answer questions about home events.
|
|
"""
|
|
|
|
import os
|
|
import requests
|
|
import chromadb
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI(title="Home Log Search API")
|
|
|
|
CHROMA_HOST = os.getenv("CHROMA_HOST", "chromadb")
|
|
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8000"))
|
|
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://ollama:11434")
|
|
EMBED_MODEL = "nomic-embed-text"
|
|
|
|
COLLECTIONS = ["frigate_events", "ha_history", "ha_logs", "docker_logs"]
|
|
|
|
|
|
def get_embedding(text: str) -> list[float]:
|
|
r = requests.post(
|
|
f"{OLLAMA_URL}/api/embeddings",
|
|
json={"model": EMBED_MODEL, "prompt": text[:4000]},
|
|
timeout=60,
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()["embedding"]
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
n_results: int = 10
|
|
sources: list[str] = [] # filter by source, empty = all
|
|
|
|
|
|
@app.post("/search")
|
|
def search(req: SearchRequest):
|
|
client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
|
|
embedding = get_embedding(req.query)
|
|
|
|
results = []
|
|
per_col = max(2, req.n_results // len(COLLECTIONS))
|
|
|
|
for col_name in COLLECTIONS:
|
|
try:
|
|
col = client.get_or_create_collection(col_name)
|
|
if col.count() == 0:
|
|
continue
|
|
|
|
where = None
|
|
if req.sources:
|
|
source_map = {
|
|
"frigate_events": "frigate",
|
|
"ha_history": "ha_history",
|
|
"ha_logs": "ha_log",
|
|
"docker_logs": "docker",
|
|
}
|
|
src = source_map.get(col_name)
|
|
if src and src not in req.sources:
|
|
continue
|
|
|
|
r = col.query(
|
|
query_embeddings=[embedding],
|
|
n_results=min(per_col, col.count()),
|
|
include=["documents", "metadatas", "distances"],
|
|
)
|
|
for doc, meta, dist in zip(
|
|
r["documents"][0], r["metadatas"][0], r["distances"][0]
|
|
):
|
|
results.append({
|
|
"text": doc,
|
|
"meta": meta,
|
|
"score": round(1 - dist, 3),
|
|
"collection": col_name,
|
|
})
|
|
except Exception as e:
|
|
pass
|
|
|
|
results.sort(key=lambda x: x["score"], reverse=True)
|
|
return results[: req.n_results]
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.get("/stats")
|
|
def stats():
|
|
client = chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)
|
|
out = {}
|
|
for col_name in COLLECTIONS:
|
|
try:
|
|
col = client.get_or_create_collection(col_name)
|
|
out[col_name] = col.count()
|
|
except Exception:
|
|
out[col_name] = "error"
|
|
return out
|