From 4345dfb9cb19f60374e44c3200dee6378b0cdfbb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Thu, 18 Dec 2025 18:00:37 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20es=20=E5=AE=B9=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docker-compose.yml | 19 +++----------- requirements.txt | 6 ++--- search.es | 0 src/pipeline/config.py | 4 ++- src/pipeline/core/es_client.py | 47 +++++++++++++++++++++++++--------- src/pipeline/core/llm.py | 47 +++++++++++++++++++--------------- src/pipeline/core/nodes.py | 45 +++++++++++++++++++++++++++++--- src/tests/test_nodes.py | 11 +++++--- 8 files changed, 121 insertions(+), 58 deletions(-) create mode 100644 search.es diff --git a/docker-compose.yml b/docker-compose.yml index 27a1311..6c5effa 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -17,24 +17,13 @@ services: - ELASTIC_PASSWORD=12345 # 设置密码 # - xpack.security.http.ssl.enabled=false # 关闭 HTTPS # - ES_JAVA_OPTS=-Xms1g -Xmx1g # JVM 初始堆内存大小,1GB,最大堆内存大小,1GB + # elasticvue 配置 + - http.cors.enabled=true + - http.cors.allow-origin=http://localhost:8210 ports: - "9210:9200" volumes: - es-data:/usr/share/elasticsearch/data restart: unless-stopped networks: - - ai-pipeline - - qdrant: - image: qdrant/qdrant:latest - container_name: qdrant - environment: - QDRANT__SERVICE__API_KEY: 12345 - ports: - - "6333:6333" # HTTP - - "6334:6334" # gRPC - volumes: - - qdrant-storage:/qdrant/storage - restart: unless-stopped - networks: - - ai-pipeline + - ai-pipeline \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 2499b92..e268aef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,5 @@ fastapi uvicorn[standard] -httpx[socks] -qdrant-client[httpx] python-dotenv pytest-asyncio pymupdf @@ -9,4 +7,6 @@ python-docx scikit-learn aiofiles pillow -loguru \ No newline at end of file +loguru +aiohttp +elasticsearch>=8.0.0,<9.0.0 \ No newline at end of file diff --git a/search.es b/search.es new file mode 100644 index 0000000..e69de29 diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 2f3130b..7dd43b0 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -17,6 +17,7 @@ class Config(TypedDict): embedding_api_key: str embedding_api_host: str embedding_model: str + embedding_dims: int es_host: str es_port: int es_user: str @@ -38,9 +39,10 @@ def _read_config() -> Config: "embedding_api_host": os.getenv("EMBEDDING_API_HOST"), "embedding_api_key": os.getenv("EMBEDDING_API_KEY"), "embedding_model": os.getenv("EMBEDDING_MODEL"), + "embedding_dims": int(os.getenv("EMBEDDING_DIMS")), "es_host": os.getenv("ES_HOST"), "es_port": int(os.getenv("ES_PORT")), - "es_host": os.getenv("ES_USER") or "elastic", + "es_user": os.getenv("ES_USER") or "elastic", "es_password": os.getenv("ES_PASSWORD") or "", "qdrant_host": os.getenv("QDRANT_HOST"), "qdrant_port": int(os.getenv("QDRANT_PORT")), diff --git a/src/pipeline/core/es_client.py b/src/pipeline/core/es_client.py index d8707eb..97f2502 100644 --- a/src/pipeline/core/es_client.py +++ b/src/pipeline/core/es_client.py @@ -1,6 +1,8 @@ +from codecs import ignore_errors from datetime import datetime, timezone -from pipeline.config import config -from elasticsearch import AsyncElasticsearch +from venv import logger +from src.pipeline.config import config +from elasticsearch import AsyncElasticsearch, ApiError class AsyncES: @@ -27,7 +29,7 @@ class AsyncES: "content": {"type": "text", "analyzer": "ik_smart"}, # 切片内容 "embedding": { # 向量字段 "type": "dense_vector", - "dims": 768, + "dims": 1024, # 必须和 embedding 模型的纬度一样 "index": True, "similarity": "cosine", }, @@ -39,20 +41,40 @@ class AsyncES: """ 创建 index(知识库) """ - is_exists = await self.client.indices.exists(index) - if is_exists: - return True - else: - await self.client.indices.create(index=index, body=self.mapping) + try: + await self.client.indices.create( + index=index, + body=self.mapping, + ) return True + except ApiError as e: + logger.error(e.message) + return False - 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 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, + ): """ 写入文档(普通文本) """ - 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()} + if created_at is None: + created_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 hybrid_search_es(self, index: str, query_text: str, query_vector: list, top_k: int = 5, doc_type: str = None): @@ -99,4 +121,5 @@ class AsyncES: for hit in res["hits"]["hits"] ] + es_client = AsyncES() diff --git a/src/pipeline/core/llm.py b/src/pipeline/core/llm.py index 71c1aba..aab22d6 100644 --- a/src/pipeline/core/llm.py +++ b/src/pipeline/core/llm.py @@ -1,22 +1,29 @@ from src.pipeline.config import config -import httpx +import aiohttp -async def get_embedding(text, timeout: int = 30): - try: - async with httpx.AsyncClient(timeout=timeout, http2=False, trust_env=False) as client: - url = config["embedding_api_host"] - body = { - "model": config["embedding_model"], - "input": text, - } - headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {config['embedding_api_key']}", - } - res = await client.post(url, headers=headers, json=body) - res.raise_for_status() - data = res.json() - return data["data"][0]["embedding"] - except Exception as e: - print(f"get_embedding[ERROR]: {e}") - return [] +_embedding_session: aiohttp.ClientSession | None = None + +async def get_embedding(text: str, timeout: int = 30): + global _embedding_session + + if _embedding_session is None or _embedding_session.closed: + _embedding_session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=timeout) + ) + + async with _embedding_session.post( + config["embedding_api_host"], + json={ + "model": config["embedding_model"], + "input": text, + }, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {config['embedding_api_key']}", + }, + ) as resp: + if resp.status >= 400: + return [] + + data = await resp.json() + return data["data"][0]["embedding"] diff --git a/src/pipeline/core/nodes.py b/src/pipeline/core/nodes.py index ada12d4..19dae23 100644 --- a/src/pipeline/core/nodes.py +++ b/src/pipeline/core/nodes.py @@ -1,7 +1,8 @@ import uuid from src.pipeline.core.pocket_flow import AsyncBatchNode -from src.pipeline.core.utils import fixed_size_chunk, load_document +from src.pipeline.core.utils import fixed_size_chunk, load_document, logger from src.pipeline.core.llm import get_embedding +from src.pipeline.core.es_client import es_client import re @@ -21,7 +22,15 @@ class ChunkDocumentsNode(AsyncBatchNode): text = re.sub(r" +", " ", text) # 去除首尾空格 text = text.strip() - return [{"text": x, "file_name": document["file_name"], "uuid": uuid.uuid4().hex} for x in fixed_size_chunk(text, chunk_size=10000)] + return [ + { + "text": x, + "file_name": document["file_name"], + "file_type": document["file_type"], + "uuid": uuid.uuid4().hex, + } + for x in fixed_size_chunk(text, chunk_size=500) + ] async def post_async(self, shared, prep_res, exec_res_list): all_chunks = [] @@ -39,7 +48,10 @@ class EmbeddingDocumentsNode(AsyncBatchNode): """ :param document: {text, file_name} """ - return {**document, "embedding": await get_embedding(document["text"])} + logger.debug(f"开始 embedding: {document["text"].strip()[:10]}...") + res = {**document, "embedding": await get_embedding(document["text"])} + logger.debug(f"结束 embedding: {document["text"].strip()[:10]}...") + return res async def post_async(self, shared, prep_res, exec_res_list): @@ -60,6 +72,7 @@ class ReadDocumentNode(AsyncBatchNode): "file_name": file_path.split("/")[-1], "text": document_text, "text_length": len(document_text), + "file_type": str(file_path.split(".")[-1]).strip().lower(), "status": "done", "message": "", } @@ -69,6 +82,7 @@ class ReadDocumentNode(AsyncBatchNode): "file_name": file_path.split("/")[-1], "text": "", "text_length": 0, + "file_type": str(file_path.split(".")[-1]).strip().lower(), "status": "error", "message": str(e), } @@ -76,3 +90,28 @@ class ReadDocumentNode(AsyncBatchNode): async def post_async(self, shared, prep_res, exec_res): shared["documents"] = exec_res return "default" + + +class WriteToElasticsearchNode(AsyncBatchNode): + async def prep_async(self, shared): + index = shared["index"] + await es_client.create_index(index) + return [ + { + "index": index, + "es_id": x["uuid"], + "doc_id": x["file_name"], + "embedding": x["embedding"], + "title": x["file_name"], + "doc_type": x["file_type"], + "content": x["text"], + } + for x in shared["documents"] + ] + + async def exec_async(self, prep_res): + await es_client.add_doc(**prep_res) + return True + + async def post_async(self, shared, prep_res, exec_res): + return "default" diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index 247fb23..354e479 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -1,6 +1,6 @@ import pytest import json -from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode +from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode, WriteToElasticsearchNode from src.pipeline.core.pocket_flow import AsyncFlow from src.pipeline.core.utils import logger @@ -14,15 +14,18 @@ async def test_embedding(): "files": [ # "./files/带图片的.pdf", # "./files/大白智问-API接入文档-V1.2.2.pdf", - "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", + "./files/山海经01.txt", # "./files/deepsearch状态.txt", - ] + ], + "documents": [], + "index": "test_kb", } readNode = ReadDocumentNode() chunkNode = ChunkDocumentsNode() embeddingNode = EmbeddingDocumentsNode() - readNode >> chunkNode >> embeddingNode + writeToESNode = WriteToElasticsearchNode() + readNode >> chunkNode >> embeddingNode >> writeToESNode flow = AsyncFlow(readNode) await flow.run_async(shared)