Skip to content

Commit

Permalink
Refactor message content handling in schemas.py to add capability for…
Browse files Browse the repository at this point in the history
… image content
  • Loading branch information
Wenlin88 committed Jun 20, 2024
1 parent df2124e commit 90fa1cc
Showing 1 changed file with 14 additions and 6 deletions.
20 changes: 14 additions & 6 deletions schemas.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,31 @@
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
messages: List[OpenAIChatMessage]

model_config = ConfigDict(extra="allow")


class FilterForm(BaseModel):
body: dict
user: Optional[dict] = None
model_config = ConfigDict(extra="allow")
model_config = ConfigDict(extra="allow")

0 comments on commit 90fa1cc

Please sign in to comment.