import json
import chromadb
import requests

rows = [json.loads(line) for line in open("chunks.jsonl")]
client = chromadb.HttpClient(host="localhost", port=8010)
col = client.get_or_create_collection("docs", configuration={"hnsw": {"space": "cosine"}})
col.upsert(ids=[r["id"] for r in rows], documents=[r["text"] for r in rows],
           embeddings=[r["embedding"] for r in rows], metadatas=[{"url": r["url"]} for r in rows])
print("count:", col.count())
q = "search_query: How do I limit how many pages a crawl fetches?"
body = {"model": "nomic-embed-text", "input": q}
vec = requests.post("http://localhost:11434/api/embed", json=body).json()["embeddings"][0]
res = col.query(query_embeddings=[vec], n_results=3)
for m, d in zip(res["metadatas"][0], res["distances"][0]):
    print(round(d, 3), m["url"])