diff --git a/src/pipeline/core/es_client.py b/src/pipeline/core/es_client.py index 1718597..d8707eb 100644 --- a/src/pipeline/core/es_client.py +++ b/src/pipeline/core/es_client.py @@ -1,45 +1,102 @@ -import httpx +from datetime import datetime, timezone from pipeline.config import config +from elasticsearch import AsyncElasticsearch + class AsyncES: def __init__(self): 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): """ 创建 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() + is_exists = await self.client.indices.exists(index) + if is_exists: + return True + else: + await self.client.indices.create(index=index, body=self.mapping) + return True - 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: - resp = await client.post(f"{self.base}/{index}/_doc", auth=self.auth, json={"text": text, "kb": kb}) - return resp.json() + if create_at is None: + create_at = datetime.now(timezone.utc) + 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: - resp = await client.post(f"{self.base}/{index}/_search", auth=self.auth, json={"size": top_k, "query": {"match": {"text": query}}}) - return resp.json() + # 基础查询:中文 BM25 匹配 content 或 title + must_clauses = [{"multi_match": {"query": query_text, "fields": ["content", "title"]}}] + # 可选类型过滤 + 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() diff --git a/src/pipeline/core/qd_client.py b/src/pipeline/core/qd_client.py deleted file mode 100644 index 413280a..0000000 --- a/src/pipeline/core/qd_client.py +++ /dev/null @@ -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() diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index b2af5be..247fb23 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -12,9 +12,9 @@ async def test_embedding(): shared = { "files": [ - "./files/带图片的.pdf", + # "./files/带图片的.pdf", # "./files/大白智问-API接入文档-V1.2.2.pdf", - # "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", + "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", # "./files/deepsearch状态.txt", ] }