Aasher commited on
Commit
3361c41
·
1 Parent(s): b68a6a7

refactor(message): return MessageRead schema instead of raw model

Browse files

Convert the return type of create_messages_for_turn to use MessageRead schema for better data encapsulation and consistency with API responses.

Files changed (1) hide show
  1. db/crud/message.py +6 -5
db/crud/message.py CHANGED
@@ -5,6 +5,7 @@ from datetime import datetime, timezone
5
 
6
  from db.models.chat import Chat
7
  from db.models.message import Message
 
8
  from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
9
 
10
 
@@ -14,7 +15,7 @@ async def create_messages_for_turn(
14
  new_lc_messages: list[BaseMessage],
15
  final_answer: str,
16
  links: list[str],
17
- ) -> Message:
18
  """
19
  Creates all new messages generated during a single agent turn.
20
 
@@ -31,7 +32,7 @@ async def create_messages_for_turn(
31
  links: The list of parsed source links from the agent's output.
32
 
33
  Returns:
34
- The final AI Message
35
  """
36
  # First, get the parent chat to update its timestamp.
37
  # This also ensures the chat exists before we try to add messages to it.
@@ -76,11 +77,11 @@ async def create_messages_for_turn(
76
  db.add_all(db_messages_to_add)
77
  await db.commit()
78
 
79
- # Only refresh the single message we need to return to the client.
80
  final_ai_message = db_messages_to_add[-1]
81
  await db.refresh(final_ai_message)
82
-
83
- return final_ai_message
 
84
 
85
 
86
  async def get_messages_by_chat(db: AsyncSession, chat_id: uuid.UUID) -> list[Message]:
 
5
 
6
  from db.models.chat import Chat
7
  from db.models.message import Message
8
+ from db.schemas.message import MessageRead
9
  from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
10
 
11
 
 
15
  new_lc_messages: list[BaseMessage],
16
  final_answer: str,
17
  links: list[str],
18
+ ) -> MessageRead:
19
  """
20
  Creates all new messages generated during a single agent turn.
21
 
 
32
  links: The list of parsed source links from the agent's output.
33
 
34
  Returns:
35
+ The final AI Message, converted to a MessageRead schema.
36
  """
37
  # First, get the parent chat to update its timestamp.
38
  # This also ensures the chat exists before we try to add messages to it.
 
77
  db.add_all(db_messages_to_add)
78
  await db.commit()
79
 
 
80
  final_ai_message = db_messages_to_add[-1]
81
  await db.refresh(final_ai_message)
82
+
83
+ # Return the final AI message as a MessageRead schema
84
+ return MessageRead.model_validate(final_ai_message)
85
 
86
 
87
  async def get_messages_by_chat(db: AsyncSession, chat_id: uuid.UUID) -> list[Message]: