TheSmallHanCat commited on
Commit
77eaf1d
·
1 Parent(s): a1634c8

feat: 图片模型连续对话自动上一次响应作为参考图

Browse files
Files changed (1) hide show
  1. src/api/routes.py +66 -1
src/api/routes.py CHANGED
@@ -1,13 +1,17 @@
1
  """API routes - OpenAI compatible endpoints"""
2
  from fastapi import APIRouter, Depends, HTTPException
3
  from fastapi.responses import StreamingResponse, JSONResponse
4
- from typing import List
5
  import base64
6
  import re
7
  import json
 
 
 
8
  from ..core.auth import verify_api_key_header
9
  from ..core.models import ChatCompletionRequest
10
  from ..services.generation_handler import GenerationHandler, MODEL_CONFIG
 
11
 
12
  router = APIRouter()
13
 
@@ -21,6 +25,40 @@ def set_generation_handler(handler: GenerationHandler):
21
  generation_handler = handler
22
 
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  @router.get("/v1/models")
25
  async def list_models(api_key: str = Depends(verify_api_key_header)):
26
  """List available models"""
@@ -92,6 +130,33 @@ async def create_chat_completion(
92
  image_bytes = base64.b64decode(image_base64)
93
  images.append(image_bytes)
94
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
95
  if not prompt:
96
  raise HTTPException(status_code=400, detail="Prompt cannot be empty")
97
 
 
1
  """API routes - OpenAI compatible endpoints"""
2
  from fastapi import APIRouter, Depends, HTTPException
3
  from fastapi.responses import StreamingResponse, JSONResponse
4
+ from typing import List, Optional
5
  import base64
6
  import re
7
  import json
8
+ import time
9
+ from urllib.parse import urlparse
10
+ from curl_cffi.requests import AsyncSession
11
  from ..core.auth import verify_api_key_header
12
  from ..core.models import ChatCompletionRequest
13
  from ..services.generation_handler import GenerationHandler, MODEL_CONFIG
14
+ from ..core.logger import debug_logger
15
 
16
  router = APIRouter()
17
 
 
25
  generation_handler = handler
26
 
27
 
28
+ async def retrieve_image_data(url: str) -> Optional[bytes]:
29
+ """
30
+ 智能获取图片数据:
31
+ 1. 优先检查是否为本地 /tmp/ 缓存文件,如果是则直接读取磁盘
32
+ 2. 如果本地不存在或是外部链接,则进行网络下载
33
+ """
34
+ # 优先尝试本地读取
35
+ try:
36
+ if "/tmp/" in url and generation_handler and generation_handler.file_cache:
37
+ path = urlparse(url).path
38
+ filename = path.split("/tmp/")[-1]
39
+ local_file_path = generation_handler.file_cache.cache_dir / filename
40
+
41
+ if local_file_path.exists() and local_file_path.is_file():
42
+ data = local_file_path.read_bytes()
43
+ if data:
44
+ return data
45
+ except Exception as e:
46
+ debug_logger.log_warning(f"[CONTEXT] 本地缓存读取失败: {str(e)}")
47
+
48
+ # 回退逻辑:网络下载
49
+ try:
50
+ async with AsyncSession() as session:
51
+ response = await session.get(url, timeout=30, impersonate="chrome110", verify=False)
52
+ if response.status_code == 200:
53
+ return response.content
54
+ else:
55
+ debug_logger.log_warning(f"[CONTEXT] 图片下载失败,状态码: {response.status_code}")
56
+ except Exception as e:
57
+ debug_logger.log_error(f"[CONTEXT] 图片下载异常: {str(e)}")
58
+
59
+ return None
60
+
61
+
62
  @router.get("/v1/models")
63
  async def list_models(api_key: str = Depends(verify_api_key_header)):
64
  """List available models"""
 
130
  image_bytes = base64.b64decode(image_base64)
131
  images.append(image_bytes)
132
 
133
+ # 自动参考图:仅对图片模型生效
134
+ model_config = MODEL_CONFIG.get(request.model)
135
+
136
+ if model_config and model_config["type"] == "image" and not images and len(request.messages) > 1:
137
+ debug_logger.log_info(f"[CONTEXT] 开始查找历史参考图,消息数量: {len(request.messages)}")
138
+
139
+ # 如果当前请求没有上传图片,则尝试从历史记录中寻找最近的一张生成图
140
+ for msg in reversed(request.messages[:-1]):
141
+ if msg.role == "assistant" and isinstance(msg.content, str):
142
+ # 匹配 Markdown 图片格式: ![...](http...)
143
+ matches = re.findall(r"!\[.*?\]\((.*?)\)", msg.content)
144
+ if matches:
145
+ last_image_url = matches[-1]
146
+
147
+ if last_image_url.startswith("http"):
148
+ try:
149
+ downloaded_bytes = await retrieve_image_data(last_image_url)
150
+ if downloaded_bytes and len(downloaded_bytes) > 0:
151
+ images.append(downloaded_bytes)
152
+ debug_logger.log_info(f"[CONTEXT] ✅ 自动使用历史参考图: {last_image_url}")
153
+ break
154
+ else:
155
+ debug_logger.log_warning(f"[CONTEXT] 图片下载失败或为空,尝试下一个: {last_image_url}")
156
+ except Exception as e:
157
+ debug_logger.log_error(f"[CONTEXT] 处理参考图时出错: {str(e)}")
158
+ # 继续尝试下一个图片
159
+
160
  if not prompt:
161
  raise HTTPException(status_code=400, detail="Prompt cannot be empty")
162