links/qdrant_sync.py

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://192.168.1.1:8000/api/links/")
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()