from decimal import Decimal from fields.message_fields import ExploreMessageListItem, MessageListItem, RetrieverResource, WebMessageListItem def _retrieved_video_resource() -> dict[str, object]: # Shaped like a real stored retriever_resources entry (core/rag/entities/citation_metadata.py's # RetrievalSourceMetadata, the write-side model) for a video citation with a start/end offset. return { "position": 1, "dataset_id": "af57b949-2259-4fe4-b58a-0117fd3a6b92", "document_id": "9196b0c1-898e-4e8e-9d78-18cb07896b89", "document_name": "training.mp4", "retriever_from": "workflow", "score": 0.86, "content": "transcript excerpt", "page": None, "title": "training.mp4", "files": None, "doc_metadata": {"start_seconds": 1036.0, "end_seconds": 1054.6, "document_name": "training.mp4"}, } def _base_kwargs() -> dict[str, object]: inputs: dict[str, object] = {} retriever_resources: list[object] = [] agent_thoughts: list[object] = [] message_files: list[object] = [] extra_contents: list[object] = [] return { "id": "m1", "conversation_id": "c1", "inputs": inputs, "query": "hi", "answer": "answer", "retriever_resources": retriever_resources, "agent_thoughts": agent_thoughts, "message_files": message_files, "status": "normal", "extra_contents": extra_contents, } class TestExploreMessageListItem: def test_retriever_resource_preserves_knowledge_fs_citation_identity(self) -> None: kwargs = _base_kwargs() payload = MessageListItem.model_validate( { **kwargs, "retriever_resources": [ { "position": 1, "document_id": "11111111-1111-4111-8111-111111111111", "document_asset_id": "22222222-2222-4222-8222-222222222222", "document_revision": 4, "document_version": 7, } ], } ).model_dump(mode="json")["retriever_resources"][0] assert payload["document_id"] == "11111111-1111-4111-8111-111111111111" assert payload["document_asset_id"] == "22222222-2222-4222-8222-222222222222" assert payload["document_revision"] == 4 assert payload["document_version"] == 7 def test_exposes_metadata_for_history_rehydration(self) -> None: # The Explore/installed-app surface must surface message_metadata (incl. reasoning) # so the chat-with-history client can rehydrate the thinking panel on reload. item = ExploreMessageListItem.model_validate( {**_base_kwargs(), "metadata": {"reasoning": {"llm": "thinking..."}}} ) payload = item.model_dump(mode="json") assert payload["metadata"] == {"reasoning": {"llm": "thinking..."}} def test_metadata_defaults_to_none(self) -> None: item = ExploreMessageListItem.model_validate(_base_kwargs()) assert item.model_dump(mode="json")["metadata"] is None def test_base_message_list_item_has_no_metadata(self) -> None: # Guard the public service-API contract: the base item must not leak metadata. payload = MessageListItem.model_validate(_base_kwargs()).model_dump(mode="json") assert "metadata" not in payload def test_message_list_item_exposes_usage_fields(self) -> None: payload = MessageListItem.model_validate( { **_base_kwargs(), "message_tokens": 7, "answer_tokens": 11, "provider_response_latency": 1.25, "total_price": Decimal("0.0001234"), "currency": "USD", } ).model_dump(mode="json") assert payload["message_tokens"] == 7 assert payload["answer_tokens"] == 11 assert payload["total_tokens"] == 18 assert payload["provider_response_latency"] == 1.25 assert payload["total_price"] == "0.0001234" assert payload["currency"] == "USD" def test_web_message_list_item_exposes_usage_and_metadata(self) -> None: payload = WebMessageListItem.model_validate( { **_base_kwargs(), "metadata": {"usage": {"total_tokens": 18}}, "message_tokens": 7, "answer_tokens": 11, } ).model_dump(mode="json") assert payload["metadata"] == {"usage": {"total_tokens": 18}} assert payload["message_tokens"] == 7 assert payload["answer_tokens"] == 11 assert payload["total_tokens"] == 18 class TestRetrieverResourceDocMetadata: def test_doc_metadata_survives_validation_from_a_stored_dict(self) -> None: # GET /messages re-validates the raw dicts stored in Message.retriever_resources # (JSON-backed, written from RetrievalSourceMetadata) against RetrieverResource. # Before this field existed, ResponseModel's extra="ignore" silently dropped # doc_metadata here -- the video player on a reloaded conversation had nothing # to read start_seconds/end_seconds from, even though the data was never lost # in storage. resource = RetrieverResource.model_validate(_retrieved_video_resource()) payload = resource.model_dump(mode="json") assert payload["doc_metadata"] == { "start_seconds": 1036.0, "end_seconds": 1054.6, "document_name": "training.mp4", } assert payload["title"] == "training.mp4" assert payload["retriever_from"] == "workflow" def test_message_list_item_propagates_doc_metadata_through_retriever_resources(self) -> None: payload = MessageListItem.model_validate( {**_base_kwargs(), "retriever_resources": [_retrieved_video_resource()]} ).model_dump(mode="json") assert payload["retriever_resources"][0]["doc_metadata"]["start_seconds"] == 1036.0