feat: 日志模块

This commit is contained in:
李如威 2025-12-11 10:28:50 +08:00
parent 9686f24b60
commit 1e18cd2e76
11 changed files with 112 additions and 100 deletions

View File

@ -9,3 +9,4 @@ python-docx
scikit-learn scikit-learn
aiofiles aiofiles
pillow pillow
loguru

View File

@ -1,3 +1,4 @@
#!/usr/bin/env bash #!/usr/bin/env bash
export $(cat .env | xargs) export $(cat .env | xargs)
pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py pytest -s -W ignore::DeprecationWarning src/tests/test_nodes.py

View File

@ -15,10 +15,12 @@ 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
logger_level: str
def _read_config() -> Config: def _read_config() -> Config:
return { return {
"logger_level": os.getenv("LOGGER_LEVEL", "DEBUG"),
"host": os.getenv("HOST"), "host": os.getenv("HOST"),
"port": int(os.getenv("PORT")), "port": int(os.getenv("PORT")),
"version": os.getenv("VERSION"), "version": os.getenv("VERSION"),

View File

@ -1,10 +1,9 @@
from src.pipeline.config import config from src.pipeline.config import config
import httpx import httpx
from urllib.parse import urljoin
async def get_embedding(text, timeout: int = 30): async def get_embedding(text, timeout: int = 30):
try: try:
async with httpx.AsyncClient(timeout=timeout) as client: async with httpx.AsyncClient(timeout=timeout, http2=False, trust_env=False) as client:
url = config["embedding_api_host"] url = config["embedding_api_host"]
body = { body = {
"model": config["embedding_model"], "model": config["embedding_model"],
@ -14,8 +13,6 @@ async def get_embedding(text, timeout: int = 30):
"Content-Type": "application/json", "Content-Type": "application/json",
"Authorization": f"Bearer {config['embedding_api_key']}", "Authorization": f"Bearer {config['embedding_api_key']}",
} }
print(url)
print(headers)
res = await client.post(url, headers=headers, json=body) res = await client.post(url, headers=headers, json=body)
res.raise_for_status() res.raise_for_status()
data = res.json() data = res.json()

View File

@ -0,0 +1,78 @@
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.llm import get_embedding
import re
class ChunkDocumentsNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["documents"]
async def exec_async(self, document):
"""
:param document: {text, file_name}
"""
# print(f"document: {document}")
text = document["text"]
# 先将所有制表符等替换为一个空格
text = re.sub(r"[ \t]+", " ", text)
# 再将多个空格替换为一个空格
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)]
async def post_async(self, shared, prep_res, exec_res_list):
all_chunks = []
for chunks in exec_res_list:
all_chunks.extend(chunks)
shared["documents"] = all_chunks
return "default"
class EmbeddingDocumentsNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["documents"]
async def exec_async(self, document):
"""
:param document: {text, file_name}
"""
return {**document, "embedding": await get_embedding(document["text"])}
async def post_async(self, shared, prep_res, exec_res_list):
shared["documents"] = exec_res_list
return "default"
class ReadDocumentNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["files"]
async def exec_async(self, file_path):
try:
document_text = await load_document(file_path)
return {
"file_path": file_path,
"file_name": file_path.split("/")[-1],
"text": document_text,
"text_length": len(document_text),
"status": "done",
"message": "",
}
except Exception as e:
return {
"file_path": file_path,
"file_name": file_path.split("/")[-1],
"text": "",
"text_length": 0,
"status": "error",
"message": str(e),
}
async def post_async(self, shared, prep_res, exec_res):
shared["documents"] = exec_res
return "default"

View File

@ -3,8 +3,30 @@ import docx
import fitz # PyMuPDF import fitz # PyMuPDF
import aiofiles import aiofiles
import io import io
import os
import sys
from pathlib import Path from pathlib import Path
from PIL import Image from PIL import Image
from loguru import logger
from src.pipeline.config import config
# -----------------------------
# 日志
# -----------------------------
logger.remove() # 清除已有 handler包括 pytest 导致的重复加载
logger.add(sys.stdout, level=config["logger_level"], colorize=True)
logger.add(
"logs/pipeline.log",
level=config["logger_level"],
rotation="10 MB", # 自动分割
retention="7 days", # 保留时间
compression="zip", # 自动压缩
enqueue=True, # 多线程安全
colorize=False, # 文件不需要颜色
backtrace=True,
diagnose=True,
)
# ----------------------------- # -----------------------------
# 文本加载 # 文本加载

View File

@ -1,4 +0,0 @@
from src.pipeline.nodes.chunk_document_node import ChunkDocumentsNode
from src.pipeline.nodes.read_document_node import ReadDocumentNode
__all__ = ["ChunkDocumentsNode", "ReadDocumentNode"]

View File

@ -1,32 +0,0 @@
import uuid
from src.pipeline.core.pocket_flow import AsyncBatchNode
from src.pipeline.core.utils import fixed_size_chunk
import re
class ChunkDocumentsNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["documents"]
async def exec_async(self, document):
"""
:param document: {text, file_name}
"""
# print(f"document: {document}")
text = document["text"]
# 先将所有制表符等替换为一个空格
text = re.sub(r"[ \t]+", " ", text)
# 再将多个空格替换为一个空格
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)]
async def post_async(self, shared, prep_res, exec_res_list):
all_chunks = []
for chunks in exec_res_list:
all_chunks.extend(chunks)
shared["documents"] = all_chunks
return "default"

View File

@ -1,19 +0,0 @@
from src.pipeline.core.pocket_flow import AsyncBatchNode
from src.pipeline.llm import get_embedding
class EmbeddingDocumentsNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["documents"]
async def exec_async(self, document):
"""
:param document: {text, file_name}
"""
return {**document, "embedding": await get_embedding(document["text"])}
async def post_async(self, shared, prep_res, exec_res_list):
shared["documents"] = exec_res_list
return "default"

View File

@ -1,32 +0,0 @@
from src.pipeline.core.pocket_flow import AsyncBatchNode
from src.pipeline.core.utils import load_document
class ReadDocumentNode(AsyncBatchNode):
async def prep_async(self, shared):
return shared["files"]
async def exec_async(self, file_path):
try:
document_text = await load_document(file_path)
return {
"file_path": file_path,
"file_name": file_path.split("/")[-1],
"text": document_text,
"text_length": len(document_text),
"status": "done",
"message": "",
}
except Exception as e:
return {
"file_path": file_path,
"file_name": file_path.split("/")[-1],
"text": "",
"text_length": 0,
"status": "error",
"message": str(e),
}
async def post_async(self, shared, prep_res, exec_res):
shared["documents"] = exec_res
return "default"

View File

@ -1,10 +1,8 @@
import pytest import pytest
import json import json
from src.pipeline.nodes.read_document_node import ReadDocumentNode from src.pipeline.core.nodes import ReadDocumentNode, ChunkDocumentsNode, EmbeddingDocumentsNode
from src.pipeline.nodes.chunk_document_node import ChunkDocumentsNode
from src.pipeline.nodes.embedding_document_node import EmbeddingDocumentsNode
# from src.pipeline.nodes import ReadDocumentNode, ChunkDocumentsNode
from src.pipeline.core.pocket_flow import AsyncFlow from src.pipeline.core.pocket_flow import AsyncFlow
from src.pipeline.core.utils import logger
@pytest.mark.asyncio @pytest.mark.asyncio
@ -15,9 +13,9 @@ async def test_embedding():
shared = { shared = {
"files": [ "files": [
"./files/带图片的.pdf", "./files/带图片的.pdf",
"./files/大白智问-API接入文档-V1.2.2.pdf", # "./files/大白智问-API接入文档-V1.2.2.pdf",
"./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md", # "./files/我来帮您创建一个美观简洁的微信小程序订单详情页面。首先让我了解一下当前的项目结构.md",
"./files/deepsearch状态.txt", # "./files/deepsearch状态.txt",
] ]
} }
@ -29,4 +27,4 @@ async def test_embedding():
await flow.run_async(shared) await flow.run_async(shared)
print(json.dumps(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))