exa-py 1.14.20__py3-none-any.whl → 1.15.1__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/websets/enrichments/client.py

@@ -4,6 +4,7 @@ from typing import Dict, Any, Union
 
 from ..types import (
     CreateEnrichmentParameters,
+    UpdateEnrichmentParameters,
     WebsetEnrichment,
 )
 from ..core.base import WebsetsBaseClient
@@ -40,6 +41,20 @@ class WebsetEnrichmentsClient(WebsetsBaseClient):
         response = self.request(f"/v0/websets/{webset_id}/enrichments/{id}", method="GET")
         return WebsetEnrichment.model_validate(response)
 
+    def update(self, webset_id: str, id: str, params: Union[Dict[str, Any], UpdateEnrichmentParameters]) -> WebsetEnrichment:
+        """Update an Enrichment.
+        
+        Args:
+            webset_id (str): The id of the Webset.
+            id (str): The id of the Enrichment.
+            params (UpdateEnrichmentParameters): The parameters for updating an enrichment.
+        
+        Returns:
+            WebsetEnrichment: The updated enrichment.
+        """
+        response = self.request(f"/v0/websets/{webset_id}/enrichments/{id}", data=params, method="PATCH")
+        return WebsetEnrichment.model_validate(response)
+
     def delete(self, webset_id: str, id: str) -> WebsetEnrichment:
         """Delete an Enrichment.
         

exa_py/websets/events/client.py

@@ -15,6 +15,13 @@ from ..types import (
     WebsetSearchUpdatedEvent,
     WebsetSearchCanceledEvent,
     WebsetSearchCompletedEvent,
+    ImportCreatedEvent,
+    ImportCompletedEvent,
+    MonitorCreatedEvent,
+    MonitorUpdatedEvent,
+    MonitorDeletedEvent,
+    MonitorRunCreatedEvent,
+    MonitorRunCompletedEvent,
 )
 from ..core.base import WebsetsBaseClient
 
@@ -30,6 +37,13 @@ Event = Union[
     WebsetSearchUpdatedEvent,
     WebsetSearchCanceledEvent,
     WebsetSearchCompletedEvent,
+    ImportCreatedEvent,
+    ImportCompletedEvent,
+    MonitorCreatedEvent,
+    MonitorUpdatedEvent,
+    MonitorDeletedEvent,
+    MonitorRunCreatedEvent,
+    MonitorRunCompletedEvent,
 ]
 
 class EventsClient(WebsetsBaseClient):
@@ -89,6 +103,13 @@ class EventsClient(WebsetsBaseClient):
             'webset.search.updated': WebsetSearchUpdatedEvent,
             'webset.search.canceled': WebsetSearchCanceledEvent,
             'webset.search.completed': WebsetSearchCompletedEvent,
+            'import.created': ImportCreatedEvent,
+            'import.completed': ImportCompletedEvent,
+            'monitor.created': MonitorCreatedEvent,
+            'monitor.updated': MonitorUpdatedEvent,
+            'monitor.deleted': MonitorDeletedEvent,
+            'monitor.run.created': MonitorRunCreatedEvent,
+            'monitor.run.completed': MonitorRunCompletedEvent,
         }
         
         event_class = event_type_map.get(event_type)