From 90fa1cc3b30253bed1ace217d64fbd7df542c987 Mon Sep 17 00:00:00 2001 From: Henri Wenlin Date: Thu, 20 Jun 2024 12:15:24 +0300 Subject: [PATCH] Refactor message content handling in schemas.py to add capability for image content --- schemas.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/schemas.py b/schemas.py index da4122a7..57517dfd 100644 --- a/schemas.py +++ b/schemas.py @@ -1,14 +1,23 @@ -from typing import List, Optional -from pydantic import BaseModel, ConfigDict +from typing import List, Union, Optional +from pydantic import BaseModel, RootModel, ConfigDict +class ImageContent(BaseModel): + type: str + image_url: dict + +class TextContent(BaseModel): + type: str + text: str + +class MessageContent(RootModel): + root: Union[TextContent, ImageContent] class OpenAIChatMessage(BaseModel): role: str - content: str | List + content: Union[str, List[MessageContent]] model_config = ConfigDict(extra="allow") - class OpenAIChatCompletionForm(BaseModel): stream: bool = True model: str @@ -16,8 +25,7 @@ class OpenAIChatCompletionForm(BaseModel): model_config = ConfigDict(extra="allow") - class FilterForm(BaseModel): body: dict user: Optional[dict] = None - model_config = ConfigDict(extra="allow") + model_config = ConfigDict(extra="allow") \ No newline at end of file