feat: es 容器
This commit is contained in:
parent
63480e312e
commit
4345dfb9cb
|
|
@ -17,24 +17,13 @@ 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:
|
||||||
- es-data:/usr/share/elasticsearch/data
|
- es-data:/usr/share/elasticsearch/data
|
||||||
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
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -9,4 +7,6 @@ python-docx
|
||||||
scikit-learn
|
scikit-learn
|
||||||
aiofiles
|
aiofiles
|
||||||
pillow
|
pillow
|
||||||
loguru
|
loguru
|
||||||
|
aiohttp
|
||||||
|
elasticsearch>=8.0.0,<9.0.0
|
||||||
|
|
@ -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")),
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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"]
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue