rag-pipeline-docker/scripts/test_qdrant_search.py

#!/usr/bin/env python3
"""
test_qdrant_search.py — Тестовый скрипт поиска по Qdrant.
Берёт текст, получает эмбеддинг через Ollama, ищет в Qdrant.

Использование: python3 scripts/test_qdrant_search.py --query "текст"
"""
import argparse
import json
import sys
import urllib.request
from pathlib import Path

OLLAMA_URL = "http://localhost:11434/api/embeddings"
EMBEDDING_MODEL = "nomic-embed-text:latest"
QDRANT_SEARCH_URL = "http://localhost:6333/collections/knowledge_base/points/search"
TOP_K = 5


def get_embedding(text: str) -> list[float]:
    payload = json.dumps({"model": EMBEDDING_MODEL, "prompt": text}).encode()
    req = urllib.request.Request(OLLAMA_URL, data=payload,
                                  headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read())
    return data.get("embedding")


def search_qdrant(vector: list[float], top_k: int = TOP_K) -> list[dict]:
    payload = json.dumps({
        "vector": {"name": "dense", "vector": vector},
        "limit": top_k,
        "with_payload": True,
        "with_vector": False,
    }).encode()
    req = urllib.request.Request(QDRANT_SEARCH_URL, data=payload,
                                  headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=15) as resp:
        data = json.loads(resp.read())
    return data.get("result", [])


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--query", "-q", required=True)
    args = parser.parse_args()

    vector = get_embedding(args.query)
    print(f"Embedding: {len(vector)} dims")

    results = search_qdrant(vector)
    print(f"Results: {len(results)}")
    for i, r in enumerate(results, 1):
        payload = r.get("payload", {})
        text = (payload.get("text") or payload.get("content", ""))[:200]
        print(f"  #{i} score={r['score']:.4f} | {payload.get('source','?')} | {text}...")


if __name__ == "__main__":
    main()