refactor(message): return MessageRead schema instead of raw model
Browse filesConvert the return type of create_messages_for_turn to use MessageRead schema for better data encapsulation and consistency with API responses.
- 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 |
-
) ->
|
| 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 |
-
|
|
|
|
| 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]:
|