feat: 新增文档模型类，调整配置与解析逻辑，优化日志及导入
移除日志设置与冗余代码，优化导入、类型提示及OCR后端管理 统一调整各文件模块导入路径为绝对导入 调整导入路径，移除部分导入，优化日志及注释 升级文档解析器为 Docx2Parser，优化超时与图片处理逻辑
This commit is contained in:
begoniezhao
+59
-46
@@ -3,11 +3,10 @@ import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Union
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
import requests
|
||||
import ollama
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -158,11 +157,16 @@ class CaptionChatResp:
|
||||
Returns:
|
||||
The content string from the first choice, or empty string if no choices
|
||||
"""
|
||||
if self.choices:
|
||||
logger.info("Retrieving content from first choice")
|
||||
return self.choices[0].message.content
|
||||
logger.warning("No choices available in response")
|
||||
return ""
|
||||
if (
|
||||
not self.choices
|
||||
or not self.choices[0]
|
||||
or not self.choices[0].message
|
||||
or not self.choices[0].message.content
|
||||
):
|
||||
logger.warning("No choices available in response")
|
||||
return ""
|
||||
logger.info("Retrieving content from first choice")
|
||||
return self.choices[0].message.content
|
||||
|
||||
|
||||
class Caption:
|
||||
@@ -171,33 +175,43 @@ class Caption:
|
||||
Uses an external API to process images and return textual descriptions.
|
||||
"""
|
||||
|
||||
def __init__(self, vlm_config=None):
|
||||
"""Initialize the Caption service with configuration from parameters or environment variables."""
|
||||
def __init__(self, vlm_config: Optional[Dict[str, str]] = None):
|
||||
"""
|
||||
Initialize the Caption service with configuration
|
||||
from parameters or environment variables.
|
||||
"""
|
||||
logger.info("Initializing Caption service")
|
||||
self.prompt = """简单凝炼的描述图片的主要内容"""
|
||||
|
||||
# Use provided VLM config if available, otherwise fall back to environment variables
|
||||
self.timeout = 30
|
||||
|
||||
# Use provided VLM config if available,
|
||||
# otherwise fall back to environment variables
|
||||
if vlm_config and vlm_config.get("base_url") and vlm_config.get("model_name"):
|
||||
self.completion_url = vlm_config.get("base_url", "") + "/chat/completions"
|
||||
self.model = vlm_config.get("model_name", "")
|
||||
self.api_key = vlm_config.get("api_key", "")
|
||||
self.interface_type = vlm_config.get("interface_type", "openai").lower()
|
||||
else:
|
||||
if os.getenv("VLM_MODEL_BASE_URL") == "" or os.getenv("VLM_MODEL_NAME") == "":
|
||||
base_url = os.getenv("VLM_MODEL_BASE_URL")
|
||||
model_name = os.getenv("VLM_MODEL_NAME")
|
||||
if not base_url or not model_name:
|
||||
logger.error("VLM_MODEL_BASE_URL or VLM_MODEL_NAME is not set")
|
||||
return
|
||||
self.completion_url = os.getenv("VLM_MODEL_BASE_URL") + "/chat/completions"
|
||||
self.model = os.getenv("VLM_MODEL_NAME")
|
||||
self.api_key = os.getenv("VLM_MODEL_API_KEY")
|
||||
self.completion_url = base_url + "/chat/completions"
|
||||
self.model = model_name
|
||||
self.api_key = os.getenv("VLM_MODEL_API_KEY", "")
|
||||
self.interface_type = os.getenv("VLM_INTERFACE_TYPE", "openai").lower()
|
||||
|
||||
|
||||
# 验证接口类型
|
||||
if self.interface_type not in ["ollama", "openai"]:
|
||||
logger.warning(f"Unknown interface type: {self.interface_type}, defaulting to openai")
|
||||
logger.warning(
|
||||
f"Unknown interface type: {self.interface_type}, defaulting to openai"
|
||||
)
|
||||
self.interface_type = "openai"
|
||||
|
||||
|
||||
logger.info(
|
||||
f"Service configured with model: {self.model}, endpoint: {self.completion_url}, interface: {self.interface_type}"
|
||||
f"Configured with model: {self.model}, "
|
||||
f"endpoint: {self.completion_url}, interface: {self.interface_type}"
|
||||
)
|
||||
|
||||
def _call_caption_api(self, image_data: str) -> Optional[CaptionChatResp]:
|
||||
@@ -210,8 +224,8 @@ class Caption:
|
||||
Returns:
|
||||
CaptionChatResp object if successful, None otherwise
|
||||
"""
|
||||
logger.info(f"Calling Caption API for image captioning")
|
||||
logger.info(f"Processing image data: {image_data[:50] if len(image_data) > 50 else image_data}")
|
||||
logger.info("Calling Caption API for image captioning")
|
||||
logger.info(f"Processing image data: {image_data[:50]}...")
|
||||
|
||||
# 根据接口类型选择调用方式
|
||||
if self.interface_type == "ollama":
|
||||
@@ -226,39 +240,35 @@ class Caption:
|
||||
|
||||
client = ollama.Client(
|
||||
host=host,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
logger.info(f"Calling Ollama API with model: {self.model}")
|
||||
|
||||
|
||||
# 调用Ollama API，使用images参数传递base64编码的图片
|
||||
response = client.generate(
|
||||
model=self.model,
|
||||
prompt="简单凝炼的描述图片的主要内容",
|
||||
images=[image_base64], # image_base64是base64编码的图片数据
|
||||
images=[image_base64], # image_base64是base64编码的图片数据
|
||||
options={"temperature": 0.1},
|
||||
stream=False,
|
||||
)
|
||||
|
||||
|
||||
# 构造响应对象
|
||||
caption_resp = CaptionChatResp(
|
||||
id="ollama_response",
|
||||
created=int(time.time()),
|
||||
model=self.model,
|
||||
model=Model(id=self.model),
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
Choice(
|
||||
message=Message(
|
||||
role="assistant",
|
||||
content=response.response
|
||||
)
|
||||
)
|
||||
]
|
||||
Choice(message=Message(role="assistant", content=response.response))
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
logger.info("Successfully received response from Ollama API")
|
||||
return caption_resp
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error calling Ollama API: {e}")
|
||||
return None
|
||||
@@ -266,13 +276,16 @@ class Caption:
|
||||
def _call_openai_api(self, image_base64: str) -> Optional[CaptionChatResp]:
|
||||
"""Call OpenAI-compatible API for image captioning."""
|
||||
logger.info(f"Calling OpenAI-compatible API with model: {self.model}")
|
||||
|
||||
|
||||
user_msg = UserMessage(
|
||||
role="user",
|
||||
content=[
|
||||
Content(type="text", text=self.prompt),
|
||||
Content(
|
||||
type="image_url", image_url=ImageUrl(url="data:image/png;base64," + image_base64, detail="auto")
|
||||
type="image_url",
|
||||
image_url=ImageUrl(
|
||||
url="data:image/png;base64," + image_base64, detail="auto"
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
@@ -295,23 +308,23 @@ class Caption:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
try:
|
||||
logger.info(f"Sending request to OpenAI-compatible API with model: {self.model}")
|
||||
logger.info(
|
||||
f"Sending request to OpenAI-compatible API with model: {self.model}"
|
||||
)
|
||||
response = requests.post(
|
||||
self.completion_url,
|
||||
data=json.dumps(gpt_req, default=lambda o: o.__dict__, indent=4),
|
||||
headers=headers,
|
||||
timeout=30,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
logger.error(
|
||||
f"OpenAI-compatible API returned non-200 status code: {response.status_code}"
|
||||
f"OpenAI API returned non-200 status code: {response.status_code}"
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
logger.info(
|
||||
f"Successfully received response from OpenAI-compatible API with status: {response.status_code}"
|
||||
)
|
||||
logger.info(f"Converting response to CaptionChatResp object")
|
||||
logger.info(f"Received from OpenAI with status: {response.status_code}")
|
||||
logger.info("Converting response to CaptionChatResp object")
|
||||
caption_resp = CaptionChatResp.from_json(response.json())
|
||||
|
||||
if caption_resp.usage:
|
||||
@@ -322,7 +335,7 @@ class Caption:
|
||||
|
||||
return caption_resp
|
||||
except requests.exceptions.Timeout:
|
||||
logger.error(f"Timeout while calling OpenAI-compatible API after 30 seconds")
|
||||
logger.error("Timeout while calling OpenAI-compatible API after 30 seconds")
|
||||
return None
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Request error calling OpenAI-compatible API: {e}")
|
||||
|
||||
Reference in New Issue
Block a user