|
1 | 1 | import copy |
| 2 | +import mimetypes |
2 | 3 | import os |
3 | 4 | import uuid |
4 | 5 | from collections.abc import AsyncIterator, Sequence |
5 | 6 | from typing import TYPE_CHECKING, Any, Literal, Self, Unpack, cast |
6 | 7 |
|
7 | 8 | import httpx |
8 | | -from openai import AsyncOpenAI, AsyncStream, OpenAIError, omit |
| 9 | +from openai import AsyncOpenAI, AsyncStream, BaseModel, OpenAIError, omit |
| 10 | +from openai._types import RequestFiles, RequestOptions |
9 | 11 | from openai.types.chat import ( |
10 | 12 | ChatCompletion, |
11 | 13 | ChatCompletionChunk, |
|
24 | 26 | TokenUsage, |
25 | 27 | ) |
26 | 28 | from kosong.chat_provider.openai_common import convert_error, tool_to_openai |
27 | | -from kosong.message import ContentPart, Message, TextPart, ThinkPart, ToolCall, ToolCallPart |
| 29 | +from kosong.message import ( |
| 30 | + ContentPart, |
| 31 | + Message, |
| 32 | + TextPart, |
| 33 | + ThinkPart, |
| 34 | + ToolCall, |
| 35 | + ToolCallPart, |
| 36 | + VideoURLPart, |
| 37 | +) |
28 | 38 | from kosong.tooling import Tool |
29 | 39 |
|
30 | 40 | if TYPE_CHECKING: |
@@ -216,6 +226,50 @@ def model_parameters(self) -> dict[str, Any]: |
216 | 226 | model_parameters.update(self._generation_kwargs) |
217 | 227 | return model_parameters |
218 | 228 |
|
| 229 | + @property |
| 230 | + def files(self) -> "KimiFiles": |
| 231 | + return KimiFiles(self.client) |
| 232 | + |
| 233 | + |
| 234 | +class KimiFiles: |
| 235 | + def __init__(self, client: AsyncOpenAI) -> None: |
| 236 | + self._client = client |
| 237 | + |
| 238 | + async def upload_video(self, *, data: bytes, mime_type: str) -> VideoURLPart: |
| 239 | + """Upload a video to Kimi files API and return a video URL content part.""" |
| 240 | + if not mime_type.startswith("video/"): |
| 241 | + raise ChatProviderError(f"Expected a video mime type, got {mime_type}") |
| 242 | + url = await self._upload_file(data=data, mime_type=mime_type, purpose="video") |
| 243 | + return VideoURLPart(video_url=VideoURLPart.VideoURL(url=url)) |
| 244 | + |
| 245 | + async def _upload_file(self, *, data: bytes, mime_type: str, purpose: "KimiFilePurpose") -> str: |
| 246 | + filename = _guess_filename(mime_type) |
| 247 | + files: RequestFiles = {"file": (filename, data, mime_type)} |
| 248 | + options: RequestOptions = {"headers": {"Content-Type": "multipart/form-data"}} |
| 249 | + try: |
| 250 | + response: KimiFileObject = await self._client.post( |
| 251 | + "/files", |
| 252 | + cast_to=KimiFileObject, |
| 253 | + body={"purpose": purpose}, |
| 254 | + files=files, |
| 255 | + options=options, |
| 256 | + ) |
| 257 | + except (OpenAIError, httpx.HTTPError) as e: |
| 258 | + raise convert_error(e) from e |
| 259 | + return f"ms://{response.id}" |
| 260 | + |
| 261 | + |
| 262 | +class KimiFileObject(BaseModel): |
| 263 | + id: str |
| 264 | + |
| 265 | + |
| 266 | +type KimiFilePurpose = Literal["video", "image"] |
| 267 | + |
| 268 | + |
| 269 | +def _guess_filename(mime_type: str) -> str: |
| 270 | + extension = mimetypes.guess_extension(mime_type) or ".bin" |
| 271 | + return f"upload{extension}" |
| 272 | + |
219 | 273 |
|
220 | 274 | def _convert_message(message: Message) -> ChatCompletionMessageParam: |
221 | 275 | message = message.model_copy(deep=True) |
|
0 commit comments