Add chat image upload support
This commit is contained in:
Xin Wang
@@ -2,6 +2,8 @@
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import mimetypes
|
||||
from pathlib import Path
|
||||
import weakref
|
||||
from typing import Any, Dict, Literal, Union
|
||||
|
||||
@@ -351,6 +353,245 @@ class AsyncChatClient(AsyncFastGPTClient):
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
async def _send_first_available_request(
|
||||
self,
|
||||
method: str,
|
||||
endpoints: list[str],
|
||||
json: Dict[str, Any] | None = None,
|
||||
params: Dict[str, Any] | None = None,
|
||||
):
|
||||
"""Try compatible FastGPT endpoint variants and return the first non-404 response."""
|
||||
last_error = None
|
||||
|
||||
for endpoint in endpoints:
|
||||
try:
|
||||
return await self._send_request(method, endpoint, json=json, params=params)
|
||||
except APIError as exc:
|
||||
last_error = exc
|
||||
if exc.status_code != 404:
|
||||
raise
|
||||
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
|
||||
raise APIError("No endpoint candidates provided")
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_response_data(data: Any) -> Any:
|
||||
"""Unwrap FastGPT's standard envelope while preserving raw OpenAI-style responses."""
|
||||
if isinstance(data, dict) and "data" in data and {"code", "message", "statusText"} & set(data):
|
||||
return data.get("data")
|
||||
return data
|
||||
|
||||
async def _get_chat_file_select_config(self, appId: str, chatId: str) -> dict[str, Any] | None:
|
||||
"""Best-effort lookup for the app's chat file selection config."""
|
||||
try:
|
||||
response = await self.get_chat_init(appId=appId, chatId=chatId)
|
||||
data = self._unwrap_response_data(response.json())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
|
||||
app = data.get("app")
|
||||
if not isinstance(app, dict):
|
||||
return None
|
||||
|
||||
chat_config = app.get("chatConfig")
|
||||
if not isinstance(chat_config, dict):
|
||||
return None
|
||||
|
||||
file_select_config = chat_config.get("fileSelectConfig")
|
||||
return file_select_config if isinstance(file_select_config, dict) else None
|
||||
|
||||
@staticmethod
|
||||
def _default_file_select_config(file_type: str) -> dict[str, Any]:
|
||||
"""Fallback config for servers that require fileSelectConfig in presign requests."""
|
||||
is_image = file_type == "img"
|
||||
return {
|
||||
"canSelectFile": not is_image,
|
||||
"canSelectImg": is_image,
|
||||
"canSelectVideo": False,
|
||||
"canSelectAudio": False,
|
||||
"canSelectCustomFileExtension": False,
|
||||
"customFileExtensionList": [],
|
||||
}
|
||||
|
||||
async def presign_chat_file_post_url(
|
||||
self,
|
||||
*,
|
||||
appId: str,
|
||||
chatId: str,
|
||||
filename: str,
|
||||
fileSelectConfig: dict[str, Any] | None = None,
|
||||
outLinkAuthData: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Request a presigned URL for uploading a chat file."""
|
||||
payload: dict[str, Any] = {
|
||||
"appId": appId,
|
||||
"chatId": chatId,
|
||||
"filename": filename,
|
||||
"fileSelectConfig": fileSelectConfig or self._default_file_select_config("img"),
|
||||
}
|
||||
|
||||
if outLinkAuthData:
|
||||
payload["outLinkAuthData"] = outLinkAuthData
|
||||
|
||||
response = await self._send_first_available_request(
|
||||
"POST",
|
||||
[
|
||||
"/api/core/chat/presignChatFilePostUrl",
|
||||
"/api/core/chat/file/presignChatFilePostUrl",
|
||||
],
|
||||
json=payload,
|
||||
)
|
||||
data = self._unwrap_response_data(response.json())
|
||||
if not isinstance(data, dict):
|
||||
raise APIError("Invalid presigned upload response")
|
||||
return data
|
||||
|
||||
async def presign_chat_file_get_url(
|
||||
self,
|
||||
*,
|
||||
appId: str,
|
||||
key: str,
|
||||
outLinkAuthData: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Request a temporary preview/download URL for an uploaded chat file."""
|
||||
payload: dict[str, Any] = {
|
||||
"appId": appId,
|
||||
"key": key,
|
||||
}
|
||||
|
||||
if outLinkAuthData:
|
||||
payload["outLinkAuthData"] = outLinkAuthData
|
||||
|
||||
response = await self._send_first_available_request(
|
||||
"POST",
|
||||
[
|
||||
"/api/core/chat/presignChatFileGetUrl",
|
||||
"/api/core/chat/file/presignChatFileGetUrl",
|
||||
],
|
||||
json=payload,
|
||||
)
|
||||
data = self._unwrap_response_data(response.json())
|
||||
if not isinstance(data, str):
|
||||
raise APIError("Invalid presigned preview response")
|
||||
return data
|
||||
|
||||
async def upload_to_presigned_url(
|
||||
self,
|
||||
*,
|
||||
upload_url: str,
|
||||
file_path: str | Path,
|
||||
headers: dict[str, str] | None = None,
|
||||
fields: dict[str, Any] | None = None,
|
||||
method: Literal["POST", "PUT"] | None = None,
|
||||
) -> httpx.Response:
|
||||
"""Upload a file to a presigned S3/MinIO URL."""
|
||||
path = Path(file_path)
|
||||
content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
|
||||
upload_method = method or ("POST" if fields else "PUT")
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
if upload_method == "POST":
|
||||
with path.open("rb") as file_obj:
|
||||
response = await client.post(
|
||||
upload_url,
|
||||
data=fields or {},
|
||||
files={"file": (path.name, file_obj, content_type)},
|
||||
)
|
||||
elif upload_method == "PUT":
|
||||
upload_headers = dict(headers or {})
|
||||
upload_headers.setdefault("Content-Type", content_type)
|
||||
with path.open("rb") as file_obj:
|
||||
response = await client.put(
|
||||
upload_url,
|
||||
content=file_obj,
|
||||
headers=upload_headers,
|
||||
)
|
||||
else:
|
||||
raise ValueError("method must be 'POST' or 'PUT'")
|
||||
|
||||
response.raise_for_status()
|
||||
return response
|
||||
|
||||
async def upload_chat_file(
|
||||
self,
|
||||
*,
|
||||
appId: str,
|
||||
chatId: str,
|
||||
file_path: str | Path,
|
||||
file_type: Literal["img", "doc"] = "doc",
|
||||
fileSelectConfig: dict[str, Any] | None = None,
|
||||
outLinkAuthData: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Upload a chat file and return a normalized file reference."""
|
||||
path = Path(file_path)
|
||||
selected_config = (
|
||||
fileSelectConfig
|
||||
or await self._get_chat_file_select_config(appId=appId, chatId=chatId)
|
||||
or self._default_file_select_config(file_type)
|
||||
)
|
||||
|
||||
presigned = await self.presign_chat_file_post_url(
|
||||
appId=appId,
|
||||
chatId=chatId,
|
||||
filename=path.name,
|
||||
fileSelectConfig=selected_config,
|
||||
outLinkAuthData=outLinkAuthData,
|
||||
)
|
||||
|
||||
max_size = presigned.get("maxSize")
|
||||
if max_size is not None and path.stat().st_size > max_size:
|
||||
raise ValueError(f"File too large: {path.stat().st_size} bytes > {max_size} bytes")
|
||||
|
||||
await self.upload_to_presigned_url(
|
||||
upload_url=presigned["url"],
|
||||
file_path=path,
|
||||
headers=presigned.get("headers"),
|
||||
fields=presigned.get("fields"),
|
||||
)
|
||||
|
||||
fields = presigned.get("fields") if isinstance(presigned.get("fields"), dict) else {}
|
||||
key = presigned.get("key") or fields.get("key")
|
||||
preview_url = presigned.get("previewUrl")
|
||||
|
||||
if not preview_url and key:
|
||||
preview_url = await self.presign_chat_file_get_url(
|
||||
appId=appId,
|
||||
key=key,
|
||||
outLinkAuthData=outLinkAuthData,
|
||||
)
|
||||
|
||||
return {
|
||||
"type": file_type,
|
||||
"name": path.name,
|
||||
"key": key,
|
||||
"url": preview_url,
|
||||
"previewUrl": preview_url,
|
||||
}
|
||||
|
||||
async def upload_chat_image(
|
||||
self,
|
||||
*,
|
||||
appId: str,
|
||||
chatId: str,
|
||||
file_path: str | Path,
|
||||
fileSelectConfig: dict[str, Any] | None = None,
|
||||
outLinkAuthData: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Upload an image for use in chat `image_url` content parts."""
|
||||
return await self.upload_chat_file(
|
||||
appId=appId,
|
||||
chatId=chatId,
|
||||
file_path=file_path,
|
||||
file_type="img",
|
||||
fileSelectConfig=fileSelectConfig,
|
||||
outLinkAuthData=outLinkAuthData,
|
||||
)
|
||||
|
||||
async def get_chat_histories(
|
||||
self,
|
||||
appId: str,
|
||||
|
||||
Reference in New Issue
Block a user