32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
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"))
|