feat: es_client
This commit is contained in:
parent
78965ded3e
commit
63480e312e
|
|
@ -1,45 +1,102 @@
|
||||||
import httpx
|
from datetime import datetime, timezone
|
||||||
from pipeline.config import config
|
from pipeline.config import config
|
||||||
|
from elasticsearch import AsyncElasticsearch
|
||||||
|
|
||||||
|
|
||||||
class AsyncES:
|
class AsyncES:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.base = f"{config['es_host']}:{config['es_port']}"
|
self.base = f"{config['es_host']}:{config['es_port']}"
|
||||||
self.auth = (config["es_user"], config['es_password'])
|
self.auth = (config["es_user"], config["es_password"])
|
||||||
|
self.client = AsyncElasticsearch(self.base, basic_auth=self.auth)
|
||||||
|
self.mapping = {
|
||||||
|
"settings": {
|
||||||
|
"analysis": {
|
||||||
|
"analyzer": {
|
||||||
|
"ik_smart": {
|
||||||
|
"tokenizer": "ik_smart",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"mappings": {
|
||||||
|
"properties": {
|
||||||
|
"type": {"type": "keyword"}, # 分类/类型
|
||||||
|
"doc_id": {"type": "keyword"}, # 业务文档 ID,便于聚合
|
||||||
|
"title": {"type": "text", "analyzer": "ik_smart"}, # 文档标题
|
||||||
|
"created_at": {"type": "date"}, # 创建时间
|
||||||
|
"content": {"type": "text", "analyzer": "ik_smart"}, # 切片内容
|
||||||
|
"embedding": { # 向量字段
|
||||||
|
"type": "dense_vector",
|
||||||
|
"dims": 768,
|
||||||
|
"index": True,
|
||||||
|
"similarity": "cosine",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
async def create_index(self, index: str):
|
async def create_index(self, index: str):
|
||||||
"""
|
"""
|
||||||
创建 index(知识库)
|
创建 index(知识库)
|
||||||
"""
|
"""
|
||||||
async with httpx.AsyncClient() as client:
|
is_exists = await self.client.indices.exists(index)
|
||||||
resp = await client.put(
|
if is_exists:
|
||||||
f"{self.base}/{index}",
|
return True
|
||||||
auth=self.auth,
|
else:
|
||||||
json={
|
await self.client.indices.create(index=index, body=self.mapping)
|
||||||
"settings": {"analysis": {"analyzer": {"default": {"type": "standard"}}}},
|
return True
|
||||||
"mappings": {
|
|
||||||
"properties": {
|
|
||||||
"text": {"type": "text"},
|
|
||||||
"kb": {"type": "keyword"},
|
|
||||||
}
|
|
||||||
},
|
|
||||||
},
|
|
||||||
)
|
|
||||||
return resp.json()
|
|
||||||
|
|
||||||
async def add_doc(self, index: str, text: str, kb: str):
|
async def add_doc(self, index: str, es_id: str, doc_id: str, content: str, embedding: list, title: str = "", doc_type: str = "", created_at: datetime = None):
|
||||||
"""
|
"""
|
||||||
写入文档(普通文本)
|
写入文档(普通文本)
|
||||||
"""
|
"""
|
||||||
async with httpx.AsyncClient() as client:
|
if create_at is None:
|
||||||
resp = await client.post(f"{self.base}/{index}/_doc", auth=self.auth, json={"text": text, "kb": kb})
|
create_at = datetime.now(timezone.utc)
|
||||||
return resp.json()
|
doc = {"doc_id": doc_id, "title": title, "type": doc_type, "content": content, "embedding": embedding, "created_at": created_at.isoformat()}
|
||||||
|
await self.client.index(index=index, id=es_id, document=doc)
|
||||||
|
|
||||||
async def bm25_search(self, index: str, query: str, top_k: int = 10):
|
async def hybrid_search_es(self, index: str, query_text: str, query_vector: list, top_k: int = 5, doc_type: str = None):
|
||||||
"""
|
"""
|
||||||
BM25 搜索
|
在 Elasticsearch 内执行 BM25 + 向量混合检索
|
||||||
"""
|
"""
|
||||||
async with httpx.AsyncClient() as client:
|
# 基础查询:中文 BM25 匹配 content 或 title
|
||||||
resp = await client.post(f"{self.base}/{index}/_search", auth=self.auth, json={"size": top_k, "query": {"match": {"text": query}}})
|
must_clauses = [{"multi_match": {"query": query_text, "fields": ["content", "title"]}}]
|
||||||
return resp.json()
|
# 可选类型过滤
|
||||||
|
if doc_type:
|
||||||
|
must_clauses.append({"term": {"type": doc_type}})
|
||||||
|
body = {
|
||||||
|
"size": top_k,
|
||||||
|
"query": {
|
||||||
|
"function_score": {
|
||||||
|
"query": {"bool": {"must": must_clauses}},
|
||||||
|
"functions": [
|
||||||
|
{
|
||||||
|
"script_score": {
|
||||||
|
"script": {
|
||||||
|
# cosineSimilarity 返回 [-1,1],加 1 保证非负
|
||||||
|
"source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
|
||||||
|
"params": {"query_vector": query_vector},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"boost_mode": "sum", # BM25 得分 + 向量相似度相加
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
es_client = AsyncES
|
res = await self.client.search(index=index, body=body)
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"es_id": hit["_id"],
|
||||||
|
"doc_id": hit["_source"].get("doc_id"),
|
||||||
|
"title": hit["_source"].get("title"),
|
||||||
|
"type": hit["_source"].get("type"),
|
||||||
|
"created_at": hit["_source"].get("created_at"),
|
||||||
|
"score": hit["_score"],
|
||||||
|
"content": hit["_source"].get("content"),
|
||||||
|
}
|
||||||
|
for hit in res["hits"]["hits"]
|
||||||
|
]
|
||||||
|
|
||||||
|
es_client = AsyncES()
|
||||||
|
|
|
||||||
|
|
@ -1,53 +0,0 @@
|
||||||
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()
|
|
||||||
|
|
@ -12,9 +12,9 @@ async def test_embedding():
|
||||||
|
|
||||||
shared = {
|
shared = {
|
||||||
"files": [
|
"files": [
|
||||||
"./files/带图片的.pdf",
|
# "./files/带图片的.pdf",
|
||||||
# "./files/大白智问-API接入文档-V1.2.2.pdf",
|
# "./files/大白智问-API接入文档-V1.2.2.pdf",
|
||||||
# "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
|
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
|
||||||
# "./files/deepsearch状态.txt",
|
# "./files/deepsearch状态.txt",
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue