feat: es 容器
This commit is contained in:
parent
63480e312e
commit
4345dfb9cb
|
|
@ -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
|
||||
|
|
@ -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
|
||||
loguru
|
||||
aiohttp
|
||||
elasticsearch>=8.0.0,<9.0.0
|
||||
|
|
@ -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")),
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue