zl-q
69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
|
|
from v1.agent.utils import convert_message_to_history
|
|
|
|
|
|
class _FakeMessage:
|
|
def __init__(self, *, role: str, metadata: dict[str, object] | None) -> None:
|
|
self.id = uuid4()
|
|
self.seq = 1
|
|
self.role = role
|
|
self.content = "content"
|
|
self.metadata = metadata
|
|
self.timestamp = datetime.now(timezone.utc)
|
|
|
|
|
|
def test_convert_message_to_history_does_not_attach_ui_schema_for_tool_message() -> (
|
|
None
|
|
):
|
|
message = _FakeMessage(
|
|
role="tool",
|
|
metadata={"tool_agent_output": {"result": "done"}},
|
|
)
|
|
|
|
result = convert_message_to_history(message) # type: ignore[arg-type]
|
|
|
|
assert "ui_schema" not in result
|
|
assert "uiSchema" not in result
|
|
|
|
|
|
def test_convert_message_to_history_uses_ui_schema_key_for_assistant_message() -> None:
|
|
message = _FakeMessage(
|
|
role="assistant",
|
|
metadata={
|
|
"agent_output": {"ui_schema": {"version": "2.0", "root": {"type": "stack"}}}
|
|
},
|
|
)
|
|
|
|
result = convert_message_to_history(message) # type: ignore[arg-type]
|
|
|
|
assert "ui_schema" in result
|
|
assert "uiSchema" not in result
|
|
assert result["ui_schema"] == {"version": "2.0", "root": {"type": "stack"}}
|
|
|
|
|
|
def test_convert_message_to_history_returns_multiple_user_attachments() -> None:
|
|
message = _FakeMessage(
|
|
role="user",
|
|
metadata={
|
|
"user_message_attachments": [
|
|
{"bucket": "bucket-a", "path": "path/a.png", "mime_type": "image/png"},
|
|
{"bucket": "bucket-a", "path": "path/b.jpg", "mime_type": "image/jpeg"},
|
|
]
|
|
},
|
|
)
|
|
|
|
def _signed(payload: dict[str, str]) -> str:
|
|
return f"https://signed.example/{payload['bucket']}/{payload['path']}"
|
|
|
|
result = convert_message_to_history(message, _signed) # type: ignore[arg-type]
|
|
|
|
attachments = result.get("attachments")
|
|
assert isinstance(attachments, list)
|
|
assert len(attachments) == 2
|
|
assert attachments[0]["mimeType"] == "image/png"
|
|
assert attachments[1]["mimeType"] == "image/jpeg"
|