diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f480d4..1d1763b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/src/backend/chat/clients/async_to_sync.py b/src/backend/chat/clients/async_to_sync.py new file mode 100644 index 0000000..9436dbc --- /dev/null +++ b/src/backend/chat/clients/async_to_sync.py @@ -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() diff --git a/src/backend/chat/clients/openai.py b/src/backend/chat/clients/openai.py index eaf66bd..ee44b4f 100644 --- a/src/backend/chat/clients/openai.py +++ b/src/backend/chat/clients/openai.py @@ -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.""" diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py new file mode 100644 index 0000000..3fe4f39 --- /dev/null +++ b/src/backend/chat/clients/pydantic_ai.py @@ -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() diff --git a/src/backend/chat/clients/pydantic_ui_message_converter.py b/src/backend/chat/clients/pydantic_ui_message_converter.py new file mode 100644 index 0000000..c398a34 --- /dev/null +++ b/src/backend/chat/clients/pydantic_ui_message_converter.py @@ -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)}") diff --git a/src/backend/chat/tests/clients/pydantic_ui_message_converter/__init__.py b/src/backend/chat/tests/clients/pydantic_ui_message_converter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_model_message_to_ui_message.py b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_model_message_to_ui_message.py new file mode 100644 index 0000000..de5ed8e --- /dev/null +++ b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_model_message_to_ui_message.py @@ -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) diff --git a/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_model_message.py b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_model_message.py new file mode 100644 index 0000000..8967e8b --- /dev/null +++ b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_model_message.py @@ -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: " + ): + 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: " + ): + 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: " + ): + ui_message_to_model_message(ui_message) diff --git a/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_user_content.py b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_user_content.py new file mode 100644 index 0000000..a11e8a1 --- /dev/null +++ b/src/backend/chat/tests/clients/pydantic_ui_message_converter/test_ui_message_to_user_content.py @@ -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" diff --git a/src/backend/chat/tests/views/chat/test_conversation.py b/src/backend/chat/tests/views/chat/test_conversation.py index 681b9e1..f064b3f 100644 --- a/src/backend/chat/tests/views/chat/test_conversation.py +++ b/src/backend/chat/tests/views/chat/test_conversation.py @@ -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, + }, + ] diff --git a/src/backend/chat/tools/__init__.py b/src/backend/chat/tools/__init__.py index aaa720c..63e9c6a 100644 --- a/src/backend/chat/tools/__init__.py +++ b/src/backend/chat/tools/__init__.py @@ -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 diff --git a/src/backend/chat/tools/fake_current_weather.py b/src/backend/chat/tools/fake_current_weather.py index dd947a8..aec190d 100644 --- a/src/backend/chat/tools/fake_current_weather.py +++ b/src/backend/chat/tools/fake_current_weather.py @@ -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, diff --git a/src/backend/chat/views.py b/src/backend/chat/views.py index 1407fe7..786eb81 100644 --- a/src/backend/chat/views.py +++ b/src/backend/chat/views.py @@ -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__) diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 3b833de..4cba562 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -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", diff --git a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx index 4177945..9b6d465 100644 --- a/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx +++ b/src/frontend/apps/conversations/src/features/chat/components/Chat.tsx @@ -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} )} - + {message.parts - ?.filter((part) => part.type === 'tool-invocation') - .map((part: ToolInvocationUIPart) => ( - - {`${part.toolInvocation.toolName}(${JSON.stringify(part.toolInvocation.args, null, 2)})`} - - ))} + ?.filter( + (part) => + part.type === 'reasoning' || + part.type === 'tool-invocation', + ) + .map((part: ReasoningUIPart | ToolInvocationUIPart) => + part.type === 'reasoning' ? ( + + {part.reasoning} + + ) : part.type === 'tool-invocation' ? ( + + {`${part.toolInvocation.toolName}(${JSON.stringify(part.toolInvocation.args, null, 2)})`} + + ) : null, + )} {/* Show attachments if present */} {message.experimental_attachments?.map( (attachment: Attachment, index: number) =>