vector_embed.py

import os
import uuid
from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer
from qdrant_client.models import Distance

# --- 1. Configuration ---
QDRANT_URL = "https://655acf9d-ceb0-4da8-9b50-762b58ece340.eu-central-1-0.aws.cloud.qdrant.io:6333"
QDRANT_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.tDd1xW9MwCFbiOdeOvMfLo0zFnkBehJORWXfvzoVspI"
COLLECTION_NAME = "void_vera"

file_path = r"D:\void\data\router_knowledge_base.txt"

# --- 2. Load and Chunk the Data ---
try:
    with open(file_path, 'r', encoding='utf-8') as f:
        full_text = f.read()
except FileNotFoundError:
    print("Error: router_knowledge_base.txt not found. Please make sure the file is in the correct location.")
    exit()

# Use separator for chunking (already split as per your file)
text_chunks = full_text.split('==================================')
text_chunks = [chunk.strip() for chunk in text_chunks if chunk.strip()]
print(f"Loaded and split into {len(text_chunks)} chunks.")

# --- 3. Initialize Embedding Model ---
print("Loading sentence-transformer model...")
model = SentenceTransformer('all-MiniLM-L6-v2')
print("Model loaded.")

# --- 4. Connect to Qdrant ---
client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
print("Connected to Qdrant.")

# --- 5. Create or Get Collection ---
vector_size = model.get_sentence_embedding_dimension()

try:
    client.get_collection(collection_name=COLLECTION_NAME)
    print(f"Collection '{COLLECTION_NAME}' exists. Adding new data.")
except Exception:
    print(f"Collection '{COLLECTION_NAME}' not found. Creating new one.")
    client.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE)
    )

# --- 6. Embed and Upsert ---
print(f"Embedding and uploading {len(text_chunks)} chunks to Qdrant...")
batch_size = 100
for i in range(0, len(text_chunks), batch_size):
    batch_to_upload = text_chunks[i:i + batch_size]
    points_to_upload = [
        models.PointStruct(
            id=str(uuid.uuid4()),
            vector=model.encode(chunk).tolist(),
            payload={"text": chunk}
        )
        for chunk in batch_to_upload
    ]
    client.upload_points(
        collection_name=COLLECTION_NAME,
        points=points_to_upload,
        wait=True
    )
    print(f"  - Uploaded batch {i // batch_size + 1}/{(len(text_chunks) + batch_size - 1) // batch_size}")

print("\n🎉 All content for 'void_vera' has been successfully embedded and added to Qdrant!")

# --- 7. Test Search Example ---
print("\n--- Running a test search ---")
test_query = "How does the router knowledge base organize troubleshooting steps?"
query_vector = model.encode(test_query).tolist()
search_results = client.search(
    collection_name=COLLECTION_NAME,
    query_vector=query_vector,
    limit=3
)

print(f"Query: '{test_query}'")
print("\nTop 3 relevant chunks:")
for i, result in enumerate(search_results):
    print(f"--- Result {i+1} (Score: {result.score:.4f}) ---")
    print(result.payload['text'])
    print("\n")