#!/usr/bin/env python3
"""
Qdrant 与 Ollama 嵌入模型集成示例
此脚本演示如何使用 Ollama 作为嵌入模型提供者与 Qdrant 向量数据库集成
"""

from langchain_ollama import OllamaEmbeddings
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
import os

def main():
    # 1. 初始化 Ollama 嵌入模型
    # 使用 nomic-embed-text 模型，这是 Ollama 推荐的嵌入模型
    print("初始化 Ollama 嵌入模型...")
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url="http://localhost:11434"  # Ollama 默认地址
    )
    
    # 2. 初始化 Qdrant 客户端
    print("连接到 Qdrant 数据库...")
    client = QdrantClient(
        url="http://localhost:6333",  # Qdrant 默认地址
        api_key="313131"  # 从之前查看的配置中获取的 API 密钥
    )
    
    # 3. 创建集合（如果不存在）
    collection_name = "ollama_integration_test"
    print(f"创建或检查集合: {collection_name}")
    
    # 首先检查集合是否已存在
    collections = client.get_collections().collections
    collection_exists = any(collection.name == collection_name for collection in collections)
    
    if not collection_exists:
        # 创建新集合
        # 首先获取嵌入模型的维度
        sample_embedding = embeddings.embed_query("sample text")
        vector_size = len(sample_embedding)
        
        client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(
                size=vector_size,
                distance=Distance.COSINE
            )
        )
        print(f"已创建新集合，向量维度: {vector_size}")
    else:
        print("集合已存在")
    
    # 4. 准备示例数据
    documents = [
        "Qdrant 是一个高性能的向量搜索引擎",
        "Ollama 是一个本地运行大语言模型的工具",
        "向量数据库用于存储和检索高维向量",
        "嵌入模型将文本转换为数值向量表示"
    ]
    
    metadata = [
        {"source": "qdrant_docs", "category": "database"},
        {"source": "ollama_docs", "category": "llm"},
        {"source": "vector_db_docs", "category": "database"},
        {"source": "embedding_docs", "category": "ml"}
    ]
    
    # 5. 使用 Ollama 生成嵌入并存储到 Qdrant
    print("生成嵌入并存储到 Qdrant...")
    points = []
    
    for idx, (doc, meta) in enumerate(zip(documents, metadata)):
        # 使用 Ollama 生成嵌入
        embedding = embeddings.embed_query(doc)
        
        # 创建 Qdrant 点
        point = PointStruct(
            id=idx,
            vector=embedding,
            payload={
                "text": doc,
                "metadata": meta
            }
        )
        points.append(point)
    
    # 上传点到 Qdrant
    client.upsert(
        collection_name=collection_name,
        points=points
    )
    print(f"已上传 {len(points)} 个文档到 Qdrant")
    
    # 6. 执行相似性搜索
    query = "什么是向量数据库？"
    print(f"\n执行搜索查询: '{query}'")
    
    # 使用 Ollama 生成查询嵌入
    query_embedding = embeddings.embed_query(query)
    
    # 在 Qdrant 中搜索
    search_result = client.search(
        collection_name=collection_name,
        query_vector=query_embedding,
        limit=2
    )
    
    # 7. 显示搜索结果
    print("\n搜索结果:")
    for i, hit in enumerate(search_result, 1):
        print(f"{i}. {hit.payload['text']} (得分: {hit.score:.4f})")
        print(f"   元数据: {hit.payload['metadata']}")
    
    print("\n集成测试完成！")

if __name__ == "__main__":
    main()