Files
links/qdrant_sync.py
T

98 lines
3.3 KiB
Python

import requests
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import os
# Configuration from environment variables
DJANGO_API = os.getenv("DJANGO_API", "http://links.apps.svc.cluster.local/api/posts/")
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://192.168.1.1:11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
QDRANT_HOST = os.getenv("QDRANT_HOST", "192.168.1.2")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
COLLECTION_NAME = os.getenv("COLLECTION_NAME", "links")
def get_embedding(text):
"""Generate embedding using Ollama"""
response = requests.post(
f"{OLLAMA_URL}/api/embeddings",
json={"model": OLLAMA_MODEL, "prompt": text}
)
response.raise_for_status()
return response.json()["embedding"]
def fetch_links():
"""Fetch links from Django API"""
# Assuming the API uses pagination or returns a list
response = requests.get(DJANGO_API)
response.raise_for_status()
data = response.json()
if isinstance(data, dict) and "results" in data:
return data["results"]
return data
def sync_to_qdrant():
"""Sync links from Django API to Qdrant"""
print(f"Connecting to Qdrant at {QDRANT_HOST}:{QDRANT_PORT}")
client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
print(f"Fetching links from {DJANGO_API}")
links = fetch_links()
if not links:
print("No links to sync")
return
print(f"Found {len(links)} links. Preparing for vectorization...")
# Get embedding dimension
sample_text = f"{links[0].get('title', '')} {links[0].get('url', '')}"
sample_embedding = get_embedding(sample_text)
vector_size = len(sample_embedding)
# Ensure collection exists
collections = client.get_collections().collections
exists = any(c.name == COLLECTION_NAME for c in collections)
if not exists:
print(f"Creating collection: {COLLECTION_NAME}")
client.create_collection(
collection_name=COLLECTION_NAME,
vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE)
)
# Prepare points
points = []
for idx, link in enumerate(links):
# Combine text fields for embedding
text = f"{link.get('title', '')} {link.get('description', '') or ''} {link.get('url', '')}"
try:
embedding = get_embedding(text)
# Use link id if available, otherwise index
point_id = link.get("id", idx)
point = PointStruct(
id=point_id,
vector=embedding,
payload={
"id": link.get("id"),
"title": link.get("title"),
"url": link.get("url"),
"description": link.get("description"),
"tags": link.get("tags", [])
}
)
points.append(point)
if (idx + 1) % 10 == 0:
print(f"Processed {idx + 1}/{len(links)} links")
except Exception as e:
print(f"Error processing link {idx}: {e}")
# Upload to Qdrant
if points:
client.upsert(collection_name=COLLECTION_NAME, points=points)
print(f"Successfully synced {len(points)} links to Qdrant")
if __name__ == "__main__":
sync_to_qdrant()