add main

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"))