From fcaac9c0f102b06ace2cf77a3855f4deb80949d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Fri, 19 Dec 2025 11:12:29 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=B9=E7=94=A8=20httpx=20=E5=AE=9E=E7=8E=B0?= =?UTF-8?q?=E6=89=80=E6=9C=89=20client=20=E4=BA=A4=E4=BA=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- requirements.txt | 3 +- src/pipeline/core/es.py | 191 +++++++++++++++++++++++++++++++++ src/pipeline/core/es_client.py | 125 --------------------- src/pipeline/core/llm.py | 77 +++++++++---- src/pipeline/core/nodes.py | 10 +- src/tests/test_nodes.py | 10 +- 6 files changed, 258 insertions(+), 158 deletions(-) create mode 100644 src/pipeline/core/es.py delete mode 100644 src/pipeline/core/es_client.py diff --git a/requirements.txt b/requirements.txt index e268aef..a6f17ac 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,5 +8,4 @@ scikit-learn aiofiles pillow loguru -aiohttp -elasticsearch>=8.0.0,<9.0.0 \ No newline at end of file +httpx \ No newline at end of file diff --git a/src/pipeline/core/es.py b/src/pipeline/core/es.py new file mode 100644 index 0000000..06e4468 --- /dev/null +++ b/src/pipeline/core/es.py @@ -0,0 +1,191 @@ +import httpx +from datetime import datetime, timezone +from src.pipeline.config import config +from src.pipeline.core.utils import logger + + +class AsyncES: + def __init__( + self, + timeout: float = 30.0, + max_connections: int = 50, + max_keepalive: int = 10, + ): + self.base = f"{config['es_host']}:{config['es_port']}".rstrip("/") + self.auth = (config["es_user"], config["es_password"]) + + self.client = httpx.AsyncClient( + http2=False, + trust_env=False, + timeout=httpx.Timeout(timeout), + limits=httpx.Limits( + max_connections=max_connections, + max_keepalive_connections=max_keepalive, + ), + auth=self.auth, + headers={ + "Content-Type": "application/json", + "Accept": "application/json", + }, + ) + + self.mapping = { + "settings": { + "analysis": { + "analyzer": { + "ik_smart": { + "tokenizer": "ik_smart", + } + } + } + }, + "mappings": { + "properties": { + "type": {"type": "keyword"}, + "doc_id": {"type": "keyword"}, + "title": {"type": "text", "analyzer": "ik_smart"}, + "created_at": {"type": "date"}, + "content": {"type": "text", "analyzer": "ik_smart"}, + "embedding": { + "type": "dense_vector", + "dims": config['embedding_dims'], + "index": True, + "similarity": "cosine", + }, + } + }, + } + + # ------------------------- + # Index + # ------------------------- + + async def create_index(self, index: str) -> bool: + url = f"{self.base}/{index}" + try: + resp = await self.client.put(url, json=self.mapping) + if resp.status_code in (200, 201): + return True + if resp.status_code == 400 and "resource_already_exists" in resp.text: + return True + + logger.error(f"Create index failed: {resp.text}") + return False + + except Exception: + logger.exception("Create index error") + return False + + # ------------------------- + # Index document + # ------------------------- + + 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 = None, + ): + 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(), + } + + url = f"{self.base}/{index}/_doc/{es_id}" + + try: + resp = await self.client.put(url, json=doc) + resp.raise_for_status() + except Exception: + logger.exception("Index document failed") + + # ------------------------- + # Hybrid search + # ------------------------- + + async def hybrid_search_es( + self, + index: str, + query_text: str, + query_vector: list, + top_k: int = 5, + doc_type: str | None = None, + ): + 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": { + "source": ("cosineSimilarity(params.query_vector, 'embedding') + 1.0"), + "params": {"query_vector": query_vector}, + } + } + } + ], + "boost_mode": "sum", + } + }, + } + + url = f"{self.base}/{index}/_search" + + resp = await self.client.post(url, json=body) + resp.raise_for_status() + + hits = resp.json()["hits"]["hits"] + + return [ + { + "es_id": h["_id"], + "doc_id": h["_source"].get("doc_id"), + "title": h["_source"].get("title"), + "type": h["_source"].get("type"), + "created_at": h["_source"].get("created_at"), + "score": h["_score"], + "content": h["_source"].get("content"), + } + for h in hits + ] + + + async def close(self): + await self.client.aclose() + + +client: AsyncES | None = None + + +async def init_client(): + global client + client = AsyncES() + +async def close_client(): + await client.close() diff --git a/src/pipeline/core/es_client.py b/src/pipeline/core/es_client.py deleted file mode 100644 index 97f2502..0000000 --- a/src/pipeline/core/es_client.py +++ /dev/null @@ -1,125 +0,0 @@ -from codecs import ignore_errors -from datetime import datetime, timezone -from venv import logger -from src.pipeline.config import config -from elasticsearch import AsyncElasticsearch, ApiError - - -class AsyncES: - def __init__(self): - self.base = f"{config['es_host']}:{config['es_port']}" - 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": 1024, # 必须和 embedding 模型的纬度一样 - "index": True, - "similarity": "cosine", - }, - }, - }, - } - - async def create_index(self, index: str): - """ - 创建 index(知识库) - """ - 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, - ): - """ - 写入文档(普通文本) - """ - 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): - """ - 在 Elasticsearch 内执行 BM25 + 向量混合检索 - """ - # 基础查询:中文 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 得分 + 向量相似度相加 - } - }, - } - - 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/llm.py b/src/pipeline/core/llm.py index aab22d6..5d1082a 100644 --- a/src/pipeline/core/llm.py +++ b/src/pipeline/core/llm.py @@ -1,29 +1,60 @@ +import httpx from src.pipeline.config import config -import aiohttp +from src.pipeline.core.utils import logger -_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) +class AsyncLLm: + def __init__( + self, + timeout: float = 30.0, + max_connections: int = 100, + max_keepalive: int = 20, + ): + self.embedding_api = config["embedding_api_host"].rstrip("/") + "/embeddings" + self.embedding_model = config["embedding_model"] + self.api_key = config["embedding_api_key"] + logger.debug(self.embedding_api) + self.embedding_client = httpx.AsyncClient( + http2=False, + trust_env=False, + timeout=httpx.Timeout(timeout), + limits=httpx.Limits( + max_connections=max_connections, + max_keepalive_connections=max_keepalive, + ), + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self.api_key}", + }, ) - 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 [] + async def embedding(self, text: str) -> list[float]: + try: + resp = await self.embedding_client.post( + self.embedding_api, + json={"model": self.embedding_model, "input": text}, + ) - data = await resp.json() - return data["data"][0]["embedding"] + resp.raise_for_status() + data = resp.json() + return data["data"][0]["embedding"] + + except httpx.HTTPStatusError as e: + logger.error(e) + logger.error(f"Embedding HTTP error: {e.response.text}") + except Exception as e: + logger.exception("Embedding request failed") + + return [] + + async def close(self): + await self.embedding_client.aclose() + +client: AsyncLLm | None = None + +async def init_client(): + global client + client = AsyncLLm() + +async def close_client(): + await client.close() diff --git a/src/pipeline/core/nodes.py b/src/pipeline/core/nodes.py index 19dae23..5d5d871 100644 --- a/src/pipeline/core/nodes.py +++ b/src/pipeline/core/nodes.py @@ -1,8 +1,8 @@ import uuid from src.pipeline.core.pocket_flow import AsyncBatchNode 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 +from src.pipeline.core import llm +from src.pipeline.core import es import re @@ -49,7 +49,7 @@ class EmbeddingDocumentsNode(AsyncBatchNode): :param document: {text, file_name} """ logger.debug(f"开始 embedding: {document["text"].strip()[:10]}...") - res = {**document, "embedding": await get_embedding(document["text"])} + res = {**document, "embedding": await llm.client.embedding(document["text"])} logger.debug(f"结束 embedding: {document["text"].strip()[:10]}...") return res @@ -95,7 +95,7 @@ class ReadDocumentNode(AsyncBatchNode): class WriteToElasticsearchNode(AsyncBatchNode): async def prep_async(self, shared): index = shared["index"] - await es_client.create_index(index) + await es.client.create_index(index) return [ { "index": index, @@ -110,7 +110,7 @@ class WriteToElasticsearchNode(AsyncBatchNode): ] async def exec_async(self, prep_res): - await es_client.add_doc(**prep_res) + await es.client.add_doc(**prep_res) return True async def post_async(self, shared, prep_res, exec_res): diff --git a/src/tests/test_nodes.py b/src/tests/test_nodes.py index 354e479..43a4df8 100644 --- a/src/tests/test_nodes.py +++ b/src/tests/test_nodes.py @@ -3,19 +3,20 @@ import json 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 +from src.pipeline.core import llm, es @pytest.mark.asyncio async def test_embedding(): + await llm.init_client() + await es.init_client() + print("\n\ntest_embedding:\n") shared = { "files": [ - # "./files/带图片的.pdf", - # "./files/大白智问-API接入文档-V1.2.2.pdf", "./files/山海经01.txt", - # "./files/deepsearch状态.txt", ], "documents": [], "index": "test_kb", @@ -31,3 +32,6 @@ async def test_embedding(): await flow.run_async(shared) logger.debug(json.dumps([{**x, "embedding":x["embedding"][:4]} for x in shared["documents"]], indent=4, ensure_ascii=False)) + + await llm.close_client() + await es.close_client()