改用 httpx 实现所有 client 交互
This commit is contained in:
parent
4345dfb9cb
commit
fcaac9c0f1
|
|
@ -8,5 +8,4 @@ scikit-learn
|
||||||
aiofiles
|
aiofiles
|
||||||
pillow
|
pillow
|
||||||
loguru
|
loguru
|
||||||
aiohttp
|
httpx
|
||||||
elasticsearch>=8.0.0,<9.0.0
|
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -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()
|
|
||||||
|
|
@ -1,29 +1,60 @@
|
||||||
|
import httpx
|
||||||
from src.pipeline.config import config
|
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):
|
class AsyncLLm:
|
||||||
global _embedding_session
|
def __init__(
|
||||||
|
self,
|
||||||
if _embedding_session is None or _embedding_session.closed:
|
timeout: float = 30.0,
|
||||||
_embedding_session = aiohttp.ClientSession(
|
max_connections: int = 100,
|
||||||
timeout=aiohttp.ClientTimeout(total=timeout)
|
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(
|
async def embedding(self, text: str) -> list[float]:
|
||||||
config["embedding_api_host"],
|
try:
|
||||||
json={
|
resp = await self.embedding_client.post(
|
||||||
"model": config["embedding_model"],
|
self.embedding_api,
|
||||||
"input": text,
|
json={"model": self.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()
|
resp.raise_for_status()
|
||||||
return data["data"][0]["embedding"]
|
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()
|
||||||
|
|
|
||||||
|
|
@ -1,8 +1,8 @@
|
||||||
import uuid
|
import uuid
|
||||||
from src.pipeline.core.pocket_flow import AsyncBatchNode
|
from src.pipeline.core.pocket_flow import AsyncBatchNode
|
||||||
from src.pipeline.core.utils import fixed_size_chunk, load_document, logger
|
from src.pipeline.core.utils import fixed_size_chunk, load_document, logger
|
||||||
from src.pipeline.core.llm import get_embedding
|
from src.pipeline.core import llm
|
||||||
from src.pipeline.core.es_client import es_client
|
from src.pipeline.core import es
|
||||||
import re
|
import re
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -49,7 +49,7 @@ class EmbeddingDocumentsNode(AsyncBatchNode):
|
||||||
:param document: {text, file_name}
|
:param document: {text, file_name}
|
||||||
"""
|
"""
|
||||||
logger.debug(f"开始 embedding: {document["text"].strip()[:10]}...")
|
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]}...")
|
logger.debug(f"结束 embedding: {document["text"].strip()[:10]}...")
|
||||||
return res
|
return res
|
||||||
|
|
||||||
|
|
@ -95,7 +95,7 @@ class ReadDocumentNode(AsyncBatchNode):
|
||||||
class WriteToElasticsearchNode(AsyncBatchNode):
|
class WriteToElasticsearchNode(AsyncBatchNode):
|
||||||
async def prep_async(self, shared):
|
async def prep_async(self, shared):
|
||||||
index = shared["index"]
|
index = shared["index"]
|
||||||
await es_client.create_index(index)
|
await es.client.create_index(index)
|
||||||
return [
|
return [
|
||||||
{
|
{
|
||||||
"index": index,
|
"index": index,
|
||||||
|
|
@ -110,7 +110,7 @@ class WriteToElasticsearchNode(AsyncBatchNode):
|
||||||
]
|
]
|
||||||
|
|
||||||
async def exec_async(self, prep_res):
|
async def exec_async(self, prep_res):
|
||||||
await es_client.add_doc(**prep_res)
|
await es.client.add_doc(**prep_res)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
async def post_async(self, shared, prep_res, exec_res):
|
async def post_async(self, shared, prep_res, exec_res):
|
||||||
|
|
|
||||||
|
|
@ -3,19 +3,20 @@ import json
|
||||||
from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode, WriteToElasticsearchNode
|
from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode, WriteToElasticsearchNode
|
||||||
from src.pipeline.core.pocket_flow import AsyncFlow
|
from src.pipeline.core.pocket_flow import AsyncFlow
|
||||||
from src.pipeline.core.utils import logger
|
from src.pipeline.core.utils import logger
|
||||||
|
from src.pipeline.core import llm, es
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_embedding():
|
async def test_embedding():
|
||||||
|
|
||||||
|
await llm.init_client()
|
||||||
|
await es.init_client()
|
||||||
|
|
||||||
print("\n\ntest_embedding:\n")
|
print("\n\ntest_embedding:\n")
|
||||||
|
|
||||||
shared = {
|
shared = {
|
||||||
"files": [
|
"files": [
|
||||||
# "./files/带图片的.pdf",
|
|
||||||
# "./files/大白智问-API接入文档-V1.2.2.pdf",
|
|
||||||
"./files/山海经01.txt",
|
"./files/山海经01.txt",
|
||||||
# "./files/deepsearch状态.txt",
|
|
||||||
],
|
],
|
||||||
"documents": [],
|
"documents": [],
|
||||||
"index": "test_kb",
|
"index": "test_kb",
|
||||||
|
|
@ -31,3 +32,6 @@ async def test_embedding():
|
||||||
await flow.run_async(shared)
|
await flow.run_async(shared)
|
||||||
|
|
||||||
logger.debug(json.dumps([{**x, "embedding":x["embedding"][:4]} for x in shared["documents"]], indent=4, ensure_ascii=False))
|
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()
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue