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 # 设置密码 - ELASTIC_PASSWORD=12345 # 设置密码
# - xpack.security.http.ssl.enabled=false # 关闭 HTTPS # - xpack.security.http.ssl.enabled=false # 关闭 HTTPS
# - ES_JAVA_OPTS=-Xms1g -Xmx1g # JVM 初始堆内存大小1GB最大堆内存大小1GB # - ES_JAVA_OPTS=-Xms1g -Xmx1g # JVM 初始堆内存大小1GB最大堆内存大小1GB
# elasticvue 配置
- http.cors.enabled=true
- http.cors.allow-origin=http://localhost:8210
ports: ports:
- "9210:9200" - "9210:9200"
volumes: volumes:
@ -24,17 +27,3 @@ services:
restart: unless-stopped restart: unless-stopped
networks: networks:
- ai-pipeline - 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 fastapi
uvicorn[standard] uvicorn[standard]
httpx[socks]
qdrant-client[httpx]
python-dotenv python-dotenv
pytest-asyncio pytest-asyncio
pymupdf pymupdf
@ -10,3 +8,5 @@ scikit-learn
aiofiles aiofiles
pillow pillow
loguru 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_key: str
embedding_api_host: str embedding_api_host: str
embedding_model: str embedding_model: str
embedding_dims: int
es_host: str es_host: str
es_port: int es_port: int
es_user: str es_user: str
@ -38,9 +39,10 @@ def _read_config() -> Config:
"embedding_api_host": os.getenv("EMBEDDING_API_HOST"), "embedding_api_host": os.getenv("EMBEDDING_API_HOST"),
"embedding_api_key": os.getenv("EMBEDDING_API_KEY"), "embedding_api_key": os.getenv("EMBEDDING_API_KEY"),
"embedding_model": os.getenv("EMBEDDING_MODEL"), "embedding_model": os.getenv("EMBEDDING_MODEL"),
"embedding_dims": int(os.getenv("EMBEDDING_DIMS")),
"es_host": os.getenv("ES_HOST"), "es_host": os.getenv("ES_HOST"),
"es_port": int(os.getenv("ES_PORT")), "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 "", "es_password": os.getenv("ES_PASSWORD") or "",
"qdrant_host": os.getenv("QDRANT_HOST"), "qdrant_host": os.getenv("QDRANT_HOST"),
"qdrant_port": int(os.getenv("QDRANT_PORT")), "qdrant_port": int(os.getenv("QDRANT_PORT")),

View File

@ -1,6 +1,8 @@
from codecs import ignore_errors
from datetime import datetime, timezone from datetime import datetime, timezone
from pipeline.config import config from venv import logger
from elasticsearch import AsyncElasticsearch from src.pipeline.config import config
from elasticsearch import AsyncElasticsearch, ApiError
class AsyncES: class AsyncES:
@ -27,7 +29,7 @@ class AsyncES:
"content": {"type": "text", "analyzer": "ik_smart"}, # 切片内容 "content": {"type": "text", "analyzer": "ik_smart"}, # 切片内容
"embedding": { # 向量字段 "embedding": { # 向量字段
"type": "dense_vector", "type": "dense_vector",
"dims": 768, "dims": 1024, # 必须和 embedding 模型的纬度一样
"index": True, "index": True,
"similarity": "cosine", "similarity": "cosine",
}, },
@ -39,20 +41,40 @@ class AsyncES:
""" """
创建 index知识库 创建 index知识库
""" """
is_exists = await self.client.indices.exists(index) try:
if is_exists: await self.client.indices.create(
return True index=index,
else: body=self.mapping,
await self.client.indices.create(index=index, body=self.mapping) )
return True 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: if created_at is None:
create_at = datetime.now(timezone.utc) 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()} 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) 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): 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"] for hit in res["hits"]["hits"]
] ]
es_client = AsyncES() es_client = AsyncES()

View File

@ -1,22 +1,29 @@
from src.pipeline.config import config from src.pipeline.config import config
import httpx import aiohttp
async def get_embedding(text, timeout: int = 30): _embedding_session: aiohttp.ClientSession | None = None
try:
async with httpx.AsyncClient(timeout=timeout, http2=False, trust_env=False) as client: async def get_embedding(text: str, timeout: int = 30):
url = config["embedding_api_host"] global _embedding_session
body = {
"model": config["embedding_model"], if _embedding_session is None or _embedding_session.closed:
"input": text, _embedding_session = aiohttp.ClientSession(
} timeout=aiohttp.ClientTimeout(total=timeout)
headers = { )
"Content-Type": "application/json",
"Authorization": f"Bearer {config['embedding_api_key']}", async with _embedding_session.post(
} config["embedding_api_host"],
res = await client.post(url, headers=headers, json=body) json={
res.raise_for_status() "model": config["embedding_model"],
data = res.json() "input": text,
return data["data"][0]["embedding"] },
except Exception as e: headers={
print(f"get_embedding[ERROR]: {e}") "Content-Type": "application/json",
return [] "Authorization": f"Bearer {config['embedding_api_key']}",
},
) as resp:
if resp.status >= 400:
return []
data = await resp.json()
return data["data"][0]["embedding"]

View File

@ -1,7 +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 from src.pipeline.core.utils import fixed_size_chunk, load_document, logger
from src.pipeline.core.llm import get_embedding from src.pipeline.core.llm import get_embedding
from src.pipeline.core.es_client import es_client
import re import re
@ -21,7 +22,15 @@ class ChunkDocumentsNode(AsyncBatchNode):
text = re.sub(r" +", " ", text) text = re.sub(r" +", " ", text)
# 去除首尾空格 # 去除首尾空格
text = text.strip() 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): async def post_async(self, shared, prep_res, exec_res_list):
all_chunks = [] all_chunks = []
@ -39,7 +48,10 @@ class EmbeddingDocumentsNode(AsyncBatchNode):
""" """
:param document: {text, file_name} :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): async def post_async(self, shared, prep_res, exec_res_list):
@ -60,6 +72,7 @@ class ReadDocumentNode(AsyncBatchNode):
"file_name": file_path.split("/")[-1], "file_name": file_path.split("/")[-1],
"text": document_text, "text": document_text,
"text_length": len(document_text), "text_length": len(document_text),
"file_type": str(file_path.split(".")[-1]).strip().lower(),
"status": "done", "status": "done",
"message": "", "message": "",
} }
@ -69,6 +82,7 @@ class ReadDocumentNode(AsyncBatchNode):
"file_name": file_path.split("/")[-1], "file_name": file_path.split("/")[-1],
"text": "", "text": "",
"text_length": 0, "text_length": 0,
"file_type": str(file_path.split(".")[-1]).strip().lower(),
"status": "error", "status": "error",
"message": str(e), "message": str(e),
} }
@ -76,3 +90,28 @@ class ReadDocumentNode(AsyncBatchNode):
async def post_async(self, shared, prep_res, exec_res): async def post_async(self, shared, prep_res, exec_res):
shared["documents"] = exec_res shared["documents"] = exec_res
return "default" 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 pytest
import json 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.pocket_flow import AsyncFlow
from src.pipeline.core.utils import logger from src.pipeline.core.utils import logger
@ -14,15 +14,18 @@ async def test_embedding():
"files": [ "files": [
# "./files/带图片的.pdf", # "./files/带图片的.pdf",
# "./files/大白智问-API接入文档-V1.2.2.pdf", # "./files/大白智问-API接入文档-V1.2.2.pdf",
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", "./files/山海经01.txt",
# "./files/deepsearch状态.txt", # "./files/deepsearch状态.txt",
] ],
"documents": [],
"index": "test_kb",
} }
readNode = ReadDocumentNode() readNode = ReadDocumentNode()
chunkNode = ChunkDocumentsNode() chunkNode = ChunkDocumentsNode()
embeddingNode = EmbeddingDocumentsNode() embeddingNode = EmbeddingDocumentsNode()
readNode >> chunkNode >> embeddingNode writeToESNode = WriteToElasticsearchNode()
readNode >> chunkNode >> embeddingNode >> writeToESNode
flow = AsyncFlow(readNode) flow = AsyncFlow(readNode)
await flow.run_async(shared) await flow.run_async(shared)