from qdrant_client import QdrantClient
from ollama import Ollama
import os

# Initialize Qdrant client (local)
qdrant = QdrantClient(path="./qdrant")
# Create collection if not exists
if "rag_collection" not in [c.name for c in qdrant.get_collections()]:
    qdrant.create_collection(name="rag_collection", vectors_config={"size": 768, "distance": "Cosine"})

# Initialize Ollama client (local)
ol = Ollama()

# Simple function to add text to vector store

def add_document(text: str):
    # Embed using ollama embedding model
    embed = ol.embeddings(model="llama2", input=[text])['embeddings'][0]
    qdrant.add_points(collection_name="rag_collection", points=[{"id": len(qdrant.get_points(collection_name="rag_collection")) + 1, "vector": embed, "payload": {"text": text}}])

# Simple query function

def query(text: str):
    embed = ol.embeddings(model="llama2", input=[text])['embeddings'][0]
    results = qdrant.search(collection_name="rag_collection", query_vector=embed, limit=3)
    return [r.payload["text"] for r in results]

if __name__ == "__main__":
    # Example usage
    add_document("Hello world example.")
    print(query("world"))