Skip to content

Commit f2d5b2a

Browse files
committed
fix:confort ruff
1 parent 2b27f86 commit f2d5b2a

File tree

5 files changed

+20
-9
lines changed

5 files changed

+20
-9
lines changed

enhancers/base_models.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@
22
from typing import Any
33

44
from openai import AsyncOpenAI
5+
from openai.types.chat import ParsedChatCompletionMessage
56
from pydantic import BaseModel
67
from tenacity import retry, stop_after_attempt, wait_exponential
8+
79
from parsers.base_models import DataItem
10+
811
MAX_RETRIES = 3
912
WAIT_TIME = 4
1013
WAIT_MAX_TIME = 15
@@ -27,7 +30,7 @@ async def enhance(self, information: DataItem) -> DataItem:
2730
pass
2831

2932
@retry(stop=stop_after_attempt(MAX_RETRIES), wait=wait_exponential(multiplier=MULTIPLIER, min=WAIT_TIME, max=WAIT_MAX_TIME))
30-
async def get_structured_response(self, user_prompt: list[dict[str, Any]], response_format: JsonResponseFormat) -> str|None:
33+
async def get_structured_response(self, user_prompt: list[dict[str, Any]], response_format: JsonResponseFormat) -> ParsedChatCompletionMessage:
3134
"""获取结构化响应"""
3235
response = await self.client.chat.completions.parse(
3336
model=self.model_name,
@@ -37,6 +40,4 @@ async def get_structured_response(self, user_prompt: list[dict[str, Any]], respo
3740
],
3841
response_format=response_format # type: ignore
3942
)
40-
if response.choices[0].message.refusal:
41-
raise ValueError(f"模型拒绝了请求: {response.choices[0].message.refusal}")
42-
return response.choices[0].message.parsed
43+
return response.choices[0].message

parsers/__init__.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
# Parsers package
22

3-
from .base_models import ChunkType, DocumentData, DocumentParser, TableDataItem, TextDataItem, ImageDataItem, FormulaDataItem
3+
from .base_models import (
4+
ChunkType,
5+
DocumentData,
6+
DocumentParser,
7+
FormulaDataItem,
8+
ImageDataItem,
9+
TableDataItem,
10+
TextDataItem,
11+
)
412
from .parser_registry import (
513
PARSER_REGISTRY,
614
get_parser,

parsers/docx_parser.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ def _extract_texts(self, texts:list[TitleItem|SectionHeaderItem|ListItem|CodeIte
232232

233233
async def _extract_images_async(self, pictures: list[PictureItem]) -> list[ImageDataItem]:
234234
"""异步提取文档中的图片"""
235-
loop = asyncio.get_event_loop()
235+
loop = asyncio.get_event_loop()
236236
return await loop.run_in_executor(None, self._extract_images, pictures)
237237

238238
async def _extract_tables_async(self, tables: list[TableItem]) -> list[TableDataItem]:

parsers/pdf_parser.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ async def _process_content_parallel(self, file_path: Path, content_list: list[di
105105
continue
106106
texts_task.append(self._process_text_async(idx, item))
107107

108-
108+
109109
texts_chunks = [item for item in (await asyncio.gather(*texts_task) if texts_task else []) if item is not None]
110110
tables_chunks = [item for item in (await asyncio.gather(*tables_task) if tables_task else []) if item is not None]
111111
images_chunks = [item for item in (await asyncio.gather(*images_task) if images_task else []) if item is not None]

worker.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33

44
from sanic import Sanic
55

6-
from enhancers import get_enhancer
7-
from parsers import TableDataItem,ImageDataItem,FormulaDataItem, TextDataItem, ChunkType, get_parser, load_all_parsers
6+
from parsers import (
7+
get_parser,
8+
load_all_parsers,
9+
)
810

911

1012
async def worker(app: Sanic) -> dict[str, Any]:

0 commit comments

Comments
 (0)