exa-py 1.14.20__py3-none-any.whl → 1.15.0__py3-none-any.whl

exa_py/research/models.py

@@ -1,120 +1,321 @@
 from __future__ import annotations
 
-import json
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
-
-# Local import placed inside TYPE_CHECKING block to avoid runtime cycles.
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:  # pragma: no cover – for static analysers only
-    from ..api import _Result  # noqa: F401
-
-
-@dataclass
-class ResearchTaskId:
-    """Structured research task ID.
-
-    Attributes
-    ----------
-    id:
-        Unique identifier for the research task.
-    """
-
-    id: str
-
-    # ---------------------------------------------------------------------
-    # Pretty representation helpers
-    # ---------------------------------------------------------------------
-    def __str__(self) -> str:  # pragma: no cover – convenience only
-        return f"ID: {self.id}\n"
-
-
-@dataclass
-class ResearchTask:
-    """Structured research task.
-
-    Attributes
-    ----------
-    id:
-        Unique identifier for the research task.
-    status:
-        Current task status
-    instructions:
-        Instructions for the task
-    schema:
-        Output schema defining the task
-    data:
-        JSON-serialisable answer generated by Exa (may be ``None`` until the task
-        completes).
-    citations:
-        Mapping from *root field* in the output schema to the list of search
-        results that were used to generate that part of the answer.
-    """
-
-    id: str
-    status: str
+from typing import Annotated, Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class CostDollars(BaseModel):
+    total: float
+    num_pages: Annotated[float, Field(alias="numPages")]
+    num_searches: Annotated[float, Field(alias="numSearches")]
+    reasoning_tokens: Annotated[float, Field(alias="reasoningTokens")]
+
+
+class Result(BaseModel):
+    url: str
+
+
+class ResearchThinkOperation(BaseModel):
+    type: Literal["think"]
+    content: str
+
+
+class ResearchSearchOperation(BaseModel):
+    type: Literal["search"]
+    search_type: Annotated[
+        Literal["neural", "keyword", "auto", "fast"], Field(alias="searchType")
+    ]
+    query: str
+    results: List[Result]
+    page_tokens: Annotated[float, Field(alias="pageTokens")]
+    goal: Optional[str] = None
+
+
+class ResearchCrawlOperation(BaseModel):
+    type: Literal["crawl"]
+    result: Result
+    page_tokens: Annotated[float, Field(alias="pageTokens")]
+    goal: Optional[str] = None
+
+
+ResearchOperation = Annotated[
+    Union[ResearchThinkOperation, ResearchSearchOperation, ResearchCrawlOperation],
+    Field(discriminator="type"),
+]
+
+
+class ResearchDefinitionEvent(BaseModel):
+    event_type: Annotated[Literal["research-definition"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    instructions: str
+    output_schema: Annotated[Optional[Dict[str, Any]], Field(alias="outputSchema")] = (
+        None
+    )
+
+
+class ResearchOutputCompleted(BaseModel):
+    output_type: Annotated[Literal["completed"], Field(alias="outputType")]
+    content: str
+    cost_dollars: Annotated[CostDollars, Field(alias="costDollars")]
+    parsed: Optional[Dict[str, Any]] = None
+
+
+class ResearchOutputFailed(BaseModel):
+    output_type: Annotated[Literal["failed"], Field(alias="outputType")]
+    error: str
+
+
+class ResearchOutputEvent(BaseModel):
+    event_type: Annotated[Literal["research-output"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    output: Annotated[
+        Union[ResearchOutputCompleted, ResearchOutputFailed],
+        Field(discriminator="output_type"),
+    ]
+
+
+class ResearchPlanDefinitionEvent(BaseModel):
+    event_type: Annotated[Literal["plan-definition"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+
+
+class ResearchPlanOperationEvent(BaseModel):
+    event_type: Annotated[Literal["plan-operation"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    operation_id: Annotated[str, Field(alias="operationId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    data: ResearchOperation
+
+
+class ResearchPlanOutputTasks(BaseModel):
+    output_type: Annotated[Literal["tasks"], Field(alias="outputType")]
+    reasoning: str
+    tasks_instructions: Annotated[List[str], Field(alias="tasksInstructions")]
+
+
+class ResearchPlanOutputStop(BaseModel):
+    output_type: Annotated[Literal["stop"], Field(alias="outputType")]
+    reasoning: str
+
+
+class ResearchPlanOutputEvent(BaseModel):
+    event_type: Annotated[Literal["plan-output"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    output: Annotated[
+        Union[ResearchPlanOutputTasks, ResearchPlanOutputStop],
+        Field(discriminator="output_type"),
+    ]
+
+
+class ResearchTaskDefinitionEvent(BaseModel):
+    event_type: Annotated[Literal["task-definition"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    task_id: Annotated[str, Field(alias="taskId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
     instructions: str
-    schema: Dict[str, Any]
-    data: Optional[Dict[str, Any]]
-    citations: Dict[str, List["_Result"]]
-
-    # ---------------------------------------------------------------------
-    # Pretty representation helpers
-    # ---------------------------------------------------------------------
-    def __str__(self) -> str:  # pragma: no cover – convenience only
-        """Human-readable representation including *all* relevant fields."""
-        schema_repr = json.dumps(self.schema, indent=2, ensure_ascii=False)
-        data_repr = (
-            json.dumps(self.data, indent=2, ensure_ascii=False)
-            if self.data is not None
-            else "None"
-        )
-
-        # Render citations grouped by the root field they belong to.
-        if self.citations:
-            # Each key is a root field, each value is a list of _Result objects.
-            citations_lines = []
-            for field, sources in self.citations.items():
-                rendered_sources = "\n    ".join(str(src) for src in sources)
-                citations_lines.append(f"{field}:\n    {rendered_sources}")
-            citations_str = "\n\n".join(citations_lines)
-        else:
-            citations_str = "None"
-
-        return (
-            f"ID: {self.id}\n"
-            f"Status: {self.status}\n"
-            f"Instructions: {self.instructions}\n"
-            f"Schema:\n{schema_repr}\n"
-            f"Data:\n{data_repr}\n\n"
-            f"Citations:\n{citations_str}"
-        )
-
-
-@dataclass
-class ListResearchTasksResponse:
-    """Paginated list of research tasks."""
-
-    data: List[ResearchTask]
-    has_more: bool
-    next_cursor: Optional[str]
-
-    # -----------------------------------------------------------------
-    # Pretty representation helpers
-    # -----------------------------------------------------------------
-    def __str__(self) -> str:  # pragma: no cover – convenience only
-        tasks_repr = "\n\n".join(str(task) for task in self.data)
-        cursor_repr = self.next_cursor or "None"
-        return (
-            f"Tasks:\n{tasks_repr}\n\n"
-            f"Has more: {self.has_more}\n"
-            f"Next cursor: {cursor_repr}"
-        )
+
+
+class ResearchTaskOperationEvent(BaseModel):
+    event_type: Annotated[Literal["task-operation"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    task_id: Annotated[str, Field(alias="taskId")]
+    operation_id: Annotated[str, Field(alias="operationId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    data: ResearchOperation
+
+
+class ResearchTaskOutput(BaseModel):
+    output_type: Annotated[Literal["completed"], Field(alias="outputType")]
+    content: str
+
+
+class ResearchTaskOutputEvent(BaseModel):
+    event_type: Annotated[Literal["task-output"], Field(alias="eventType")]
+    research_id: Annotated[str, Field(alias="researchId")]
+    plan_id: Annotated[str, Field(alias="planId")]
+    task_id: Annotated[str, Field(alias="taskId")]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    output: ResearchTaskOutput
+
+
+ResearchMetaEvent = Union[ResearchDefinitionEvent, ResearchOutputEvent]
+ResearchPlanEvent = Union[
+    ResearchPlanDefinitionEvent, ResearchPlanOperationEvent, ResearchPlanOutputEvent
+]
+ResearchTaskEvent = Union[
+    ResearchTaskDefinitionEvent, ResearchTaskOperationEvent, ResearchTaskOutputEvent
+]
+ResearchEvent = Union[ResearchMetaEvent, ResearchPlanEvent, ResearchTaskEvent]
+
+
+class ResearchOutput(BaseModel):
+    content: str
+    parsed: Optional[Dict[str, Any]] = None
+
+
+class ResearchBaseDto(BaseModel):
+    research_id: Annotated[
+        str,
+        Field(
+            alias="researchId",
+            description="The unique identifier for the research request",
+        ),
+    ]
+    created_at: Annotated[
+        float, Field(alias="createdAt", description="Milliseconds since epoch time")
+    ]
+    model: Annotated[
+        Literal["exa-research", "exa-research-pro"],
+        Field(description="The model used for the research request"),
+    ] = "exa-research"
+    instructions: Annotated[
+        str, Field(description="The instructions given to this research request")
+    ]
+
+
+class ResearchPendingDto(ResearchBaseDto):
+    status: Literal["pending"]
+
+
+class ResearchRunningDto(ResearchBaseDto):
+    status: Literal["running"]
+    events: Optional[List[ResearchEvent]] = None
+
+
+class ResearchCompletedDto(ResearchBaseDto):
+    status: Literal["completed"]
+    events: Optional[List[ResearchEvent]] = None
+    output: ResearchOutput
+    cost_dollars: Annotated[CostDollars, Field(alias="costDollars")]
+
+
+class ResearchCanceledDto(ResearchBaseDto):
+    status: Literal["canceled"]
+    events: Optional[List[ResearchEvent]] = None
+
+
+class ResearchFailedDto(ResearchBaseDto):
+    status: Literal["failed"]
+    events: Optional[List[ResearchEvent]] = None
+    error: Annotated[
+        str, Field(description="A message indicating why the request failed")
+    ]
+
+
+ResearchDto = Annotated[
+    Union[
+        ResearchPendingDto,
+        ResearchRunningDto,
+        ResearchCompletedDto,
+        ResearchCanceledDto,
+        ResearchFailedDto,
+    ],
+    Field(discriminator="status"),
+]
+
+
+class ListResearchResponseDto(BaseModel):
+    data: Annotated[
+        List[ResearchDto], Field(description="The list of research requests")
+    ]
+    has_more: Annotated[
+        bool,
+        Field(
+            alias="hasMore",
+            description="Whether there are more results to paginate through",
+        ),
+    ]
+    next_cursor: Annotated[
+        Optional[str],
+        Field(
+            alias="nextCursor",
+            description="The cursor to paginate through the next set of results",
+        ),
+    ]
+
+
+class ResearchCreateRequestDto(BaseModel):
+    model: Literal["exa-research", "exa-research-pro"] = "exa-research"
+    instructions: Annotated[
+        str,
+        Field(
+            description="Instructions for what research should be conducted",
+            max_length=4096,
+        ),
+    ]
+    output_schema: Annotated[Optional[Dict[str, Any]], Field(alias="outputSchema")] = (
+        None
+    )
+
+
+ResearchDtoClass = RootModel[ResearchDto]
+ResearchCreateRequestDtoClass = ResearchCreateRequestDto
+ResearchEventDtoClass = RootModel[ResearchEvent]
+ResearchOperationDtoClass = RootModel[ResearchOperation]
 
 
 __all__ = [
-    "ResearchTaskId",
-    "ResearchTask",
-    "ListResearchTasksResponse",
+    "CostDollars",
+    "Result",
+    "ResearchThinkOperation",
+    "ResearchSearchOperation",
+    "ResearchCrawlOperation",
+    "ResearchOperation",
+    "ResearchDefinitionEvent",
+    "ResearchOutputCompleted",
+    "ResearchOutputFailed",
+    "ResearchOutputEvent",
+    "ResearchPlanDefinitionEvent",
+    "ResearchPlanOperationEvent",
+    "ResearchPlanOutputTasks",
+    "ResearchPlanOutputStop",
+    "ResearchPlanOutputEvent",
+    "ResearchTaskDefinitionEvent",
+    "ResearchTaskOperationEvent",
+    "ResearchTaskOutput",
+    "ResearchTaskOutputEvent",
+    "ResearchMetaEvent",
+    "ResearchPlanEvent",
+    "ResearchTaskEvent",
+    "ResearchEvent",
+    "ResearchOutput",
+    "ResearchBaseDto",
+    "ResearchPendingDto",
+    "ResearchRunningDto",
+    "ResearchCompletedDto",
+    "ResearchCanceledDto",
+    "ResearchFailedDto",
+    "ResearchDto",
+    "ListResearchResponseDto",
+    "ResearchCreateRequestDto",
+    "ResearchDtoClass",
+    "ResearchCreateRequestDtoClass",
+    "ResearchEventDtoClass",
+    "ResearchOperationDtoClass",
 ]