File size: 2,401 Bytes
bcd6d73 3ef184a 7c93711 a17e5a6 3ef184a a17e5a6 3ef184a a17e5a6 3ef184a 7c93711 a17e5a6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
from datetime import datetime
from app.mapper.base_mapper import BaseMapper
from app.model.chat_model import ChatCompletion, ChatMessageModel
from app.schema.chat_schema import ChatCompletionResponse, ChatCompletionRequest, ChatMessageResponse, ChoiceResponse
from loguru import logger
def to_message_schema(model: ChatMessageModel) -> ChatMessageResponse:
"""Convert ChatMessageModel to ChatMessageResponse schema."""
return ChatMessageResponse(
message_id=model.message_id,
role=model.role,
content=model.content,
figure=model.figure,
created_date=model.created_date,
)
class ChatMapper(BaseMapper[ChatCompletion, ChatCompletionResponse]):
"""Mapper for converting between ChatCompletion model and schema objects."""
def to_schema(self, model: ChatCompletion, convert_last_message: bool = False) -> ChatCompletionResponse:
"""Convert ChatCompletion model to ChatCompletionResponse schema."""
# Convert datetime to Unix timestamp
created_timestamp = int(model.created_date.timestamp()) if model.created_date else int(datetime.now().timestamp())
choices = []
if convert_last_message:
last_message = model.messages[-1]
choices.append(ChoiceResponse(index=0, message=to_message_schema(last_message), finish_reason="stop"))
else:
index = 0
for message in model.messages:
choices.append(ChoiceResponse(index=index, message=to_message_schema(message), finish_reason="stop"))
index += 1
return ChatCompletionResponse(completion_id=model.completion_id, model=model.model, created=created_timestamp, choices=choices)
def to_model(self, schema: ChatCompletionRequest) -> ChatCompletion:
"""Convert ChatCompletionRequest schema to ChatCompletion model."""
# Convert messages to ChatMessage objects
messages = []
if schema.messages:
for msg in schema.messages:
messages.append(ChatMessageModel(role=msg.role, content=msg.content, figure=None, message_id=None))
try:
return ChatCompletion(completion_id=schema.completion_id, model=schema.model, messages=messages, stream=schema.stream)
except Exception as e:
logger.error(f"Error converting ChatCompletionRequest to ChatCompletion: {e}")
raise
|