feat: es、qd 客户端

This commit is contained in:
李如威 2025-12-15 17:32:35 +08:00
parent 9776f3c791
commit 8132b1d515
5 changed files with 131 additions and 11 deletions

View File

@ -1,17 +1,29 @@
networks:
ai-pipeline:
driver: bridge
volumes:
es-data:
qdrant-storage:
services:
elasticsearch:
image: elasticsearch:9.2.2
image: docker.elastic.co/elasticsearch/elasticsearch:9.2.2
container_name: elasticsearch
user: "1000:1000"
environment:
- discovery.type=single-node # 单节点模式
- xpack.security.enabled=true # 开启安全(才可以设密码)
- ELASTIC_PASSWORD=12345 # 设置密码(重要)
- ES_JAVA_OPTS=-Xms1g -Xmx1g
- xpack.security.enabled=true # 开启安全
- ELASTIC_PASSWORD=12345 # 设置密码
# - xpack.security.http.ssl.enabled=false # 关闭 HTTPS
# - ES_JAVA_OPTS=-Xms1g -Xmx1g # JVM 初始堆内存大小1GB最大堆内存大小1GB
ports:
- "9210:9200"
volumes:
- ../datas/es-data:/usr/share/elasticsearch/data
- es-data:/usr/share/elasticsearch/data
restart: unless-stopped
networks:
- ai-pipeline
qdrant:
image: qdrant/qdrant:latest
@ -22,5 +34,7 @@ services:
- "6333:6333" # HTTP
- "6334:6334" # gRPC
volumes:
- ../datas/qdrant-storage:/qdrant/storage
restart: unless-stopped
- qdrant-storage:/qdrant/storage
restart: unless-stopped
networks:
- ai-pipeline

View File

@ -1,6 +1,7 @@
from mimetypes import init
from typing import TypedDict
from dotenv import load_dotenv
from src.pipeline.core.utils import logger
import os
load_dotenv()
@ -17,9 +18,12 @@ class Config(TypedDict):
embedding_api_host: str
embedding_model: str
es_host: str
es_key: str
es_port: int
es_user: str
es_password: str
qdrant_host: str
qdrant_key: str
qdrant_port: int
qdrant_api_key: str
def _read_config() -> Config:
@ -35,9 +39,13 @@ def _read_config() -> Config:
"embedding_api_key": os.getenv("EMBEDDING_API_KEY"),
"embedding_model": os.getenv("EMBEDDING_MODEL"),
"es_host": os.getenv("ES_HOST"),
"es_key": os.getenv("ES_KEY") or "",
"es_port": int(os.getenv("ES_PORT")),
"es_host": os.getenv("ES_USER") or "elastic",
"es_password": os.getenv("ES_PASSWORD") or "",
"qdrant_host": os.getenv("QDRANT_HOST"),
"qdrant_host": os.getenv("QDRANT_KEY") or "",
"qdrant_port": int(os.getenv("QDRANT_PORT")),
"qdrant_api_key": os.getenv("QDRANT_API_KEY") or "",
}
config = _read_config()
logger.debug("创建全局: config")

View File

@ -0,0 +1,45 @@
import httpx
from pipeline.config import config
class AsyncES:
def __init__(self):
self.base = f"{config['es_host']}:{config['es_port']}"
self.auth = (config["es_user"], config['es_password'])
async def create_index(self, index: str):
"""
创建 index知识库
"""
async with httpx.AsyncClient() as client:
resp = await client.put(
f"{self.base}/{index}",
auth=self.auth,
json={
"settings": {"analysis": {"analyzer": {"default": {"type": "standard"}}}},
"mappings": {
"properties": {
"text": {"type": "text"},
"kb": {"type": "keyword"},
}
},
},
)
return resp.json()
async def add_doc(self, index: str, text: str, kb: str):
"""
写入文档普通文本
"""
async with httpx.AsyncClient() as client:
resp = await client.post(f"{self.base}/{index}/_doc", auth=self.auth, json={"text": text, "kb": kb})
return resp.json()
async def bm25_search(self, index: str, query: str, top_k: int = 10):
"""
BM25 搜索
"""
async with httpx.AsyncClient() as client:
resp = await client.post(f"{self.base}/{index}/_search", auth=self.auth, json={"size": top_k, "query": {"match": {"text": query}}})
return resp.json()
es_client = AsyncES

View File

@ -0,0 +1,53 @@
from qdrant_client import AsyncQdrantClient
from qdrant_client.models import (
VectorParams,
Distance,
PointStruct,
)
from pipeline.config import config
class AsyncQD:
def __init__(self):
self.client = AsyncQdrantClient(
host=config["qdrant_host"],
port=config["qdrant_port"],
api_key=config["qdrant_api_key"],
)
async def ensure_collection(self, name: str, dim: int):
"""
如果 collection 不存在则创建
适用于多知识库每个知识库一个 collection
"""
try:
await self.client.get_collection(name)
except Exception:
await self.client.recreate_collection(
collection_name=name,
vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
)
async def upsert_vectors(self, collection: str, vectors: list, payloads: list):
"""
写入向量id 自增或自行维护
- vectors: [[float...], ...]
- payloads: [{"text": "...", "kb": "..."}]
"""
points = [PointStruct(id=i, vector=vectors[i], payload=payloads[i]) for i in range(len(vectors))]
await self.client.upsert(collection_name=collection, points=points)
async def vector_search(self, collection: str, query_vector: list, top_k: int = 5):
"""
查询向量相似度
"""
result = await self.client.search(
collection_name=collection,
query_vector=query_vector,
limit=top_k,
)
return result
qd_client = AsyncQD()