diff --git a/main.py b/main.py new file mode 100644 index 0000000..7562155 --- /dev/null +++ b/main.py @@ -0,0 +1,31 @@ +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"))