117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
#!/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() |