feat: es 容器

This commit is contained in:
李如威 2025-12-18 18:00:37 +08:00
parent 63480e312e
commit 4345dfb9cb
8 changed files with 121 additions and 58 deletions

View File

@ -17,6 +17,9 @@ 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:
@ -24,17 +27,3 @@ services:
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

View File

@ -1,7 +1,5 @@
fastapi
uvicorn[standard]
httpx[socks]
qdrant-client[httpx]
python-dotenv
pytest-asyncio
pymupdf
@ -10,3 +8,5 @@ scikit-learn
aiofiles
pillow
loguru
aiohttp
elasticsearch>=8.0.0,<9.0.0

0
search.es Normal file
View File

View File

@ -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")),

View File

@ -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()

View File

@ -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 = {
_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']}",
}
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}")
},
) as resp:
if resp.status >= 400:
return []
data = await resp.json()
return data["data"][0]["embedding"]

View File

@ -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"

View File

@ -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)