diff --git a/.env.template b/.env.template index c1740c9..f4542c7 100644 --- a/.env.template +++ b/.env.template @@ -51,4 +51,13 @@ DEBUG=false LOG_LEVEL=INFO # PDF解析 -MINERU_MODEL_SOURCE=local \ No newline at end of file +MINERU_MODEL_SOURCE=local + +# 信息增强 +LLM_MODEL_NAME=gpt-4o +LLM_BASE_URL=http://192.168.120.2:4000 +LLM_API_KEY=ae + +VLLM_MODEL_NAME=qwen2.5-vl-7b-instruct +VLLM_API_KEY=sk- +VLLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 \ No newline at end of file diff --git a/config.py b/config.py index 487fe5a..66519d9 100644 --- a/config.py +++ b/config.py @@ -33,4 +33,13 @@ class Settings: MAX_FILES_PER_REQUEST: int = int(os.getenv("MAX_FILES_PER_REQUEST", "20")) TASK_TIMEOUT: int = int(os.getenv("TASK_TIMEOUT", "3600")) # 1小时 + # 模型配置 + LLM_MODEL_NAME: str = os.getenv("LLM_MODEL_NAME", "gpt-4o") + LLM_BASE_URL: str = os.getenv("LLM_BASE_URL", "http://192.168.120.2:4000") + LLM_API_KEY: str = os.getenv("LLM_API_KEY", "sk-") + + VLLM_MODEL_NAME: str = os.getenv("VLLM_MODEL_NAME", "qwen2.5-vl-7b-instruct") + VLLM_API_KEY: str = os.getenv("VLLM_API_KEY", "sk-") + VLLM_BASE_URL: str = os.getenv("VLLM_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1") + settings = Settings() diff --git a/enhancers/base_models.py b/enhancers/base_models.py index ba77fa7..189c3a4 100644 --- a/enhancers/base_models.py +++ b/enhancers/base_models.py @@ -1,32 +1,44 @@ from abc import ABC, abstractmethod +from typing import Any + +from openai import AsyncOpenAI +from pydantic import BaseModel +from tenacity import retry, stop_after_attempt, wait_exponential from parsers.base_models import ChunkData +MAX_RETRIES = 3 +WAIT_TIME = 4 +WAIT_MAX_TIME = 15 +MULTIPLIER = 1 + +class JsonResponseFormat(BaseModel): + """JSON 响应格式""" + description:str class InformationEnhancer(ABC): """信息增强器基类""" + def __init__(self, model_name: str, base_url: str, api_key: str): + self.client = AsyncOpenAI(api_key=api_key, base_url=base_url) + self.model_name = model_name + self.system_prompt = "You are a helpful assistant." + @abstractmethod async def enhance(self, information: ChunkData) -> ChunkData: """增强信息""" pass -class TableInformationEnhancer(InformationEnhancer): - """表格信息增强器""" - - async def enhance(self, information: ChunkData) -> ChunkData: - """增强信息""" - return information - -class FormulasInformationEnhancer(InformationEnhancer): - """公式信息增强器""" - - async def enhance(self, information: ChunkData) -> ChunkData: - """增强信息""" - return information - -class ImageInformationEnhancer(InformationEnhancer): - """图片信息增强器""" - - async def enhance(self, information: ChunkData) -> ChunkData: - """增强信息""" - return information + @retry(stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=MULTIPLIER, min=WAIT_TIME, max=WAIT_MAX_TIME)) + async def get_structured_response(self, user_prompt: list[dict[str, Any]], response_format: JsonResponseFormat) -> str|None: + """获取结构化响应""" + response = await self.client.chat.completions.parse( + model=self.model_name, + messages=[ + {"role": "system", "content": self.system_prompt}, + {"role": "user", "content": user_prompt} # type: ignore + ], + response_format=response_format # type: ignore + ) + if response.choices[0].message.refusal: + raise ValueError(f"模型拒绝了请求: {response.choices[0].message.refusal}") + return response.choices[0].message.parsed diff --git a/enhancers/enhancer_registry.py b/enhancers/enhancer_registry.py index 58a8d0a..01d44fe 100644 --- a/enhancers/enhancer_registry.py +++ b/enhancers/enhancer_registry.py @@ -7,6 +7,7 @@ import logging from collections.abc import Callable +from config import settings from enhancers.base_models import InformationEnhancer from parsers.base_models import ChunkType @@ -67,7 +68,11 @@ def get_enhancer(modality: ChunkType) -> InformationEnhancer | None: enhancer_class = ENHANCER_REGISTRY[modality_type] try: - return enhancer_class() + match modality_type: + case ChunkType.IMAGE.value.lower(): + return enhancer_class(settings.VLLM_MODEL_NAME, settings.VLLM_BASE_URL, settings.VLLM_API_KEY) + case _: + return enhancer_class(settings.LLM_MODEL_NAME, settings.LLM_BASE_URL, settings.LLM_API_KEY) except Exception as e: logger.error(f"创建信息增强器实例失败: {enhancer_class.__name__}, 错误: {e}") return None diff --git a/pyproject.toml b/pyproject.toml index f69e8f5..7d2c154 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "docling>=2.45.0", "mineru[core]>=2.1.11", "beautifulsoup4>=4.13.4", + "tenacity>=9.1.2", ] [dependency-groups] diff --git a/uv.lock b/uv.lock index 2361566..bfcfd7c 100644 --- a/uv.lock +++ b/uv.lock @@ -1806,6 +1806,7 @@ dependencies = [ { name = "redis" }, { name = "sanic" }, { name = "sanic-ext" }, + { name = "tenacity" }, ] [package.dev-dependencies] @@ -1852,6 +1853,7 @@ requires-dist = [ { name = "redis", specifier = ">=6.4.0" }, { name = "sanic", specifier = ">=23.12.0" }, { name = "sanic-ext", specifier = ">=23.12.0" }, + { name = "tenacity", specifier = ">=9.1.2" }, ] [package.metadata.requires-dev] @@ -3718,6 +3720,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" }, ] +[[package]] +name = "tenacity" +version = "9.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0a/d4/2b0cd0fe285e14b36db076e78c93766ff1d529d70408bd1d2a5a84f1d929/tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb", size = 48036, upload-time = "2025-04-02T08:25:09.966Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/30/643397144bfbfec6f6ef821f36f33e57d35946c44a2352d3c9f0ae847619/tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138", size = 28248, upload-time = "2025-04-02T08:25:07.678Z" }, +] + [[package]] name = "thop" version = "0.1.1.post2209072238"