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

exa_py/research/sync_client.py

@@ -0,0 +1,308 @@
+"""Synchronous Research API client."""
+
+from __future__ import annotations
+
+import time
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Generic,
+    Literal,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from pydantic import BaseModel, TypeAdapter
+
+from .base import ResearchBaseClient
+from .models import (
+    ResearchDto,
+    ResearchEvent,
+    ListResearchResponseDto,
+)
+from .utils import (
+    is_pydantic_model,
+    pydantic_to_json_schema,
+    stream_sse_events,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ResearchTyped(Generic[T]):
+    """Wrapper for typed research responses."""
+
+    def __init__(self, research: ResearchDto, parsed_output: T):
+        self.research = research
+        self.parsed_output = parsed_output
+        # Expose research fields
+        self.research_id = research.research_id
+        self.status = research.status
+        self.created_at = research.created_at
+        self.model = research.model
+        self.instructions = research.instructions
+        if hasattr(research, "events"):
+            self.events = research.events
+        if hasattr(research, "output"):
+            self.output = research.output
+        if hasattr(research, "cost_dollars"):
+            self.cost_dollars = research.cost_dollars
+        if hasattr(research, "error"):
+            self.error = research.error
+
+
+class ResearchClient(ResearchBaseClient):
+    """Synchronous client for the Research API."""
+
+    @overload
+    def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+    ) -> ResearchDto: ...
+
+    @overload
+    def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+        output_schema: Dict[str, Any],
+    ) -> ResearchDto: ...
+
+    @overload
+    def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+        output_schema: Type[T],
+    ) -> ResearchDto: ...
+
+    def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+        output_schema: Optional[Union[Dict[str, Any], Type[BaseModel]]] = None,
+    ) -> ResearchDto:
+        """Create a new research request.
+
+        Args:
+            instructions: The research instructions.
+            model: The model to use for research.
+            output_schema: Optional JSON schema or Pydantic model for structured output.
+
+        Returns:
+            The created research object.
+        """
+        payload = {
+            "instructions": instructions,
+            "model": model,
+        }
+
+        if output_schema is not None:
+            if is_pydantic_model(output_schema):
+                payload["outputSchema"] = pydantic_to_json_schema(output_schema)
+            else:
+                payload["outputSchema"] = output_schema
+
+        response = self.request("", method="POST", data=payload)
+        adapter = TypeAdapter(ResearchDto)
+        return adapter.validate_python(response)
+
+    @overload
+    def get(
+        self,
+        research_id: str,
+    ) -> ResearchDto: ...
+
+    @overload
+    def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[False] = False,
+        events: bool = False,
+    ) -> ResearchDto: ...
+
+    @overload
+    def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[True],
+        events: Optional[bool] = None,
+    ) -> Generator[ResearchEvent, None, None]: ...
+
+    @overload
+    def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[False] = False,
+        events: bool = False,
+        output_schema: Type[T],
+    ) -> ResearchTyped[T]: ...
+
+    def get(
+        self,
+        research_id: str,
+        *,
+        stream: bool = False,
+        events: bool = False,
+        output_schema: Optional[Type[BaseModel]] = None,
+    ) -> Union[ResearchDto, ResearchTyped, Generator[ResearchEvent, None, None]]:
+        """Get a research request by ID.
+
+        Args:
+            research_id: The research ID.
+            stream: Whether to stream events.
+            events: Whether to include events in non-streaming response.
+            output_schema: Optional Pydantic model for typed output validation.
+
+        Returns:
+            Research object, typed research, or event generator.
+        """
+        params = {}
+        if not stream:
+            params["stream"] = "false"
+            if events:
+                params["events"] = "true"
+        else:
+            params["stream"] = "true"
+            if events is not None:
+                params["events"] = str(events).lower()
+
+        if stream:
+            response = self.request(
+                f"/{research_id}", method="GET", params=params, stream=True
+            )
+            return stream_sse_events(response)
+        else:
+            response = self.request(f"/{research_id}", method="GET", params=params)
+            adapter = TypeAdapter(ResearchDto)
+            research = adapter.validate_python(response)
+
+            if output_schema and hasattr(research, "output") and research.output:
+                try:
+                    if research.output.parsed:
+                        parsed = output_schema.model_validate(research.output.parsed)
+                    else:
+                        import json
+
+                        parsed_data = json.loads(research.output.content)
+                        parsed = output_schema.model_validate(parsed_data)
+                    return ResearchTyped(research, parsed)
+                except Exception:
+                    # If parsing fails, return the regular research object
+                    return research
+
+            return research
+
+    def list(
+        self,
+        *,
+        cursor: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> ListResearchResponseDto:
+        """List research requests.
+
+        Args:
+            cursor: Pagination cursor.
+            limit: Maximum number of results.
+
+        Returns:
+            List of research objects with pagination info.
+        """
+        params = self.build_pagination_params(cursor, limit)
+        response = self.request("", method="GET", params=params)
+        return ListResearchResponseDto.model_validate(response)
+
+    @overload
+    def poll_until_finished(
+        self,
+        research_id: str,
+        *,
+        poll_interval: int = 1000,
+        timeout_ms: int = 600000,
+        events: bool = False,
+    ) -> ResearchDto: ...
+
+    @overload
+    def poll_until_finished(
+        self,
+        research_id: str,
+        *,
+        poll_interval: int = 1000,
+        timeout_ms: int = 600000,
+        events: bool = False,
+        output_schema: Type[T],
+    ) -> ResearchTyped[T]: ...
+
+    def poll_until_finished(
+        self,
+        research_id: str,
+        *,
+        poll_interval: int = 1000,
+        timeout_ms: int = 600000,
+        events: bool = False,
+        output_schema: Optional[Type[BaseModel]] = None,
+    ) -> Union[ResearchDto, ResearchTyped]:
+        """Poll until research is finished.
+
+        Args:
+            research_id: The research ID.
+            poll_interval: Milliseconds between polls (default 1000).
+            timeout_ms: Maximum time to wait in milliseconds (default 600000).
+            events: Whether to include events in the response.
+            output_schema: Optional Pydantic model for typed output validation.
+
+        Returns:
+            Completed research object or typed research.
+
+        Raises:
+            TimeoutError: If research doesn't complete within timeout.
+            RuntimeError: If polling fails too many times.
+        """
+        poll_interval_sec = poll_interval / 1000
+        timeout_sec = timeout_ms / 1000
+        max_consecutive_failures = 5
+        start_time = time.time()
+        consecutive_failures = 0
+
+        while True:
+            try:
+                if output_schema:
+                    result = self.get(
+                        research_id, events=events, output_schema=output_schema
+                    )
+                else:
+                    result = self.get(research_id, events=events)
+
+                consecutive_failures = 0
+
+                # Check if research is finished
+                status = result.status if hasattr(result, "status") else None
+                if status in ["completed", "failed", "canceled"]:
+                    return result
+
+            except Exception as e:
+                consecutive_failures += 1
+                if consecutive_failures >= max_consecutive_failures:
+                    raise RuntimeError(
+                        f"Polling failed {max_consecutive_failures} times in a row "
+                        f"for research {research_id}: {e}"
+                    )
+
+            if time.time() - start_time > timeout_sec:
+                raise TimeoutError(
+                    f"Research {research_id} did not complete within {timeout_ms}ms"
+                )
+
+            time.sleep(poll_interval_sec)

exa_py/research/utils.py

@@ -0,0 +1,222 @@
+"""Utilities for the Research API."""
+
+from __future__ import annotations
+
+import json
+from typing import (
+    Any,
+    AsyncGenerator,
+    Dict,
+    Generator,
+    Optional,
+    Type,
+)
+
+import httpx
+import requests
+from pydantic import BaseModel, ValidationError
+
+from .models import (
+    ResearchEvent,
+    ResearchDefinitionEvent,
+    ResearchOutputEvent,
+    ResearchPlanDefinitionEvent,
+    ResearchPlanOperationEvent,
+    ResearchPlanOutputEvent,
+    ResearchTaskDefinitionEvent,
+    ResearchTaskOperationEvent,
+    ResearchTaskOutputEvent,
+)
+
+
+def is_pydantic_model(schema: Any) -> bool:
+    """Check if the given schema is a Pydantic model class.
+
+    Args:
+        schema: The schema to check.
+
+    Returns:
+        True if schema is a Pydantic model class, False otherwise.
+    """
+    try:
+        return isinstance(schema, type) and issubclass(schema, BaseModel)
+    except (TypeError, AttributeError):
+        return False
+
+
+def pydantic_to_json_schema(model: Type[BaseModel]) -> Dict[str, Any]:
+    """Convert a Pydantic model to JSON Schema.
+
+    Args:
+        model: The Pydantic model class.
+
+    Returns:
+        JSON Schema dictionary with all references inlined.
+    """
+    # Import here to avoid circular imports
+    from exa_py.utils import _convert_schema_input
+
+    # Use the existing _convert_schema_input which already handles inlining
+    return _convert_schema_input(model)
+
+
+def parse_sse_line(line: str) -> Optional[tuple[str, str]]:
+    """Parse a single SSE line.
+
+    Args:
+        line: The SSE line to parse.
+
+    Returns:
+        Tuple of (field, value) or None if not a valid SSE line.
+    """
+    if not line or not line.strip():
+        return None
+
+    if ":" not in line:
+        return None
+
+    field, _, value = line.partition(":")
+    return field.strip(), value.strip()
+
+
+def parse_sse_event_raw(event_lines: list[str]) -> Optional[Dict[str, Any]]:
+    """Parse SSE event lines into a raw event dictionary.
+
+    Args:
+        event_lines: List of lines that make up an SSE event.
+
+    Returns:
+        Parsed event data or None if invalid.
+    """
+    event_name = None
+    event_data = None
+
+    for line in event_lines:
+        parsed = parse_sse_line(line)
+        if not parsed:
+            continue
+
+        field, value = parsed
+        if field == "event":
+            event_name = value
+        elif field == "data":
+            try:
+                event_data = json.loads(value)
+            except json.JSONDecodeError:
+                # Some events might have non-JSON data
+                event_data = value
+
+    if event_name and event_data:
+        # Add eventType to the data for consistency
+        if isinstance(event_data, dict):
+            event_data["eventType"] = event_name
+        return event_data
+
+    return None
+
+
+def parse_research_event(raw_event: Dict[str, Any]) -> Optional[ResearchEvent]:
+    """Parse a raw event dictionary into a typed ResearchEvent.
+
+    Args:
+        raw_event: Raw event dictionary with eventType field.
+
+    Returns:
+        Typed ResearchEvent or None if parsing fails.
+    """
+    event_type = raw_event.get("eventType")
+    if not event_type:
+        return None
+
+    # Map event types to their corresponding Pydantic models
+    event_models = {
+        "research-definition": ResearchDefinitionEvent,
+        "research-output": ResearchOutputEvent,
+        "plan-definition": ResearchPlanDefinitionEvent,
+        "plan-operation": ResearchPlanOperationEvent,
+        "plan-output": ResearchPlanOutputEvent,
+        "task-definition": ResearchTaskDefinitionEvent,
+        "task-operation": ResearchTaskOperationEvent,
+        "task-output": ResearchTaskOutputEvent,
+    }
+
+    model_class = event_models.get(event_type)
+    if not model_class:
+        return None
+
+    try:
+        return model_class.model_validate(raw_event)
+    except ValidationError:
+        # Log or handle validation error if needed
+        return None
+
+
+def stream_sse_events(
+    response: requests.Response,
+) -> Generator[ResearchEvent, None, None]:
+    """Stream SSE events from a requests Response.
+
+    Args:
+        response: The streaming response object.
+
+    Yields:
+        Parsed ResearchEvent objects.
+    """
+    event_lines = []
+
+    for line in response.iter_lines():
+        if not line:
+            # Empty line signals end of event
+            if event_lines:
+                raw_event = parse_sse_event_raw(event_lines)
+                if raw_event:
+                    event = parse_research_event(raw_event)
+                    if event:
+                        yield event
+                event_lines = []
+        else:
+            decoded_line = line.decode("utf-8") if isinstance(line, bytes) else line
+            event_lines.append(decoded_line)
+
+    # Handle any remaining lines
+    if event_lines:
+        raw_event = parse_sse_event_raw(event_lines)
+        if raw_event:
+            event = parse_research_event(raw_event)
+            if event:
+                yield event
+
+
+async def async_stream_sse_events(
+    response: httpx.Response,
+) -> AsyncGenerator[ResearchEvent, None]:
+    """Stream SSE events from an httpx Response.
+
+    Args:
+        response: The async streaming response object.
+
+    Yields:
+        Parsed ResearchEvent objects.
+    """
+    event_lines = []
+
+    async for line in response.aiter_lines():
+        if not line:
+            # Empty line signals end of event
+            if event_lines:
+                raw_event = parse_sse_event_raw(event_lines)
+                if raw_event:
+                    event = parse_research_event(raw_event)
+                    if event:
+                        yield event
+                event_lines = []
+        else:
+            event_lines.append(line)
+
+    # Handle any remaining lines
+    if event_lines:
+        raw_event = parse_sse_event_raw(event_lines)
+        if raw_event:
+            event = parse_research_event(raw_event)
+            if event:
+                yield event

exa_py/utils.py

@@ -1,10 +1,7 @@
 import json
-import os
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 from openai.types.chat import ChatCompletion
 
-from typing import TYPE_CHECKING
-
 from pydantic import BaseModel
 from pydantic.json_schema import GenerateJsonSchema
 

{exa_py-1.14.20.dist-info → exa_py-1.15.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: exa-py
-Version: 1.14.20
+Version: 1.15.0
 Summary: Python SDK for Exa API.
 License: MIT
 Author: Exa AI

{exa_py-1.14.20.dist-info → exa_py-1.15.0.dist-info}/RECORD

@@ -1,10 +1,13 @@
 exa_py/__init__.py,sha256=M2GC9oSdoV6m2msboW0vMWWl8wrth4o6gmEV4MYLGG8,66
-exa_py/api.py,sha256=TIGTu4ITw7SrhTZjx9uQNpwi1u3nDpjtYUNXQrgg5bo,107356
+exa_py/api.py,sha256=gUvQ2NXtUgBlrFrWrCIUwiq-z2lAlo_gWVFncN68Vhw,108947
 exa_py/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-exa_py/research/__init__.py,sha256=QeY-j6bP4QP5tF9ytX0IeQhJvd0Wn4cJCD69U8pP7kA,271
-exa_py/research/client.py,sha256=mnoTA4Qoa0TA5d8nVTR9tAU9LJElXV-MlPozgMxlUp4,12799
-exa_py/research/models.py,sha256=j7YgRoMRp2MLgnaij7775x_hJEeV5gksKpfLwmawqxY,3704
-exa_py/utils.py,sha256=eYnJRAFJonwKP_mCxzAB9TnLEqoF-88stg6wh-M-Ups,6424
+exa_py/research/__init__.py,sha256=qbWY5c3LYBM7L4yGVtiMs9nUg_kMcAVhMhT0DvN6EI4,1014
+exa_py/research/async_client.py,sha256=jP6BF8nc29Mt4iU2ZEQGAC5dGM-Z9I4k3pzpPcCKJIs,9436
+exa_py/research/base.py,sha256=0C52XGUMRwRox1JSVBrbiHZxqPI6Cgif7EFNqYxpc40,5451
+exa_py/research/models.py,sha256=2gu4jckbLYHlTmln8BhDzovuqg1fejTCncqi4r0oDow,9822
+exa_py/research/sync_client.py,sha256=LHMkgqUZgKCQWgLP41Fj6EWLbt_YfX4aAgvWGjNgbXM,9155
+exa_py/research/utils.py,sha256=YR61UkReyK-LrvOTaX_aG_neS1DsmCkv7B6YODpwHvU,5965
+exa_py/utils.py,sha256=1jwKwcJwHSbrzwgMkprXSEZ-xssCtXxwf-GTRgaFmz8,6395
 exa_py/websets/__init__.py,sha256=x7Dc0MS8raRXA7Ud6alKgnsUmLi6X9GTqfB8kOwC9iQ,179
 exa_py/websets/_generator/pydantic/BaseModel.jinja2,sha256=RUDCmPZVamoVx1WudylscYFfDhGoNNtRYlpTvKjAiuA,1276
 exa_py/websets/client.py,sha256=sKkji8QaPFnGM1-J5TB6yKJcGAEd6gk7lsnIebzXNQ8,5856
@@ -27,6 +30,6 @@ exa_py/websets/searches/client.py,sha256=X3f7axWGfecmxf-2tBTX0Yf_--xToz1X8ZHbbud
 exa_py/websets/types.py,sha256=DxO_T4Ijnd06gxFAX3f238Mt5P0_ulpY44M1kiT4y4U,47120
 exa_py/websets/webhooks/__init__.py,sha256=iTPBCxFd73z4RifLQMX6iRECx_6pwlI5qscLNjMOUHE,77
 exa_py/websets/webhooks/client.py,sha256=zS1eoWKliuiY4AIeFJdpAlPZeOINyphn7KEWANF-zaE,4384
-exa_py-1.14.20.dist-info/METADATA,sha256=tkmcVljW6I0gOgd6uLMTkMXLnmJLQOSVJmeoxLW3tQw,3827
-exa_py-1.14.20.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-exa_py-1.14.20.dist-info/RECORD,,
+exa_py-1.15.0.dist-info/METADATA,sha256=QTMlL2y63YGWPWdSrQy-BkoDSAZ-axbTMVEtPhE_H0c,3826
+exa_py-1.15.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+exa_py-1.15.0.dist-info/RECORD,,