♻️(chat) rewrite backend using Pydantic AI SDK

We chose to use Pydantic AI instead of the OpenAI SDK.
This commit should not change the chat behavior.
This commit is contained in:
Quentin BEY
2025-07-25 17:46:42 +02:00
parent 0526c9bbef
commit c5af3cdcb1
15 changed files with 2121 additions and 118 deletions
+4
View File
@@ -8,6 +8,10 @@ and this project adheres to
## [Unreleased]
### Changed
- ♻️(chat) rewrite backend using Pydantic AI SDK #4
### Added
- 🎉(conversations) bootstrap backend & frontend #1
+44
View File
@@ -0,0 +1,44 @@
"""
Helpers to manage async objects in a synchronous context.
This is not optimal, but we would prefer to stay in a synchronous context
for now.
"""
import asyncio
import queue
import threading
def convert_async_generator_to_sync(async_gen):
"""Convert an async generator to a sync generator."""
q = queue.Queue()
sentinel = object()
exc_sentinel = object()
async def run_async_gen():
try:
async for async_item in async_gen:
q.put(async_item)
except Exception as exc: # pylint: disable=broad-except #noqa: BLE001
q.put((exc_sentinel, exc))
finally:
q.put(sentinel)
def start_async_loop():
asyncio.run(run_async_gen())
thread = threading.Thread(target=start_async_loop, daemon=True)
thread.start()
try:
while True:
item = q.get()
if item is sentinel:
break
if isinstance(item, tuple) and item[0] is exc_sentinel:
# re-raise the exception in the sync context
raise item[1]
yield item
finally:
thread.join()
+1 -37
View File
@@ -1,10 +1,7 @@
"""AIAgentService class for handling AI agent interactions."""
import asyncio
import json
import logging
import queue
import threading
import uuid
from contextlib import AsyncExitStack
from typing import List
@@ -29,46 +26,13 @@ from chat.ai_sdk_types import (
ToolInvocationUIPart,
UIMessage,
)
from chat.clients.async_to_sync import convert_async_generator_to_sync
from chat.mcp_servers import get_mcp_servers
from chat.tools import get_tool_by_name
logger = logging.getLogger(__name__)
def convert_async_generator_to_sync(async_gen):
"""Convert an async generator to a sync generator."""
q = queue.Queue()
sentinel = object()
exc_sentinel = object()
async def run_async_gen():
try:
async for item in async_gen:
q.put(item)
except Exception as exc: # pylint: disable=broad-except #noqa: BLE001
q.put((exc_sentinel, exc))
finally:
q.put(sentinel)
def start_async_loop():
asyncio.run(run_async_gen())
thread = threading.Thread(target=start_async_loop, daemon=True)
thread.start()
try:
while True:
item = q.get()
if item is sentinel:
break
if isinstance(item, tuple) and item[0] is exc_sentinel:
# re-raise the exception in the sync context
raise item[1]
yield item
finally:
thread.join()
class AIAgentService:
"""Service class for AI-related operations."""
+295
View File
@@ -0,0 +1,295 @@
"""
Pydantic-AI based AIAgentService.
This file replaces the previous OpenAI-specific client with a Pydantic-AI
implementation while keeping the *exact* same public API so that no
changes are needed in views.py or tests.
"""
import dataclasses
import json
import logging
from contextlib import AsyncExitStack
from itertools import chain
from typing import Dict, List
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from asgiref.sync import sync_to_async
from pydantic_ai import Agent
from pydantic_ai.messages import (
FunctionToolCallEvent,
FunctionToolResultEvent,
ModelMessage,
ModelRequest,
ModelResponse,
PartDeltaEvent,
PartStartEvent,
RetryPromptPart,
TextPart,
TextPartDelta,
ThinkingPart,
ThinkingPartDelta,
ToolCallPart,
ToolCallPartDelta,
ToolReturnPart,
)
from pydantic_ai.models.openai import OpenAIModel, OpenAIResponsesModelSettings
from pydantic_ai.providers.openai import OpenAIProvider
from chat.ai_sdk_types import (
UIMessage,
)
from chat.clients.async_to_sync import convert_async_generator_to_sync
from chat.clients.pydantic_ui_message_converter import (
model_message_to_ui_message,
ui_message_to_model_message,
ui_message_to_user_content,
)
from chat.mcp_servers import get_mcp_servers
from chat.tools import get_pydantic_tools_by_name
logger = logging.getLogger(__name__)
def _build_pydantic_agent(mcp_servers) -> Agent[None, str]:
"""Create a Pydantic AI Agent instance with the configured settings."""
if settings.AI_BASE_URL is None or settings.AI_API_KEY is None or settings.AI_MODEL is None:
raise ImproperlyConfigured("AIChatService configuration not set")
agent = Agent(
model=OpenAIModel(
model_name=settings.AI_MODEL,
provider=OpenAIProvider(
base_url=settings.AI_BASE_URL,
api_key=settings.AI_API_KEY,
),
settings=OpenAIResponsesModelSettings(
openai_reasoning_effort="low",
openai_reasoning_summary="detailed",
),
),
system_prompt=settings.AI_AGENT_INSTRUCTIONS,
mcp_servers=mcp_servers,
tools=[get_pydantic_tools_by_name(tool_name) for tool_name in settings.AI_AGENT_TOOLS],
)
return agent
class AIAgentService:
"""Service class for AI-related operations (Pydantic-AI edition)."""
def __init__(self, conversation):
self.conversation = conversation
# --------------------------------------------------------------------- #
# Public streaming API (unchanged signatures)
# --------------------------------------------------------------------- #
def stream_text(self, messages: List[UIMessage]):
"""Return only the assistant text deltas (legacy text mode)."""
return convert_async_generator_to_sync(self.stream_text_async(messages))
def stream_data(self, messages: List[UIMessage]):
"""Return Vercel-AI-SDK formatted events."""
return convert_async_generator_to_sync(self.stream_data_async(messages))
# --------------------------------------------------------------------- #
# Async internals
# --------------------------------------------------------------------- #
async def stream_text_async(self, messages: List[UIMessage]):
"""Return only the assistant text deltas (legacy text mode)."""
async for delta in self._run_agent(messages):
if delta["type"] == "0":
yield delta["payload"]
async def stream_data_async(self, messages: List[UIMessage]):
"""Return Vercel-AI-SDK formatted events."""
async for delta in self._run_agent(messages):
yield f"{delta['type']}:{json.dumps(delta['payload'])}\n"
# --------------------------------------------------------------------- #
# Core agent runner
# --------------------------------------------------------------------- #
# pylint: disable=too-many-branches,too-many-statements
async def _run_agent(self, messages: List[UIMessage]): # noqa: PLR0912
"""Run the Pydantic AI agent and stream events."""
if messages[-1].role != "user":
return
history = [ui_message_to_model_message(message) for message in messages[:-1]]
prompt = ui_message_to_user_content(messages[-1])
usage = {"promptTokens": 0, "completionTokens": 0}
async with AsyncExitStack() as stack:
# MCP servers (if any) can be initialized here
mcp_servers = [await stack.enter_async_context(mcp) for mcp in get_mcp_servers()]
async with _build_pydantic_agent(mcp_servers).iter(
prompt, message_history=history
) as run:
async for node in run:
if Agent.is_user_prompt_node(node):
# A user prompt node => The user has provided input
pass
elif Agent.is_model_request_node(node):
# A model request node => We can stream tokens from the model's request
async with node.stream(run.ctx) as request_stream:
async for event in request_stream:
logger.debug("Received request_stream event: %s", type(event))
if isinstance(event, PartStartEvent):
logger.debug("PartStartEvent: %s", dataclasses.asdict(event))
if isinstance(event.part, TextPart):
yield {"type": "0", "payload": event.part.content}
elif isinstance(event.part, ToolCallPart):
yield {
"type": "b",
"payload": {
"toolCallId": event.part.tool_call_id,
"toolName": event.part.tool_name,
},
}
elif isinstance(event.part, ThinkingPart):
yield {"type": "g", "payload": event.part.content}
elif isinstance(event, PartDeltaEvent):
logger.debug(
"PartDeltaEvent: %s %s",
type(event),
dataclasses.asdict(event),
)
if isinstance(event.delta, TextPartDelta):
yield {"type": "0", "payload": event.delta.content_delta}
elif isinstance(event.delta, ToolCallPartDelta):
yield {
"type": "c",
"payload": {
"toolCallId": event.delta.tool_call_id,
"argsTextDelta": event.delta.args_delta,
},
}
elif isinstance(event.delta, ThinkingPartDelta):
yield {"type": "g", "payload": event.delta.content_delta}
elif Agent.is_call_tools_node(node):
# A handle-response node => The model returned some data,
# potentially calls a tool
async with node.stream(run.ctx) as handle_stream:
async for event in handle_stream:
logger.debug(
"Received request_stream event: %s, %s",
type(event),
dataclasses.asdict(event),
)
if isinstance(event, FunctionToolCallEvent):
# We are already streaming the tool call events don't yield
# the tool call again
pass
# yield {
# "type": "9",
# "payload": {
# "toolCallId": event.part.tool_call_id,
# "toolName": event.part.tool_name,
# "args": event.part.args,
# },
# }
elif isinstance(event, FunctionToolResultEvent):
if isinstance(event.result, ToolReturnPart):
yield {
"type": "a",
"payload": {
"toolCallId": event.tool_call_id,
"result": event.result.content,
},
}
elif isinstance(event.result, RetryPromptPart):
yield {
"type": "a",
"payload": {
"toolCallId": event.tool_call_id,
"result": event.result.content,
},
}
else:
logger.warning(
"Unexpected tool result type: %s %s",
type(event.result),
dataclasses.asdict(event.result),
)
elif Agent.is_end_node(node):
# Once an End node is reached, the agent run is complete
logger.debug("Received end_node event: %s", dataclasses.asdict(node))
else:
logger.warning(
"Unknown node type encountered: %s",
type(node),
)
# Final usage summary
final_usage = run.usage()
usage["promptTokens"] = final_usage.request_tokens
usage["completionTokens"] = final_usage.response_tokens
# Persist conversation
await sync_to_async(self._update_conversation)(
history, run.result.new_messages(), run.result.new_messages_json(), usage
)
# Vercel finish message
yield {
"type": "d",
"payload": {
"finishReason": "stop",
"usage": usage,
},
}
def _update_conversation(
self,
history: List[ModelMessage],
final_output: List[ModelRequest | ModelMessage],
raw_final_output: bytes,
usage: Dict[str, int],
):
"""Persist messages + usage to DB (simplified)."""
_merged_final_output_request = None
_merged_final_output_message = None
_merged_final_output_request = ModelRequest(
parts=[
part for msg in final_output if isinstance(msg, ModelRequest) for part in msg.parts
],
kind="request",
)
_merged_final_output_message = ModelResponse(
parts=[
part for msg in final_output if isinstance(msg, ModelResponse) for part in msg.parts
],
kind="response",
)
self.conversation.messages = [
model_message_to_ui_message(msg)
for msg in chain(history, [_merged_final_output_request, _merged_final_output_message])
]
for message in self.conversation.messages:
logger.debug("conversation.messages: %s %s", type(message), message)
self.conversation.messages = [
msg.model_dump(mode="json") for msg in self.conversation.messages if msg
]
self.conversation.agent_usage = usage
logger.debug(
"raw_final_output: %s %s",
raw_final_output.decode("utf-8"),
json.loads(raw_final_output.decode("utf-8")),
)
self.conversation.openai_messages += json.loads(raw_final_output.decode("utf-8"))
self.conversation.save()
@@ -0,0 +1,275 @@
"""
Utility functions to convert between UIMessage (ai_sdk_types.py)
and UserContent/ModelMessage (pydantic_ai.messages.py).
"""
import base64
import json
import logging
from dataclasses import asdict
from typing import List
from pydantic_ai.messages import (
BinaryContent,
ModelMessage,
ModelRequest,
ModelRequestPart,
ModelResponse,
ModelResponsePart,
RetryPromptPart,
SystemPromptPart,
TextPart,
ThinkingPart,
ToolCallPart,
ToolReturnPart,
UserContent,
UserPromptPart,
)
from chat.ai_sdk_types import (
Attachment,
FileUIPart,
ReasoningDetailText,
ReasoningUIPart,
TextUIPart,
ToolInvocationCall,
ToolInvocationUIPart,
UIMessage,
UIPart,
)
def ui_message_to_model_message(message: UIMessage) -> ModelMessage: # noqa: PLR0912
"""
Convert a UIMessage to a ModelMessage (ModelRequest or ModelResponse) for Pydantic-AI.
"""
# pylint: disable=too-many-branches
parts_request: List[ModelRequestPart] = []
parts_response: List[ModelResponsePart] = []
for part in message.parts:
if isinstance(part, TextUIPart):
if message.role == "user":
parts_request.append(UserPromptPart(content=part.text, timestamp=message.createdAt))
elif message.role == "assistant":
parts_response.append(TextPart(content=part.text))
elif isinstance(part, ToolInvocationUIPart):
parts_response.append(
ToolCallPart(
tool_call_id=part.toolInvocation.toolCallId,
tool_name=part.toolInvocation.toolName,
args=part.toolInvocation.args,
)
)
elif isinstance(part, ReasoningUIPart):
parts_response.append(
ThinkingPart(
content=part.reasoning,
)
)
else:
raise ValueError(f"Unsupported UIPart type: {type(part)}")
# Handle experimental attachments
for experimental_attachment in message.experimental_attachments or []:
if experimental_attachment.url.startswith("data:"):
raw_data = base64.b64decode(experimental_attachment.url.split(",")[1])
if message.role == "user":
parts_request.append(
UserPromptPart(
content=[
BinaryContent(
data=raw_data, media_type=experimental_attachment.contentType
)
]
)
)
elif message.role == "assistant":
raise ValueError(
"Experimental attachments are not supported in assistant responses."
)
else:
raise ValueError(
f"Unsupported experimental attachment URL format: {experimental_attachment.url}"
)
if message.role == "user":
return ModelRequest(parts=parts_request, kind="request")
if message.role == "assistant":
return ModelResponse(parts=parts_response)
raise ValueError(f"Unsupported message role: {message.role}")
def ui_message_to_user_content(message: UIMessage) -> List[UserContent]:
"""
Convert a UIMessage to a list of UserContent for Pydantic-AI.
"""
user_contents: List[UserContent] = []
for part in message.parts:
if isinstance(part, TextUIPart):
user_contents.append(part.text)
elif isinstance(part, FileUIPart):
user_contents.append(
BinaryContent(data=part.data.encode("utf-8"), media_type=part.mimeType)
)
elif isinstance(part, ToolInvocationUIPart):
# Tool invocations are not directly mapped to UserContent, skip or handle as needed
continue
elif isinstance(part, ReasoningUIPart):
# Reasoning parts are not directly mapped to UserContent, skip or handle as needed
continue
else:
raise ValueError(f"Unsupported UIPart type: {type(part)}")
for experimental_attachment in message.experimental_attachments or []:
if experimental_attachment.url.startswith("data:"):
# Handle data URLs
raw_data = base64.b64decode(experimental_attachment.url.split(",")[1])
user_contents.append(
BinaryContent(data=raw_data, media_type=experimental_attachment.contentType)
)
else:
raise ValueError(
f"Unsupported experimental attachment URL format: {experimental_attachment.url}"
)
return user_contents
def model_message_to_ui_message(model_message: ModelMessage) -> UIMessage: # noqa: PLR0912
"""
Convert a ModelMessage (ModelRequest or ModelResponse) to a UIMessage.
"""
# pylint: disable=too-many-nested-blocks,too-many-branches
parts: List[UIPart] = []
experimental_attachments: List[Attachment] = []
logging.getLogger(__name__).debug(
"Converting ModelMessage to UIMessage: %s %s",
type(model_message),
asdict(model_message),
)
_states = {"tool-calls": {}}
if isinstance(model_message, ModelRequest):
message_timestamp = None
for part in model_message.parts:
if isinstance(part, SystemPromptPart):
# System prompts are not included in UIMessage parts
continue
if isinstance(part, UserPromptPart):
message_timestamp = part.timestamp
if isinstance(part.content, str):
parts.append(TextUIPart(type="text", text=part.content))
elif isinstance(part.content, list):
for c in part.content:
if isinstance(c, str):
parts.append(TextUIPart(type="text", text=c))
elif isinstance(c, BinaryContent):
experimental_attachments.append(
Attachment(
contentType=c.media_type,
url=f"data:{c.media_type};base64,"
+ base64.b64encode(c.data).decode("utf-8"),
)
)
else: # ImageUrl, AudioUrl, VideoUrl, DocumentUrl, BinaryContent
raise ValueError(
f"Unsupported UserContent in UserPromptPart: {type(c)}"
)
elif isinstance(part, TextPart):
parts.append(TextUIPart(type="text", text=part.content))
elif isinstance(part, ToolReturnPart):
pass
# parts.append(ToolInvocationUIPart(
# type="tool-invocation",
# toolInvocation=ToolInvocationResult(
# state="result",
# toolCallId=part.tool_call_id,
# toolName=part.tool_name,
# args={},
# result=part.content,
# )
# ))
elif isinstance(part, ThinkingPart):
parts.append(
ReasoningUIPart(
type="reasoning",
reasoning=part.content,
details=[
ReasoningDetailText(
type="text",
text=part.content,
signature=part.signature,
)
],
)
)
elif isinstance(part, RetryPromptPart):
# Retry prompts are not included in UIMessage parts
continue
else:
raise ValueError(f"Unsupported ModelRequest part type: {type(part)}")
if not parts:
return None
return UIMessage(
id="",
role="user",
content="".join(part.text for part in parts if isinstance(part, TextUIPart)),
parts=parts,
experimental_attachments=experimental_attachments or None,
createdAt=message_timestamp,
)
if isinstance(model_message, ModelResponse):
for part in model_message.parts:
if isinstance(part, UserPromptPart):
if isinstance(part.content, str):
parts.append(TextUIPart(type="text", text=part.content))
elif isinstance(part.content, list):
for c in part.content:
if isinstance(c, str):
parts.append(TextUIPart(type="text", text=c))
else: # ImageUrl, AudioUrl, VideoUrl, DocumentUrl, BinaryContent
raise ValueError(
f"Unsupported UserContent in UserPromptPart: {type(c)}"
)
elif isinstance(part, TextPart):
parts.append(TextUIPart(type="text", text=part.content))
elif isinstance(part, ToolCallPart):
parts.append(
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId=part.tool_call_id,
toolName=part.tool_name,
args=json.loads(part.args) if isinstance(part.args, str) else part.args,
),
)
)
elif isinstance(part, ThinkingPart):
parts.append(
ReasoningUIPart(
type="reasoning",
reasoning=part.content,
details=[
ReasoningDetailText(
type="text",
text=part.content,
signature=part.signature,
)
],
)
)
else:
raise ValueError(f"Unsupported ModelMessage part type: {type(part)}")
return UIMessage(
id="",
role="assistant",
content="".join(part.text for part in parts if isinstance(part, TextUIPart)),
parts=parts,
createdAt=model_message.timestamp,
)
raise ValueError(f"Unsupported ModelMessage part type: {type(model_message)}")
@@ -0,0 +1,265 @@
"""Tests for converting ModelMessage to UIMessage using Pydantic AI types."""
import datetime
import json
from django.utils import timezone
import pytest
from freezegun import freeze_time
from pydantic_ai.messages import (
AudioUrl,
BinaryContent,
DocumentUrl,
ModelRequest,
ModelResponse,
RetryPromptPart,
SystemPromptPart,
TextPart,
ThinkingPart,
ToolCallPart,
UserPromptPart,
VideoUrl,
)
from chat.ai_sdk_types import (
Attachment,
ReasoningDetailText,
ReasoningUIPart,
TextUIPart,
ToolInvocationCall,
ToolInvocationUIPart,
UIMessage,
)
from chat.clients.pydantic_ui_message_converter import model_message_to_ui_message
def test_model_message_to_ui_message_text_user_full():
"""Test converting a ModelRequest with UserPromptPart containing text to UIMessage."""
timestamp = datetime.datetime.now()
model_message = ModelRequest(
parts=[UserPromptPart(content="Hello!", timestamp=timestamp)], kind="request"
)
expected = UIMessage(
id="",
role="user",
content="Hello!",
parts=[TextUIPart(type="text", text="Hello!")],
createdAt=timestamp,
)
result = model_message_to_ui_message(model_message)
assert result == expected
@freeze_time()
def test_model_message_to_ui_message_text_assistant_full():
"""Test converting a ModelResponse with TextPart to UIMessage."""
model_message = ModelResponse(parts=[TextPart(content="Hi there!")])
expected = UIMessage(
id="",
role="assistant",
content="Hi there!",
parts=[TextUIPart(type="text", text="Hi there!")],
createdAt=timezone.now(),
)
result = model_message_to_ui_message(model_message)
assert result == expected
@freeze_time()
def test_model_message_to_ui_message_tool_call_full():
"""Test converting a ModelResponse with ToolCallPart to UIMessage."""
args = {"foo": "bar"}
model_message = ModelResponse(
parts=[ToolCallPart(tool_call_id="id1", tool_name="tool", args=args)]
)
expected = UIMessage(
id="",
role="assistant",
content="",
parts=[
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId="id1",
toolName="tool",
args=args,
),
)
],
createdAt=timezone.now(),
)
result = model_message_to_ui_message(model_message)
assert result == expected
@freeze_time()
def test_model_message_to_ui_message_reasoning_full():
"""Test converting a ModelResponse with ThinkingPart to UIMessage."""
model_message = ModelResponse(parts=[ThinkingPart(content="reason", signature="sig")])
expected = UIMessage(
id="",
role="assistant",
content="",
parts=[
ReasoningUIPart(
type="reasoning",
reasoning="reason",
details=[ReasoningDetailText(type="text", text="reason", signature="sig")],
)
],
createdAt=timezone.now(),
)
result = model_message_to_ui_message(model_message)
assert result.id == expected.id
assert result.role == expected.role
assert result.content == expected.content
assert result.createdAt == expected.createdAt
assert len(result.parts) == 1
parts_list = list(result.parts)
part = parts_list[0]
assert isinstance(part, ReasoningUIPart)
assert part.reasoning == "reason"
assert part.details[0].type == "text"
assert part.details[0].text == "reason"
assert part.details[0].signature == "sig"
def test_model_message_to_ui_message_binary_content():
"""Test converting a ModelRequest with BinaryContent to UIMessage."""
bin_data = b"bin"
model_message = ModelRequest(
parts=[
UserPromptPart(
content=[
"What do you see?",
BinaryContent(media_type="application/octet-stream", data=bin_data),
]
),
],
kind="request",
)
result = model_message_to_ui_message(model_message)
assert result.role == "user"
assert result.parts == [TextUIPart(type="text", text="What do you see?")]
assert result.experimental_attachments == [
Attachment(
name=None,
contentType="application/octet-stream",
url="data:application/octet-stream;base64,Ymlu",
),
]
def test_model_message_to_ui_message_file_parts_full():
"""Test handling unsupported file parts in UserPromptPart content."""
for part_type in [AudioUrl, VideoUrl, DocumentUrl]:
model_message = ModelRequest(
parts=[
UserPromptPart(
content=[
"Check this file",
part_type(url="http://example.com/file"),
],
timestamp=None,
),
],
kind="request",
)
with pytest.raises(ValueError, match="Unsupported UserContent in UserPromptPart"):
model_message_to_ui_message(model_message)
def test_model_message_to_ui_message_empty_parts():
"""Test converting a ModelRequest with no valid parts returns None."""
model_message = ModelRequest(parts=[], kind="request")
assert model_message_to_ui_message(model_message) is None
def test_model_message_to_ui_message_unsupported_part():
"""Test handling unsupported part types in ModelRequest."""
model_message = ModelRequest(parts=[SystemPromptPart(content="sys")], kind="request")
assert model_message_to_ui_message(model_message) is None
model_message = ModelRequest(parts=[RetryPromptPart(content="retry")], kind="request")
assert model_message_to_ui_message(model_message) is None
def test_model_message_to_ui_message_invalid_content_type():
"""Test handling invalid content type in UserPromptPart."""
class DummyContent:
"""Dummy class for testing invalid content types."""
model_message = ModelRequest(
parts=[UserPromptPart(content=[DummyContent()], timestamp=None)], kind="request"
)
with pytest.raises(ValueError, match="Unsupported UserContent in UserPromptPart"):
model_message_to_ui_message(model_message)
def test_model_message_to_ui_message_invalid_response_part():
"""Test handling invalid part type in ModelResponse."""
class DummyPart:
"""Dummy class for testing invalid part types."""
model_message = ModelResponse(parts=[DummyPart()])
with pytest.raises(ValueError, match="Unsupported ModelMessage part type"):
model_message_to_ui_message(model_message)
def test_model_message_to_ui_message_multiple_text_parts():
"""Test converting a ModelResponse with multiple TextParts to UIMessage."""
model_message = ModelResponse(parts=[TextPart(content="A"), TextPart(content="B")])
result = model_message_to_ui_message(model_message)
assert result.role == "assistant"
parts_list = list(result.parts)
assert [p.text for p in parts_list if isinstance(p, TextUIPart)] == ["A", "B"]
assert result.content == "AB"
def test_model_message_to_ui_message_userpromptpart_list_of_str():
"""Test converting a ModelRequest with UserPromptPart containing list of strings."""
model_message = ModelRequest(
parts=[UserPromptPart(content=["A", "B"], timestamp=None)], kind="request"
)
result = model_message_to_ui_message(model_message)
assert result.role == "user"
parts_list = list(result.parts)
assert [p.text for p in parts_list if isinstance(p, TextUIPart)] == ["A", "B"]
assert result.content == "AB"
def test_model_message_to_ui_message_tool_call_args_str():
"""Test converting a ModelResponse with ToolCallPart containing JSON string args."""
args = {"foo": "bar"}
model_message = ModelResponse(
parts=[ToolCallPart(tool_call_id="id1", tool_name="tool", args=json.dumps(args))]
)
result = model_message_to_ui_message(model_message)
parts_list = list(result.parts)
part = parts_list[0]
assert isinstance(part, ToolInvocationUIPart)
assert part.toolInvocation.args == args
def test_model_message_to_ui_message_with_reasoning_signature_none():
"""Test converting a ModelResponse with ThinkingPart having signature=None."""
model_message = ModelResponse(parts=[ThinkingPart(content="reason", signature=None)])
result = model_message_to_ui_message(model_message)
parts_list = list(result.parts)
part = parts_list[0]
assert isinstance(part, ReasoningUIPart)
assert part.details[0].signature is None
def test_model_message_to_ui_message_created_at_response():
"""Test converting a ModelResponse with a specific timestamp."""
model_message = ModelResponse(
parts=[TextPart(content="Hi!")], timestamp=datetime.datetime(2024, 1, 1, 12, 0, 0)
)
result = model_message_to_ui_message(model_message)
assert result.createdAt == datetime.datetime(2024, 1, 1, 12, 0, 0)
@@ -0,0 +1,365 @@
"""Tests for the conversion between UI messages and Pydantic AI messages."""
import base64
import datetime
import pytest
from pydantic_ai.messages import (
BinaryContent,
ModelRequest,
ModelResponse,
TextPart,
ThinkingPart,
ToolCallPart,
UserPromptPart,
)
from pydantic_ai.usage import Usage
from chat.ai_sdk_types import (
Attachment,
FileUIPart,
LanguageModelV1Source,
ReasoningUIPart,
SourceUIPart,
StepStartUIPart,
TextUIPart,
ToolInvocationCall,
ToolInvocationUIPart,
UIMessage,
)
from chat.clients.pydantic_ui_message_converter import ui_message_to_model_message
def test_user_message_with_text():
"""Test conversion of a user message with text only."""
timestamp = datetime.datetime.now()
ui_message = UIMessage(
id="msg1",
role="user",
content="Hello, how are you?",
parts=[TextUIPart(type="text", text="Hello, how are you?")],
createdAt=timestamp,
)
result = ui_message_to_model_message(ui_message)
assert isinstance(result, ModelRequest)
assert result.parts == [
UserPromptPart(content="Hello, how are you?", timestamp=timestamp),
]
assert result.instructions is None
def test_assistant_message_with_text():
"""Test conversion of an assistant message with text only."""
ui_message = UIMessage(
id="msg2",
role="assistant",
content="I'm doing well, thank you!",
parts=[TextUIPart(type="text", text="I'm doing well, thank you!")],
)
result = ui_message_to_model_message(ui_message)
assert isinstance(result, ModelResponse)
assert result.parts == [TextPart(content="I'm doing well, thank you!")]
assert result.usage == Usage()
assert result.timestamp is not None
def test_assistant_message_with_tool_call():
"""Test conversion of an assistant message with a tool call."""
tool_args = {"location": "Paris", "unit": "celsius"}
ui_message = UIMessage(
id="msg8",
role="assistant",
content="Let me check the weather for you.",
parts=[
TextUIPart(type="text", text="Let me check the weather for you."),
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId="call123",
toolName="get_weather",
args=tool_args,
),
),
],
)
result = ui_message_to_model_message(ui_message)
assert isinstance(result, ModelResponse)
assert len(result.parts) == 2
assert isinstance(result.parts[0], TextPart)
assert result.parts[0].content == "Let me check the weather for you."
assert isinstance(result.parts[1], ToolCallPart)
assert result.parts[1].tool_call_id == "call123"
assert result.parts[1].tool_name == "get_weather"
assert result.parts[1].args == tool_args
def test_assistant_message_with_reasoning():
"""Test conversion of an assistant message with reasoning."""
# Arrange
reasoning_text = "I need to think about this problem step by step..."
ui_message = UIMessage(
id="msg9",
role="assistant",
content="The answer is 42.",
parts=[
ReasoningUIPart(
type="reasoning",
reasoning=reasoning_text,
details=[],
),
TextUIPart(type="text", text="The answer is 42."),
],
)
# Act
result = ui_message_to_model_message(ui_message)
# Assert
assert isinstance(result, ModelResponse)
assert len(result.parts) == 2
assert isinstance(result.parts[0], ThinkingPart)
assert result.parts[0].content == reasoning_text
assert isinstance(result.parts[1], TextPart)
assert result.parts[1].content == "The answer is 42."
def test_multiple_text_parts():
"""Test conversion of a message with multiple text parts."""
# Arrange
ui_message = UIMessage(
id="msg10",
role="user",
content="Hello world! How are you today?",
parts=[
TextUIPart(type="text", text="Hello world!"),
TextUIPart(type="text", text=" How are you today?"),
],
)
# Act
result = ui_message_to_model_message(ui_message)
# Assert
assert isinstance(result, ModelRequest)
assert len(result.parts) == 2
assert isinstance(result.parts[0], UserPromptPart)
assert result.parts[0].content == "Hello world!"
assert isinstance(result.parts[1], UserPromptPart)
assert result.parts[1].content == " How are you today?"
def test_complex_message():
"""
Test conversion of a conversation with user and assistant messages,
including tool call and thinking.
"""
user_message = UIMessage(
id="msg_user",
role="user",
content="What's the weather in Paris?",
parts=[TextUIPart(type="text", text="What's the weather in Paris?")],
createdAt=datetime.datetime.now(),
)
user_result = ui_message_to_model_message(user_message)
assert isinstance(user_result, ModelRequest)
assert len(user_result.parts) == 1
assert isinstance(user_result.parts[0], UserPromptPart)
assert user_result.parts[0].content == "What's the weather in Paris?"
# Assistant message with text, tool call, and thinking
tool_args = {"location": "Paris", "unit": "celsius"}
reasoning_text = "Looking up the weather for Paris."
assistant_message = UIMessage(
id="msg_assistant",
role="assistant",
content="Let me check the weather for you.",
parts=[
TextUIPart(type="text", text="Let me check the weather for you."),
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId="call123",
toolName="get_weather",
args=tool_args,
),
),
ReasoningUIPart(
type="reasoning",
reasoning=reasoning_text,
details=[],
),
],
)
assistant_result = ui_message_to_model_message(assistant_message)
assert isinstance(assistant_result, ModelResponse)
assert len(assistant_result.parts) == 3
assert isinstance(assistant_result.parts[0], TextPart)
assert assistant_result.parts[0].content == "Let me check the weather for you."
assert isinstance(assistant_result.parts[1], ToolCallPart)
assert assistant_result.parts[1].tool_call_id == "call123"
assert assistant_result.parts[1].tool_name == "get_weather"
assert assistant_result.parts[1].args == tool_args
assert isinstance(assistant_result.parts[2], ThinkingPart)
assert assistant_result.parts[2].content == reasoning_text
def test_assistant_message_with_string_tool_args():
"""Test conversion of an assistant message with string tool arguments."""
ui_message = UIMessage(
id="msg14",
role="assistant",
content="Let me check the weather for you.",
parts=[
TextUIPart(type="text", text="Let me check the weather for you."),
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId="call123",
toolName="get_weather",
args={"location": "Paris", "unit": "celsius"},
),
),
],
)
result = ui_message_to_model_message(ui_message)
assert isinstance(result, ModelResponse)
assert len(result.parts) == 2
assert isinstance(result.parts[1], ToolCallPart)
assert result.parts[1].args == {"location": "Paris", "unit": "celsius"}
def test_user_message_with_attachment():
"""Test conversion of a user message with text only."""
timestamp = datetime.datetime.now()
ui_message = UIMessage(
id="msg1",
role="user",
content="Hello, how are you?",
parts=[TextUIPart(type="text", text="What do you see?")],
experimental_attachments=[
Attachment(
name="image.png",
contentType="image/png",
url=(
"data:image/png;base64,"
"iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ"
"3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII="
),
)
],
createdAt=timestamp,
)
result = ui_message_to_model_message(ui_message)
assert isinstance(result, ModelRequest)
assert len(result.parts) == 2
assert result.parts[0] == UserPromptPart(content="What do you see?", timestamp=timestamp)
assert isinstance(result.parts[1], UserPromptPart)
assert len(result.parts[1].content) == 1
assert isinstance(result.parts[1].content[0], BinaryContent)
assert result.parts[1].content[0].media_type == "image/png"
assert base64.b64encode(result.parts[1].content[0].data) == (
b"iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ"
b"3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII="
)
assert result.instructions is None
def test_assistant_message_with_attachment():
"""Test conversion of a user message with text only."""
timestamp = datetime.datetime.now()
ui_message = UIMessage(
id="msg1",
role="assistant",
content="Hello, how are you?",
parts=[TextUIPart(type="text", text="What do you see?")],
experimental_attachments=[
Attachment(
name="image.png",
contentType="image/png",
url=(
"data:image/png;base64,"
"iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ"
"3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII="
),
)
],
createdAt=timestamp,
)
with pytest.raises(ValueError, match="Experimental attachments are not supported"):
ui_message_to_model_message(ui_message)
@pytest.mark.parametrize("role", ("system", "data"))
def test_unsupported_role(role):
"""Test conversion with an unsupported role."""
# Arrange
ui_message = UIMessage(
id="msg12",
role=role, # Unsupported role
content="You are a helpful assistant",
parts=[TextUIPart(type="text", text="You are a helpful assistant")],
)
with pytest.raises(ValueError, match=f"Unsupported message role: {role}"):
ui_message_to_model_message(ui_message)
@pytest.mark.parametrize("role", ("user", "assistant"))
def test_message_with_source_part(role):
"""Test conversion of a user/assistant message with SourceUIPart (should raise)."""
ui_message = UIMessage(
id="msg_source_user",
role=role,
content="source info",
parts=[
SourceUIPart(
type="source", source=LanguageModelV1Source(source_type="test", details={})
)
],
)
with pytest.raises(
ValueError, match="Unsupported UIPart type: <class 'chat.ai_sdk_types.SourceUIPart'>"
):
ui_message_to_model_message(ui_message)
@pytest.mark.parametrize("role", ("user", "assistant"))
def test_message_with_step_start_part(role):
"""Test conversion of a user/assistant message with StepStartUIPart (should raise)."""
ui_message = UIMessage(
id="msg_step_user",
role=role,
content="step start",
parts=[StepStartUIPart(type="step-start")],
)
with pytest.raises(
ValueError, match="Unsupported UIPart type: <class 'chat.ai_sdk_types.StepStartUIPart'>"
):
ui_message_to_model_message(ui_message)
@pytest.mark.parametrize("role", ("user", "assistant"))
def test_assistant_message_with_file_part(role):
"""Test conversion of a user/assistant message with FileUIPart (should raise)."""
ui_message = UIMessage(
id="msg_file_assistant",
role=role,
content="file part",
parts=[FileUIPart(type="file", mimeType="image/png", data="http://example.com/image.png")],
)
with pytest.raises(
ValueError, match="Unsupported UIPart type: <class 'chat.ai_sdk_types.FileUIPart'>"
):
ui_message_to_model_message(ui_message)
@@ -0,0 +1,261 @@
"""Tests for the conversion from UIMessage to UserContent list."""
import base64
import datetime
import pytest
from pydantic_ai.messages import BinaryContent
from chat.ai_sdk_types import (
Attachment,
FileUIPart,
ReasoningUIPart,
TextUIPart,
ToolInvocationCall,
ToolInvocationUIPart,
UIMessage,
)
from chat.clients.pydantic_ui_message_converter import ui_message_to_user_content
def test_user_message_with_text():
"""Test conversion of a user message with text only."""
timestamp = datetime.datetime.now()
ui_message = UIMessage(
id="msg1",
role="user",
content="Hello, how are you?",
parts=[TextUIPart(type="text", text="Hello, how are you?")],
createdAt=timestamp,
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 1
assert isinstance(result[0], str)
assert result[0] == "Hello, how are you?"
def test_user_message_with_multiple_text_parts():
"""Test conversion of a user message with multiple text parts."""
ui_message = UIMessage(
id="msg2",
role="user",
content="Hello world! How are you today?",
parts=[
TextUIPart(type="text", text="Hello world!"),
TextUIPart(type="text", text=" How are you today?"),
],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 2
assert result[0] == "Hello world!"
assert result[1] == " How are you today?"
def test_user_message_with_file():
"""Test conversion of a user message with a file part."""
file_content = "This is a text file content"
mime_type = "text/plain"
ui_message = UIMessage(
id="msg3",
role="user",
content="Check this file",
parts=[
TextUIPart(type="text", text="Check this file"),
FileUIPart(type="file", data=file_content, mimeType=mime_type, name="example.txt"),
],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 2
assert result[0] == "Check this file"
assert isinstance(result[1], BinaryContent)
assert result[1].data == file_content.encode("utf-8")
assert result[1].media_type == mime_type
def test_user_message_with_experimental_attachment():
"""Test conversion of a user message with an experimental attachment."""
content_type = "image/png"
sample_data = b"sample image data"
base64_data = base64.b64encode(sample_data).decode("utf-8")
data_url = f"data:{content_type};base64,{base64_data}"
ui_message = UIMessage(
id="msg4",
role="user",
content="Check this image",
parts=[TextUIPart(type="text", text="Check this image")],
experimental_attachments=[
Attachment(
contentType=content_type,
url=data_url,
)
],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 2
assert result[0] == "Check this image"
assert isinstance(result[1], BinaryContent)
assert result[1].data == sample_data
assert result[1].media_type == content_type
def test_user_message_with_multiple_attachments():
"""Test conversion of a user message with multiple attachments of different types."""
# First attachment - text file
file_content = "This is a text file content"
file_mime_type = "text/plain"
# Second attachment - image
image_content_type = "image/png"
image_data = b"sample image data"
base64_image = base64.b64encode(image_data).decode("utf-8")
image_data_url = f"data:{image_content_type};base64,{base64_image}"
ui_message = UIMessage(
id="msg5",
role="user",
content="Check these files",
parts=[
TextUIPart(type="text", text="Check these files"),
FileUIPart(type="file", data=file_content, mimeType=file_mime_type, name="example.txt"),
],
experimental_attachments=[
Attachment(
contentType=image_content_type,
url=image_data_url,
)
],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 3
assert result[0] == "Check these files"
assert isinstance(result[1], BinaryContent)
assert result[1].data == file_content.encode("utf-8")
assert result[1].media_type == file_mime_type
assert isinstance(result[2], BinaryContent)
assert result[2].data == image_data
assert result[2].media_type == image_content_type
def test_user_message_with_tool_invocation():
"""Test conversion of a user message with a tool invocation part."""
tool_args = {"query": "weather in Paris", "unit": "celsius"}
ui_message = UIMessage(
id="msg6",
role="user",
content="Check the weather",
parts=[
TextUIPart(type="text", text="Check the weather"),
ToolInvocationUIPart(
type="tool-invocation",
toolInvocation=ToolInvocationCall(
state="call",
toolCallId="call123",
toolName="get_weather",
args=tool_args,
),
),
],
)
result = ui_message_to_user_content(ui_message)
# Tool invocation parts are skipped in the conversion
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == "Check the weather"
def test_user_message_with_reasoning():
"""Test conversion of a user message with a reasoning part."""
reasoning_text = "I need to think about this..."
ui_message = UIMessage(
id="msg7",
role="user",
content="Let me think",
parts=[
TextUIPart(type="text", text="Let me think"),
ReasoningUIPart(
type="reasoning",
reasoning=reasoning_text,
details=[],
),
],
)
result = ui_message_to_user_content(ui_message)
# Reasoning parts are skipped in the conversion
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == "Let me think"
def test_unsupported_experimental_attachment_url():
"""Test error handling for unsupported experimental attachment URL format."""
ui_message = UIMessage(
id="msg8",
role="user",
content="Check this file",
parts=[TextUIPart(type="text", text="Check this file")],
experimental_attachments=[
Attachment(
contentType="text/plain",
url="https://example.com/file.txt", # Not a data URL
)
],
)
with pytest.raises(ValueError) as excinfo:
ui_message_to_user_content(ui_message)
assert "Unsupported experimental attachment URL format" in str(excinfo.value)
def test_empty_message():
"""Test conversion of a user message with no parts."""
ui_message = UIMessage(
id="msg10",
role="user",
content="",
parts=[],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 0
def test_assistant_message():
"""Test conversion of an assistant message."""
ui_message = UIMessage(
id="msg11",
role="assistant",
content="I'm an assistant",
parts=[TextUIPart(type="text", text="I'm an assistant")],
)
result = ui_message_to_user_content(ui_message)
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == "I'm an assistant"
@@ -1,10 +1,14 @@
"""Unit tests for chat conversation actions in the chat API view."""
# pylint: disable=too-many-lines
import json
from django.utils import timezone
import httpx
import pytest
import respx
from freezegun import freeze_time
from rest_framework import status
from core.factories import UserFactory
@@ -28,6 +32,7 @@ def ai_settings(settings):
@pytest.fixture(name="mock_openai_stream")
@freeze_time("2025-07-25T10:36:35.297675Z")
def fixture_mock_openai_stream():
"""
Fixture to mock the OpenAI stream response.
@@ -38,6 +43,8 @@ def fixture_mock_openai_stream():
"data: "
+ json.dumps(
{
"id": "chatcmpl-1234567890",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{
"delta": {"content": "Hello"},
@@ -52,6 +59,8 @@ def fixture_mock_openai_stream():
"data: "
+ json.dumps(
{
"id": "chatcmpl-1234567890",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{
"delta": {"content": " there"},
@@ -60,6 +69,11 @@ def fixture_mock_openai_stream():
}
],
"object": "chat.completion.chunk",
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
}
)
+ "\n\n"
@@ -78,6 +92,7 @@ def fixture_mock_openai_stream():
@pytest.fixture(name="mock_openai_stream_image")
@freeze_time("2025-07-25T10:36:35.297675Z")
def fixture_mock_openai_stream_image():
"""
Mock a very simple OpenAI stream that *mentions* the image
@@ -88,6 +103,8 @@ def fixture_mock_openai_stream_image():
"data: "
+ json.dumps(
{
"id": "chatcmpl-1234567890",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{
"delta": {"content": "I see a cat"},
@@ -102,6 +119,8 @@ def fixture_mock_openai_stream_image():
"data: "
+ json.dumps(
{
"id": "chatcmpl-1234567890",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{
"delta": {"content": " in the picture."},
@@ -110,6 +129,11 @@ def fixture_mock_openai_stream_image():
}
],
"object": "chat.completion.chunk",
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
}
)
+ "\n\n"
@@ -127,6 +151,7 @@ def fixture_mock_openai_stream_image():
@pytest.fixture(name="mock_openai_stream_tool")
@freeze_time("2025-07-25T10:36:35.297675Z")
def fixture_mock_openai_stream_tool():
"""
Mock both API calls in the tool call flow:
@@ -140,6 +165,7 @@ def fixture_mock_openai_stream_tool():
+ json.dumps(
{
"id": "chatcmpl-tool-call",
"created": timezone.make_naive(timezone.now()).timestamp(),
"object": "chat.completion.chunk",
"choices": [
{
@@ -152,6 +178,33 @@ def fixture_mock_openai_stream_tool():
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": "",
},
}
]
},
}
],
}
)
+ "\n\n"
"data: "
+ json.dumps(
{
"id": "chatcmpl-tool-call",
"created": timezone.make_naive(timezone.now()).timestamp(),
"object": "chat.completion.chunk",
"choices": [
{
"index": 0,
"delta": {
"tool_calls": [
{
"index": 0,
"id": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"type": "function",
"function": {
"name": "",
"arguments": '{"location":"Paris", "unit":"celsius"}',
},
}
@@ -166,7 +219,13 @@ def fixture_mock_openai_stream_tool():
+ json.dumps(
{
"id": "chatcmpl-tool-call",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [{"delta": {}, "finish_reason": "tool_calls"}],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
}
)
+ "\n\n"
@@ -179,6 +238,7 @@ def fixture_mock_openai_stream_tool():
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"object": "chat.completion.chunk",
"choices": [{"delta": {"role": "assistant"}, "index": 0}],
}
@@ -188,9 +248,15 @@ def fixture_mock_openai_stream_tool():
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{"delta": {"content": "The current weather in Paris is nice"}, "index": 0}
],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
}
)
+ "\n\n"
@@ -198,6 +264,47 @@ def fixture_mock_openai_stream_tool():
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [{"delta": {}, "finish_reason": "stop"}],
}
)
+ "\n\n"
"data: [DONE]\n\n"
)
# Second response - final answer when failing
second_response_fail = (
"data: "
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"object": "chat.completion.chunk",
"choices": [{"delta": {"role": "assistant"}, "index": 0}],
}
)
+ "\n\n"
"data: "
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [
{"delta": {"content": "I cannot give you an answer to that."}, "index": 0}
],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
},
}
)
+ "\n\n"
"data: "
+ json.dumps(
{
"id": "chatcmpl-final",
"created": timezone.make_naive(timezone.now()).timestamp(),
"choices": [{"delta": {}, "finish_reason": "stop"}],
}
)
@@ -213,10 +320,20 @@ def fixture_mock_openai_stream_tool():
for line in second_response.splitlines(keepends=True):
yield line.encode()
async def mock_second_response_failing_stream():
for line in second_response_fail.splitlines(keepends=True):
yield line.encode()
def tool_answer_side_effect(request):
if "Unknown tool name:" in request.content.decode():
# Simulate the second response with tool call failure
return httpx.Response(200, stream=mock_second_response_failing_stream())
return httpx.Response(200, stream=mock_second_response_stream())
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
side_effect=[
httpx.Response(200, stream=mock_first_response_stream()),
httpx.Response(200, stream=mock_second_response_stream()),
tool_answer_side_effect,
]
)
@@ -269,6 +386,7 @@ def test_post_conversation_invalid_protocol(api_client):
assert "Invalid protocol" in response.data["error"]
@freeze_time("2025-07-25T10:36:35.297675Z")
@respx.mock
def test_post_conversation_data_protocol(api_client, mock_openai_stream):
"""Test posting messages to a conversation using the 'data' protocol."""
@@ -317,27 +435,70 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
]
assert len(chat_conversation.messages) == 2
assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[0] == {
"annotations": None,
"content": "Hello",
"createdAt": "2025-07-03T15:22:17.105Z",
"id": "yuPoOuBkKA4FnKvk",
"experimental_attachments": None,
"id": "", # ID is not set in the response
"parts": [{"text": "Hello", "type": "text"}],
"reasoning": None,
"role": "user",
"toolInvocations": None,
}
assert chat_conversation.messages[1].pop("id") # Remove ID for comparison
assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[1] == {
"annotations": None,
"content": "Hello there",
"createdAt": None,
"experimental_attachments": None,
"id": "", # ID is not set in the response
"parts": [{"text": "Hello there", "type": "text"}],
"reasoning": None,
"role": "assistant",
"toolInvocations": None,
}
assert chat_conversation.openai_messages == [
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "You are a helpful assistant. Escape formulas or any "
"math notation between `$$`, like `$$x^2 + y^2 = "
"z^2$$` or `$$C_l$$`. You can use Markdown to format "
"your answers. ",
"dynamic_ref": None,
"part_kind": "system-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
{
"content": ["Hello"],
"part_kind": "user-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [{"content": "Hello there", "part_kind": "text"}],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
]
@freeze_time("2025-07-25T10:36:35.297675Z")
@respx.mock
def test_post_conversation_text_protocol(api_client, mock_openai_stream):
"""Test posting messages to a conversation using the 'text' protocol."""
@@ -380,27 +541,70 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
]
assert len(chat_conversation.messages) == 2
assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[0] == {
"annotations": None,
"content": "Hello",
"createdAt": "2025-07-03T15:22:17.105Z",
"id": "yuPoOuBkKA4FnKvk",
"experimental_attachments": None,
"id": "", # ID is not set in the response
"parts": [{"text": "Hello", "type": "text"}],
"reasoning": None,
"role": "user",
"toolInvocations": None,
}
assert chat_conversation.messages[1].pop("id")
assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[1] == {
"annotations": None,
"content": "Hello there",
"createdAt": None,
"experimental_attachments": None,
"id": "", # ID is not set in the response
"parts": [{"text": "Hello there", "type": "text"}],
"reasoning": None,
"role": "assistant",
"toolInvocations": None,
}
assert chat_conversation.openai_messages == [
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "You are a helpful assistant. Escape formulas or any "
"math notation between `$$`, like `$$x^2 + y^2 = "
"z^2$$` or `$$C_l$$`. You can use Markdown to format "
"your answers. ",
"dynamic_ref": None,
"part_kind": "system-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
{
"content": ["Hello"],
"part_kind": "user-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [{"content": "Hello there", "part_kind": "text"}],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
]
@freeze_time("2025-07-25T10:36:35.297675Z")
@respx.mock
def test_post_conversation_with_image(api_client, mock_openai_stream_image):
"""Ensure an image URL is correctly forwarded to the AI service."""
@@ -417,9 +621,13 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
"createdAt": "2025-07-07T15:52:27.822Z",
"experimental_attachments": [
{
"url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD",
"url": (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAA"
"ABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5C"
"YII="
),
"name": "FELV-cat.jpg",
"contentType": "image/jpeg",
"contentType": "image/png",
}
],
}
@@ -460,8 +668,12 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
{"text": "Hello, what do you see on this picture?", "type": "text"},
{
"image_url": {
"detail": "auto",
"url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD",
# "detail": "auto",
"url": (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAA"
"ABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5C"
"YII="
),
},
"type": "image_url",
},
@@ -472,7 +684,114 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
assert body["model"] == "test-model"
assert body["stream"] is True
chat_conversation.refresh_from_db()
assert chat_conversation.ui_messages == [
{
"content": "Hello, what do you see on this picture?",
"createdAt": "2025-07-07T15:52:27.822Z",
"experimental_attachments": [
{
"contentType": "image/png",
"name": "FELV-cat.jpg",
"url": (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAA"
"ABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5C"
"YII="
),
}
],
"id": "7x3hLsq6rB3xp91T",
"parts": [{"text": "Hello, what do you see on this picture?", "type": "text"}],
"role": "user",
}
]
assert len(chat_conversation.messages) == 2
assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[0] == {
"annotations": None,
"content": "Hello, what do you see on this picture?",
"experimental_attachments": [
{
"contentType": "image/png",
"name": None,
"url": (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wS"
"zIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAA"
"AASUVORK5CYII="
),
}
],
"id": "",
"parts": [{"text": "Hello, what do you see on this picture?", "type": "text"}],
"reasoning": None,
"role": "user",
"toolInvocations": None,
}
assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[1] == {
"annotations": None,
"content": "I see a cat in the picture.",
"experimental_attachments": None,
"id": "",
"parts": [{"text": "I see a cat in the picture.", "type": "text"}],
"reasoning": None,
"role": "assistant",
"toolInvocations": None,
}
assert chat_conversation.openai_messages == [
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "You are a helpful assistant. Escape formulas or any "
"math notation between `$$`, like `$$x^2 + y^2 = "
"z^2$$` or `$$C_l$$`. You can use Markdown to format "
"your answers. ",
"dynamic_ref": None,
"part_kind": "system-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
{
"content": [
"Hello, what do you see on this picture?",
{
"data": (
"iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD-wSzIAAAABlBMVEX___-_"
"v7-jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD_aNpbtEAAAAASUVORK5CYII="
),
"kind": "binary",
"media_type": "image/png",
"vendor_metadata": None,
},
],
"part_kind": "user-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [{"content": "I see a cat in the picture.", "part_kind": "text"}],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
]
@freeze_time("2025-07-25T10:36:35.297675Z")
@respx.mock
def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settings):
"""Ensure tool calls are correctly forwarded and streamed back."""
@@ -504,10 +823,12 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
# Wait for the streaming content to be fully received
response_content = b"".join(response.streaming_content).decode("utf-8")
assert response_content == (
'9:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "toolName": '
'"get_current_weather", "args": {"location": "Paris", "unit": "celsius"}}\n'
'a:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "result": '
"\"{'location': 'Paris', 'temperature': 22, 'unit': 'celsius'}\"}\n"
'b:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "toolName": '
'"get_current_weather"}\n'
'c:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "argsTextDelta": '
'"{\\"location\\":\\"Paris\\", \\"unit\\":\\"celsius\\"}"}\n'
'a:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "result": {"location": '
'"Paris", "temperature": 22, "unit": "celsius"}}\n'
'0:"The current weather in Paris is nice"\n'
'd:{"finishReason": "stop", "usage": {"promptTokens": 0, "completionTokens": '
"0}}\n"
@@ -539,60 +860,118 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
]
assert len(chat_conversation.messages) == 2
assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[0] == {
"annotations": None,
"content": "Weather in Paris?",
"createdAt": "2025-07-18T12:00:00Z",
"id": "tool-msg-1",
"experimental_attachments": None,
"id": "",
"parts": [{"text": "Weather in Paris?", "type": "text"}],
"reasoning": None,
"role": "user",
"toolInvocations": None,
}
assert chat_conversation.messages[1].pop("id")
assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[1] == {
"annotations": None,
"content": "The current weather in Paris is nice",
"createdAt": None,
"experimental_attachments": None,
"parts": [{"text": "The current weather in Paris is nice", "type": "text"}],
"id": "",
"parts": [
{
"toolInvocation": {
"args": {"location": "Paris", "unit": "celsius"},
"state": "call",
"step": None,
"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"toolName": "get_current_weather",
},
"type": "tool-invocation",
},
{"text": "The current weather in Paris is nice", "type": "text"},
],
"reasoning": None,
"role": "assistant",
"toolInvocations": None,
}
# To be fixed, because in real life, the tool invocation is added to the message...
# assert chat_conversation.messages[1] == {
# "annotations": None,
# "content": "The weather is sunny",
# "createdAt": None,
# "experimental_attachments": None,
# "parts": [
# {
# "type": "tool-invocation",
# "toolInvocation": {
# "args": {"unit": "celsius", "location": "Paris"},
# "step": 0,
# "state": "result",
# "result": "{'location': 'Paris', 'temperature': 22, 'unit': 'celsius'}",
# "toolName": "get_current_weather",
# "toolCallId": "FCBUEY5SpcsaB72P9taJR7Bcx0bAuqOu",
# },
# },
# {"text": "The weather is sunny", "type": "text"},
# ],
# "reasoning": None,
# "role": "assistant",
# "toolInvocations": [
# {
# "args": {"unit": "celsius", "location": "Paris"},
# "step": 0,
# "state": "result",
# "result": "{'location': 'Paris', 'temperature': 22, 'unit': 'celsius'}",
# "toolName": "get_current_weather",
# "toolCallId": "FCBUEY5SpcsaB72P9taJR7Bcx0bAuqOu",
# }
# ],
# }
assert chat_conversation.openai_messages == [
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "You are a helpful assistant. Escape formulas or any "
"math notation between `$$`, like `$$x^2 + y^2 = "
"z^2$$` or `$$C_l$$`. You can use Markdown to format "
"your answers. ",
"dynamic_ref": None,
"part_kind": "system-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
{
"content": ["Weather in Paris?"],
"part_kind": "user-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [
{
"args": '{"location":"Paris", "unit":"celsius"}',
"part_kind": "tool-call",
"tool_call_id": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"tool_name": "get_current_weather",
}
],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": {"location": "Paris", "temperature": 22, "unit": "celsius"},
"metadata": None,
"part_kind": "tool-return",
"timestamp": "2025-07-25T10:36:35.297675Z",
"tool_call_id": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"tool_name": "get_current_weather",
}
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [{"content": "The current weather in Paris is nice", "part_kind": "text"}],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
]
@freeze_time("2025-07-25T10:36:35.297675Z")
@respx.mock
def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool, settings):
"""Ensure tool calls are correctly forwarded and streamed back when failing."""
@@ -624,8 +1003,14 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
# Wait for the streaming content to be fully received
response_content = b"".join(response.streaming_content).decode("utf-8")
assert response_content == (
'3:"Tool get_current_weather not found in agent Conversations Assistant"\n'
'd:{"finishReason": "error", "usage": {"promptTokens": 0, "completionTokens": '
'b:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "toolName": '
'"get_current_weather"}\n'
'c:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "argsTextDelta": '
'"{\\"location\\":\\"Paris\\", \\"unit\\":\\"celsius\\"}"}\n'
'a:{"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", "result": "Unknown tool '
"name: 'get_current_weather'. No tools available.\"}\n"
'0:"I cannot give you an answer to that."\n'
'd:{"finishReason": "stop", "usage": {"promptTokens": 0, "completionTokens": '
"0}}\n"
)
@@ -654,11 +1039,112 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
}
]
assert len(chat_conversation.messages) == 1
assert len(chat_conversation.messages) == 2
assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[0] == {
"annotations": None,
"content": "Weather in Paris?",
"createdAt": "2025-07-18T12:00:00Z",
"id": "tool-msg-1",
"experimental_attachments": None,
"id": "",
"parts": [{"text": "Weather in Paris?", "type": "text"}],
"reasoning": None,
"role": "user",
"toolInvocations": None,
}
assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison
assert chat_conversation.messages[1] == {
"annotations": None,
"content": "I cannot give you an answer to that.",
"experimental_attachments": None,
"id": "",
"parts": [
{
"toolInvocation": {
"args": {"location": "Paris", "unit": "celsius"},
"state": "call",
"step": None,
"toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"toolName": "get_current_weather",
},
"type": "tool-invocation",
},
{"text": "I cannot give you an answer to that.", "type": "text"},
],
"reasoning": None,
"role": "assistant",
"toolInvocations": None,
}
assert chat_conversation.openai_messages == [
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "You are a helpful assistant. Escape formulas or any "
"math notation between `$$`, like `$$x^2 + y^2 = "
"z^2$$` or `$$C_l$$`. You can use Markdown to format "
"your answers. ",
"dynamic_ref": None,
"part_kind": "system-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
{
"content": ["Weather in Paris?"],
"part_kind": "user-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
},
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [
{
"args": '{"location":"Paris", "unit":"celsius"}',
"part_kind": "tool-call",
"tool_call_id": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"tool_name": "get_current_weather",
}
],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
{
"instructions": None,
"kind": "request",
"parts": [
{
"content": "Unknown tool name: 'get_current_weather'. No tools available.",
"part_kind": "retry-prompt",
"timestamp": "2025-07-25T10:36:35.297675Z",
"tool_call_id": "xLDcIljdsDrz0idal7tATWSMm2jhMj47",
"tool_name": "get_current_weather",
}
],
},
{
"kind": "response",
"model_name": "test-model",
"parts": [{"content": "I cannot give you an answer to that.", "part_kind": "text"}],
"timestamp": "2025-07-25T10:36:35.297675Z",
"usage": {
"details": None,
"request_tokens": 0,
"requests": 1,
"response_tokens": 0,
"total_tokens": 0,
},
"vendor_details": None,
"vendor_id": None,
},
]
+16 -3
View File
@@ -1,9 +1,12 @@
"""Tools for the chat agent."""
from agents import FunctionTool
from django.conf import settings
from .fake_current_weather import agent_get_current_weather
from .web_search_tavily import agent_web_search_tavily
from agents import FunctionTool
from pydantic_ai import Agent, Tool
from .fake_current_weather import agent_get_current_weather, get_current_weather
from .web_search_tavily import agent_web_search_tavily, tavily_web_search
def get_tool_by_name(name: str) -> FunctionTool:
@@ -14,3 +17,13 @@ def get_tool_by_name(name: str) -> FunctionTool:
}
return tool_dict[name] # will raise on purpose if name is not found
def get_pydantic_tools_by_name(name: str) -> Tool:
"""Get a Pydantic AI agent by its name."""
tool_dict = {
"get_current_weather": Tool(get_current_weather, takes_ctx=False),
"tavily_web_search": Tool(tavily_web_search, takes_ctx=False),
}
return tool_dict[name] # will raise on purpose if name is not found
+10 -1
View File
@@ -25,7 +25,16 @@ current_weather = ChatCompletionToolParam(
def get_current_weather(location: str, unit: str):
"""Get the current weather in a given location."""
"""
Get the current weather in a given location.
Args:
location (str): The city and state, e.g. San Francisco, CA.
unit (str): The unit of temperature, either 'celsius' or 'fahrenheit'.
Returns:
dict: A dictionary containing the location, temperature, and unit.
"""
return {
"location": location,
"temperature": 22 if unit == "celsius" else 72,
+1 -1
View File
@@ -13,7 +13,7 @@ from core.filters import remove_accents
from chat import models, serializers
from chat.ai_sdk_types import UIMessage
from chat.clients.openai import AIAgentService
from chat.clients.pydantic_ai import AIAgentService
logger = logging.getLogger(__name__)
+1
View File
@@ -77,6 +77,7 @@ dev = [
"freezegun==1.5.2",
"ipdb==0.13.13",
"ipython==9.3.0",
"pydantic-ai==0.4.3",
"pyfakefs==5.8.0",
"pylint-django==2.6.1",
"pylint==3.3.7",
@@ -1,4 +1,8 @@
import { Message, ToolInvocationUIPart } from '@ai-sdk/ui-utils';
import {
Message,
ReasoningUIPart,
ToolInvocationUIPart,
} from '@ai-sdk/ui-utils';
import { Loader } from '@openfun/cunningham-react';
import 'katex/dist/katex.min.css'; // `rehype-katex` does not import the CSS for you
import Image from 'next/image';
@@ -229,22 +233,39 @@ export const Chat = ({
{message.content}
</Markdown>
)}
<Box $direction="row" $gap="2">
<Box $direction="column">
{message.parts
?.filter((part) => part.type === 'tool-invocation')
.map((part: ToolInvocationUIPart) => (
<Box
as="pre"
key={part.toolInvocation.toolCallId}
$background="var(--c--theme--colors--greyscale-100)"
$color="var(--c--theme--colors--greyscale-500)"
$padding={{ all: 'sm' }}
$radius="md"
$css="font-family: monospace; font-size: 0.9em;"
>
{`${part.toolInvocation.toolName}(${JSON.stringify(part.toolInvocation.args, null, 2)})`}
</Box>
))}
?.filter(
(part) =>
part.type === 'reasoning' ||
part.type === 'tool-invocation',
)
.map((part: ReasoningUIPart | ToolInvocationUIPart) =>
part.type === 'reasoning' ? (
<Box
key={part.reasoning}
$background="var(--c--theme--colors--greyscale-100)"
$color="var(--c--theme--colors--greyscale-500)"
$padding={{ all: 'sm' }}
$radius="md"
$css="font-size: 0.9em;"
>
{part.reasoning}
</Box>
) : part.type === 'tool-invocation' ? (
<Box
as="pre"
key={part.toolInvocation.toolCallId}
$background="var(--c--theme--colors--greyscale-100)"
$color="var(--c--theme--colors--greyscale-500)"
$padding={{ all: 'sm' }}
$radius="md"
$css="font-family: monospace; font-size: 0.9em;"
>
{`${part.toolInvocation.toolName}(${JSON.stringify(part.toolInvocation.args, null, 2)})`}
</Box>
) : null,
)}
{/* Show attachments if present */}
{message.experimental_attachments?.map(
(attachment: Attachment, index: number) =>