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

exa_py/api.py

@@ -41,7 +41,7 @@ from exa_py.utils import (
 )
 from .websets import WebsetsClient
 from .websets.core.base import ExaJSONEncoder
-from .research.client import ResearchClient, AsyncResearchClient
+from .research import ResearchClient, AsyncResearchClient
 
 
 is_beta = os.getenv("IS_BETA") == "True"
@@ -133,6 +133,7 @@ SEARCH_OPTIONS_TYPES = {
     "end_published_date": [
         str
     ],  # Results before this publish date; excludes links with no date. ISO 8601 format.
+    "user_location": [str],  # Two-letter ISO country code of the user (e.g. US).
     "include_text": [
         list
     ],  # Must be present in webpage text. (One string, up to 5 words)
@@ -1186,23 +1187,37 @@ class Exa:
             # Otherwise, serialize the dictionary to JSON if it exists
             json_data = json.dumps(data, cls=ExaJSONEncoder) if data else None
 
-        if data and data.get("stream"):
-            res = requests.post(
-                self.base_url + endpoint,
-                data=json_data,
-                headers=self.headers,
-                stream=True,
-            )
-            return res
+        # Check if we need streaming (either from data for POST or params for GET)
+        needs_streaming = (data and isinstance(data, dict) and data.get("stream")) or (
+            params and params.get("stream") == "true"
+        )
 
         if method.upper() == "GET":
-            res = requests.get(
-                self.base_url + endpoint, headers=self.headers, params=params
-            )
+            if needs_streaming:
+                res = requests.get(
+                    self.base_url + endpoint,
+                    headers=self.headers,
+                    params=params,
+                    stream=True,
+                )
+                return res
+            else:
+                res = requests.get(
+                    self.base_url + endpoint, headers=self.headers, params=params
+                )
         elif method.upper() == "POST":
-            res = requests.post(
-                self.base_url + endpoint, data=json_data, headers=self.headers
-            )
+            if needs_streaming:
+                res = requests.post(
+                    self.base_url + endpoint,
+                    data=json_data,
+                    headers=self.headers,
+                    stream=True,
+                )
+                return res
+            else:
+                res = requests.post(
+                    self.base_url + endpoint, data=json_data, headers=self.headers
+                )
         elif method.upper() == "PATCH":
             res = requests.patch(
                 self.base_url + endpoint, data=json_data, headers=self.headers
@@ -1236,6 +1251,7 @@ class Exa:
         category: Optional[str] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
     ) -> SearchResponse[_Result]:
         """Perform a search with a prompt-engineered query to retrieve relevant results.
 
@@ -1251,10 +1267,11 @@ class Exa:
             include_text (List[str], optional): Strings that must appear in the page text.
             exclude_text (List[str], optional): Strings that must not appear in the page text.
             use_autoprompt (bool, optional): Convert query to Exa (default False).
-            type (str, optional): 'keyword', 'neural', 'hybrid', or 'fast' (default 'neural').
+            type (str, optional): 'keyword', 'neural', 'hybrid', 'fast', or 'auto' (default 'auto').
             category (str, optional): e.g. 'company'
             flags (List[str], optional): Experimental flags for Exa usage.
             moderation (bool, optional): If True, the search results will be moderated for safety.
+            user_location (str, optional): Two-letter ISO country code of the user (e.g. US).
 
         Returns:
             SearchResponse: The response containing search results, etc.
@@ -1312,6 +1329,7 @@ class Exa:
         category: Optional[str] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1370,6 +1388,7 @@ class Exa:
         subpage_target: Optional[Union[str, List[str]]] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1399,6 +1418,7 @@ class Exa:
         subpage_target: Optional[Union[str, List[str]]] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1427,6 +1447,7 @@ class Exa:
         subpage_target: Optional[Union[str, List[str]]] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1456,6 +1477,7 @@ class Exa:
         subpage_target: Optional[Union[str, List[str]]] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1485,6 +1507,7 @@ class Exa:
         subpage_target: Optional[Union[str, List[str]]] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         filter_empty_results: Optional[bool] = None,
@@ -1513,6 +1536,7 @@ class Exa:
         category: Optional[str] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
         livecrawl_timeout: Optional[int] = None,
         livecrawl: Optional[LIVECRAWL_OPTIONS] = None,
         subpages: Optional[int] = None,
@@ -2396,34 +2420,55 @@ class AsyncExa(Exa):
         # this may only be a
         if self._client is None:
             self._client = httpx.AsyncClient(
-                base_url=self.base_url, headers=self.headers, timeout=60
+                base_url=self.base_url, headers=self.headers, timeout=600
             )
         return self._client
 
-    async def async_request(self, endpoint: str, data):
-        """Send a POST request to the Exa API, optionally streaming if data['stream'] is True.
+    async def async_request(
+        self, endpoint: str, data=None, method: str = "POST", params=None
+    ):
+        """Send a request to the Exa API, optionally streaming if data['stream'] is True.
 
         Args:
             endpoint (str): The API endpoint (path).
-            data (dict): The JSON payload to send.
+            data (dict, optional): The JSON payload to send.
+            method (str, optional): The HTTP method to use. Defaults to "POST".
+            params (dict, optional): Query parameters.
 
         Returns:
-            Union[dict, requests.Response]: If streaming, returns the Response object.
+            Union[dict, httpx.Response]: If streaming, returns the Response object.
             Otherwise, returns the JSON-decoded response as a dict.
 
         Raises:
             ValueError: If the request fails (non-200 status code).
         """
-        if data.get("stream"):
-            request = httpx.Request(
-                "POST", self.base_url + endpoint, json=data, headers=self.headers
-            )
-            res = await self.client.send(request, stream=True)
-            return res
-
-        res = await self.client.post(
-            self.base_url + endpoint, json=data, headers=self.headers
+        # Check if we need streaming (either from data for POST or params for GET)
+        needs_streaming = (data and isinstance(data, dict) and data.get("stream")) or (
+            params and params.get("stream") == "true"
         )
+
+        if method.upper() == "GET":
+            if needs_streaming:
+                request = httpx.Request(
+                    "GET", self.base_url + endpoint, params=params, headers=self.headers
+                )
+                res = await self.client.send(request, stream=True)
+                return res
+            else:
+                res = await self.client.get(
+                    self.base_url + endpoint, params=params, headers=self.headers
+                )
+        elif method.upper() == "POST":
+            if needs_streaming:
+                request = httpx.Request(
+                    "POST", self.base_url + endpoint, json=data, headers=self.headers
+                )
+                res = await self.client.send(request, stream=True)
+                return res
+            else:
+                res = await self.client.post(
+                    self.base_url + endpoint, json=data, headers=self.headers
+                )
         if res.status_code != 200 and res.status_code != 201:
             raise ValueError(
                 f"Request failed with status code {res.status_code}: {res.text}"
@@ -2448,6 +2493,7 @@ class AsyncExa(Exa):
         category: Optional[str] = None,
         flags: Optional[List[str]] = None,
         moderation: Optional[bool] = None,
+        user_location: Optional[str] = None,
     ) -> SearchResponse[_Result]:
         """Perform a search with a prompt-engineered query to retrieve relevant results.
 
@@ -2463,10 +2509,11 @@ class AsyncExa(Exa):
             include_text (List[str], optional): Strings that must appear in the page text.
             exclude_text (List[str], optional): Strings that must not appear in the page text.
             use_autoprompt (bool, optional): Convert query to Exa (default False).
-            type (str, optional): 'keyword', 'neural', 'hybrid', or 'fast' (default 'neural').
+            type (str, optional): 'keyword', 'neural', 'hybrid', 'fast', or 'auto' (default 'auto').
             category (str, optional): e.g. 'company'
             flags (List[str], optional): Experimental flags for Exa usage.
             moderation (bool, optional): If True, the search results will be moderated for safety.
+            user_location (str, optional): Two-letter ISO country code of the user (e.g. US).
 
         Returns:
             SearchResponse: The response containing search results, etc.

exa_py/research/__init__.py

@@ -1,10 +1,39 @@
-from .client import ResearchClient, AsyncResearchClient
-from .models import ResearchTask, ListResearchTasksResponse, ResearchTaskId
+"""Research API client modules for Exa."""
+
+from .sync_client import ResearchClient, ResearchTyped
+from .async_client import AsyncResearchClient, AsyncResearchTyped
+from .models import (
+    ResearchDto,
+    ResearchEvent,
+    ResearchDefinitionEvent,
+    ResearchOutputEvent,
+    ResearchPlanDefinitionEvent,
+    ResearchPlanOperationEvent,
+    ResearchPlanOutputEvent,
+    ResearchTaskDefinitionEvent,
+    ResearchTaskOperationEvent,
+    ResearchTaskOutputEvent,
+    ListResearchResponseDto,
+    CostDollars,
+    ResearchOutput,
+)
 
 __all__ = [
     "ResearchClient",
     "AsyncResearchClient",
-    "ResearchTaskId",
-    "ResearchTask",
-    "ListResearchTasksResponse",
+    "ResearchTyped",
+    "AsyncResearchTyped",
+    "ResearchDto",
+    "ResearchEvent",
+    "ResearchDefinitionEvent",
+    "ResearchOutputEvent",
+    "ResearchPlanDefinitionEvent",
+    "ResearchPlanOperationEvent",
+    "ResearchPlanOutputEvent",
+    "ResearchTaskDefinitionEvent",
+    "ResearchTaskOperationEvent",
+    "ResearchTaskOutputEvent",
+    "ListResearchResponseDto",
+    "CostDollars",
+    "ResearchOutput",
 ]

exa_py/research/async_client.py

@@ -0,0 +1,310 @@
+"""Asynchronous Research API client."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import (
+    Any,
+    AsyncGenerator,
+    Dict,
+    Generic,
+    Literal,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from pydantic import BaseModel, TypeAdapter
+
+from .base import AsyncResearchBaseClient
+from .models import (
+    ResearchDto,
+    ResearchEvent,
+    ListResearchResponseDto,
+)
+from .utils import (
+    async_stream_sse_events,
+    is_pydantic_model,
+    pydantic_to_json_schema,
+)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class AsyncResearchTyped(Generic[T]):
+    """Wrapper for typed research responses in async context."""
+
+    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 AsyncResearchClient(AsyncResearchBaseClient):
+    """Asynchronous client for the Research API."""
+
+    @overload
+    async def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+    ) -> ResearchDto: ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+        output_schema: Dict[str, Any],
+    ) -> ResearchDto: ...
+
+    @overload
+    async def create(
+        self,
+        *,
+        instructions: str,
+        model: Literal["exa-research", "exa-research-pro"] = "exa-research",
+        output_schema: Type[T],
+    ) -> ResearchDto: ...
+
+    async 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 = await self.request("", method="POST", data=payload)
+        adapter = TypeAdapter(ResearchDto)
+        return adapter.validate_python(response)
+
+    @overload
+    async def get(
+        self,
+        research_id: str,
+    ) -> ResearchDto: ...
+
+    @overload
+    async def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[False] = False,
+        events: bool = False,
+    ) -> ResearchDto: ...
+
+    @overload
+    async def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[True],
+        events: Optional[bool] = None,
+    ) -> AsyncGenerator[ResearchEvent, None]: ...
+
+    @overload
+    async def get(
+        self,
+        research_id: str,
+        *,
+        stream: Literal[False] = False,
+        events: bool = False,
+        output_schema: Type[T],
+    ) -> AsyncResearchTyped[T]: ...
+
+    async def get(
+        self,
+        research_id: str,
+        *,
+        stream: bool = False,
+        events: bool = False,
+        output_schema: Optional[Type[BaseModel]] = None,
+    ) -> Union[ResearchDto, AsyncResearchTyped, AsyncGenerator[ResearchEvent, 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 async 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 = await self.request(
+                f"/{research_id}", method="GET", params=params, stream=True
+            )
+            return async_stream_sse_events(response)
+        else:
+            response = await 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 AsyncResearchTyped(research, parsed)
+                except Exception:
+                    # If parsing fails, return the regular research object
+                    return research
+
+            return research
+
+    async 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 = await self.request("", method="GET", params=params)
+        return ListResearchResponseDto.model_validate(response)
+
+    @overload
+    async def poll_until_finished(
+        self,
+        research_id: str,
+        *,
+        poll_interval: int = 1000,
+        timeout_ms: int = 600000,
+        events: bool = False,
+    ) -> ResearchDto: ...
+
+    @overload
+    async def poll_until_finished(
+        self,
+        research_id: str,
+        *,
+        poll_interval: int = 1000,
+        timeout_ms: int = 600000,
+        events: bool = False,
+        output_schema: Type[T],
+    ) -> AsyncResearchTyped[T]: ...
+
+    async 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, AsyncResearchTyped]:
+        """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 = asyncio.get_event_loop().time()
+        consecutive_failures = 0
+
+        while True:
+            try:
+                if output_schema:
+                    result = await self.get(
+                        research_id, events=events, output_schema=output_schema
+                    )
+                else:
+                    result = await 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 asyncio.get_event_loop().time() - start_time > timeout_sec:
+                raise TimeoutError(
+                    f"Research {research_id} did not complete within {timeout_ms}ms"
+                )
+
+            await asyncio.sleep(poll_interval_sec)