import chromadb
import requests
from langchain_chroma import Chroma
from langchain_core.embeddings import Embeddings

from common import HOST, PORT, COLLECTION, OLLAMA


class OllamaQuery(Embeddings):
    """Embeds queries with the nomic search_query prefix. Documents are already in the collection."""

    def embed_query(self, text):
        r = requests.post(OLLAMA, json={"model": "nomic-embed-text", "input": "search_query: " + text}, timeout=60)
        return r.json()["embeddings"][0]

    def embed_documents(self, texts):
        raise NotImplementedError("vectors come from chunks.jsonl")


store = Chroma(client=chromadb.HttpClient(host=HOST, port=PORT), collection_name=COLLECTION,
               embedding_function=OllamaQuery())
q = "How do I get notified when a crawl job finishes?"
for doc, score in store.similarity_search_with_score(q, k=3, filter={"section": "guides"}):
    print(f"{score:.3f}  {doc.metadata['url']}")
print("count via LangChain:", store._collection.count())