⚡️(pydantic) use stored history in agent requests
We were converting the whole message history from the frontend on each call. We replace that by direct use of the history stored in the proper format from the database. This important to do so because the raw Pydantic messages may contain more information than the ones displayed to the user (like when doing RAG web search).
This commit is contained in:
@@ -51,7 +51,6 @@ from chat.ai_sdk_types import (
|
||||
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
|
||||
@@ -178,7 +177,7 @@ class AIAgentService:
|
||||
if messages[-1].role != "user":
|
||||
return
|
||||
|
||||
history = [ui_message_to_model_message(message) for message in messages[:-1]]
|
||||
history = ModelMessagesTypeAdapter.validate_python(self.conversation.openai_messages)
|
||||
prompt = ui_message_to_user_content(messages[-1])
|
||||
usage = {"promptTokens": 0, "completionTokens": 0}
|
||||
|
||||
|
||||
@@ -14,9 +14,7 @@ from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelRequestPart,
|
||||
ModelResponse,
|
||||
ModelResponsePart,
|
||||
RetryPromptPart,
|
||||
SystemPromptPart,
|
||||
TextPart,
|
||||
@@ -40,64 +38,6 @@ from chat.ai_sdk_types import (
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
-365
@@ -1,365 +0,0 @@
|
||||
"""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,329 @@
|
||||
"""Common test fixtures for chat conversation endpoint tests."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from freezegun import freeze_time
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_uuid4")
|
||||
def mock_uuid4_fixture():
|
||||
"""Fixture to mock UUID generation for testing."""
|
||||
value = uuid.uuid4()
|
||||
with patch("uuid.uuid4", return_value=value):
|
||||
yield value
|
||||
|
||||
|
||||
@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.
|
||||
|
||||
See https://platform.openai.com/docs/api-reference/chat-streaming/streaming
|
||||
"""
|
||||
openai_stream = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": "Hello"},
|
||||
"index": 0,
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": " there"},
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_stream():
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
return_value=httpx.Response(200, stream=mock_stream())
|
||||
)
|
||||
|
||||
return route
|
||||
|
||||
|
||||
@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
|
||||
in its textual reply (the real test is that the image URL is
|
||||
forwarded in the request body to the AI service).
|
||||
"""
|
||||
openai_stream = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": "I see a cat"},
|
||||
"index": 0,
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": " in the picture."},
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_stream():
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
return_value=httpx.Response(200, stream=mock_stream())
|
||||
)
|
||||
return route
|
||||
|
||||
|
||||
@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:
|
||||
1. First call returns function call
|
||||
2. Second call returns final answer after tool execution
|
||||
"""
|
||||
|
||||
# First response - tool call
|
||||
first_response = (
|
||||
"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": "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"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ 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"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
# Second response - final answer
|
||||
second_response = (
|
||||
"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": "The current weather in Paris is nice"}, "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"}],
|
||||
}
|
||||
)
|
||||
+ "\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"}],
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_first_response_stream():
|
||||
for line in first_response.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
async def mock_second_response_stream():
|
||||
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()),
|
||||
tool_answer_side_effect,
|
||||
]
|
||||
)
|
||||
|
||||
return route
|
||||
-320
@@ -2,12 +2,9 @@
|
||||
# pylint: disable=too-many-lines
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from freezegun import freeze_time
|
||||
@@ -43,323 +40,6 @@ def ai_settings(settings):
|
||||
return settings
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_uuid4")
|
||||
def mock_uuid4_fixture():
|
||||
"""Fixture to mock UUID generation for testing."""
|
||||
value = uuid.uuid4()
|
||||
with patch("uuid.uuid4", return_value=value):
|
||||
yield value
|
||||
|
||||
|
||||
@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.
|
||||
|
||||
See https://platform.openai.com/docs/api-reference/chat-streaming/streaming
|
||||
"""
|
||||
openai_stream = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": "Hello"},
|
||||
"index": 0,
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": " there"},
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_stream():
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
return_value=httpx.Response(200, stream=mock_stream())
|
||||
)
|
||||
|
||||
return route
|
||||
|
||||
|
||||
@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
|
||||
in its textual reply (the real test is that the image URL is
|
||||
forwarded in the request body to the AI service).
|
||||
"""
|
||||
openai_stream = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": "I see a cat"},
|
||||
"index": 0,
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-1234567890",
|
||||
"created": timezone.make_naive(timezone.now()).timestamp(),
|
||||
"choices": [
|
||||
{
|
||||
"delta": {"content": " in the picture."},
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"object": "chat.completion.chunk",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_stream():
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
return_value=httpx.Response(200, stream=mock_stream())
|
||||
)
|
||||
return route
|
||||
|
||||
|
||||
@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:
|
||||
1. First call returns function call
|
||||
2. Second call returns final answer after tool execution
|
||||
"""
|
||||
|
||||
# First response - tool call
|
||||
first_response = (
|
||||
"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": "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"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: "
|
||||
+ 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"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
# Second response - final answer
|
||||
second_response = (
|
||||
"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": "The current weather in Paris is nice"}, "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"}],
|
||||
}
|
||||
)
|
||||
+ "\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"}],
|
||||
}
|
||||
)
|
||||
+ "\n\n"
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
async def mock_first_response_stream():
|
||||
for line in first_response.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
async def mock_second_response_stream():
|
||||
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()),
|
||||
tool_answer_side_effect,
|
||||
]
|
||||
)
|
||||
|
||||
return route
|
||||
|
||||
|
||||
def test_post_conversation_anonymous(api_client):
|
||||
"""Test posting messages as an anonymous user returns a 401 error."""
|
||||
chat_conversation = ChatConversationFactory()
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user