exa-py 1.12.3__py3-none-any.whl → 1.12.4__py3-none-any.whl

exa_py/api.py

@@ -39,7 +39,6 @@ from exa_py.utils import (
 from .websets import WebsetsClient
 from .websets.core.base import ExaJSONEncoder
 from .research.client import ResearchClient, AsyncResearchClient
-from .research.models import ResearchTaskResponse  # noqa: E402,F401
 
 is_beta = os.getenv("IS_BETA") == "True"
 
@@ -839,37 +838,6 @@ def nest_fields(original_dict: Dict, fields_to_nest: List[str], new_key: str):
     return original_dict
 
 
-@dataclass
-class ResearchTaskResponse:
-    """A class representing the response for a research task.
-
-    Attributes:
-        id (str): The unique identifier for the research request.
-        status (str): Status of the research request.
-        output (Optional[Dict[str, Any]]): The answer structured as JSON, if available.
-        citations (Optional[Dict[str, List[_Result]]]): List of citations used to generate the answer, grouped by root field in the output schema.
-    """
-
-    id: str
-    status: str
-    output: Optional[Dict[str, Any]]
-    citations: Dict[str, List[_Result]]
-
-    def __str__(self):
-        output_repr = (
-            json.dumps(self.output, indent=2, ensure_ascii=False)
-            if self.output is not None
-            else "None"
-        )
-        citations_str = "\n\n".join(str(src) for src in self.citations)
-        return (
-            f"ID: {self.id}\n"
-            f"Status: {self.status}\n"
-            f"Output: {output_repr}\n\n"
-            f"Citations:\n{citations_str}"
-        )
-
-
 class Exa:
     """A client for interacting with Exa API."""
 
@@ -877,7 +845,7 @@ class Exa:
         self,
         api_key: Optional[str],
         base_url: str = "https://api.exa.ai",
-        user_agent: str = "exa-py 1.12.3",
+        user_agent: str = "exa-py 1.12.4",
     ):
         """Initialize the Exa client with the provided API key and optional base URL and user agent.
 
@@ -909,7 +877,6 @@ class Exa:
         data: Optional[Union[Dict[str, Any], str]] = None,
         method: str = "POST",
         params: Optional[Dict[str, Any]] = None,
-        force_stream: Optional[bool] = False,
     ) -> Union[Dict[str, Any], requests.Response]:
         """Send a request to the Exa API, optionally streaming if data['stream'] is True.
 
@@ -934,7 +901,7 @@ 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")) or force_stream:
+        if data and data.get("stream"):
             res = requests.post(
                 self.base_url + endpoint,
                 data=json_data,
@@ -1974,9 +1941,7 @@ class AsyncExa(Exa):
             )
         return self._client
 
-    async def async_request(
-        self, endpoint: str, data, force_stream: Optional[bool] = False
-    ):
+    async def async_request(self, endpoint: str, data):
         """Send a POST request to the Exa API, optionally streaming if data['stream'] is True.
 
         Args:
@@ -1990,7 +1955,7 @@ class AsyncExa(Exa):
         Raises:
             ValueError: If the request fails (non-200 status code).
         """
-        if data.get("stream") or force_stream:
+        if data.get("stream"):
             request = httpx.Request(
                 "POST", self.base_url + endpoint, json=data, headers=self.headers
             )

exa_py/research/__init__.py

@@ -1,8 +1,9 @@
 from .client import ResearchClient, AsyncResearchClient
-from .models import ResearchTaskResponse
+from .models import ResearchTask
 
 __all__ = [
     "ResearchClient",
     "AsyncResearchClient",
-    "ResearchTaskResponse",
+    "ResearchTaskId",
+    "ResearchTask",
 ]

exa_py/research/client.py

@@ -9,14 +9,14 @@ block, but at runtime we only pay the cost if/when a helper is actually used.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict
 
 if TYPE_CHECKING:  # pragma: no cover – only for static analysers
     # Import with full type info when static type-checking.  `_Result` still
     # lives in ``exa_py.api`` but the response model moved to
     # ``exa_py.research.models``.
     from ..api import _Result  # noqa: F401
-    from .models import ResearchTaskResponse  # noqa: F401
+    from .models import ResearchTask, ResearchTaskId  # noqa: F401
 
 # ---------------------------------------------------------------------------
 # Public, user-facing clients
@@ -31,118 +31,82 @@ class ResearchClient:
         # can piggy-back on its HTTP plumbing (headers, base URL, retries, …).
         self._client = parent_client
 
-    # ------------------------------------------------------------------
-    # API surface
-    # ------------------------------------------------------------------
     def create_task(
         self,
         *,
         input_instructions: str,
         output_schema: Dict[str, Any],
-    ) -> "ResearchTaskResponse":
-        """Submit a research request to the Exa backend.
-
-        The public API remains synchronous – the function only returns once
-        the task has finished and the final structured answer is available.
-        Internally, however, the endpoint now streams *progress* updates via
-        Server-Sent Events (SSE). We therefore initiate a streaming request
-        and keep reading until we receive the terminal ``{"tag": "complete"}``
-        chunk, which carries the exact same payload shape that the blocking
-        variant returned previously.  Any ``{"tag": "progress"}`` chunks are
-        ignored, while ``{"tag": "error"}`` chunks result in an exception.
+    ) -> "ResearchTaskId":
+        """Submit a research request and return the *task identifier*."""
+        payload = {
+            "input": {"instructions": input_instructions},
+            "output": {"schema": output_schema},
+        }
+
+        raw_response: Dict[str, Any] = self._client.request("/research/tasks", payload)
+
+        # Defensive checks so that we fail loudly if the contract changes.
+        if not isinstance(raw_response, dict) or "id" not in raw_response:
+            raise RuntimeError(
+                f"Unexpected response while creating research task: {raw_response}"
+            )
+
+        # Lazily import to avoid circular deps at runtime.
+        from .models import ResearchTaskId  # noqa: WPS433 – runtime import
+
+        return ResearchTaskId(id=raw_response["id"])
+
+    def get_task(
+        self, id: str
+    ) -> "ResearchTask":  # noqa: D401 – imperative mood is fine
+        """Fetch the current status / result for a research task."""
+        endpoint = f"/research/tasks/{id}"
+
+        # The new endpoint is a simple GET.
+        raw_response: Dict[str, Any] = self._client.request(endpoint, method="GET")
+
+        return _build_research_task(raw_response)
+
+    # ------------------------------------------------------------------
+    # Convenience helpers
+    # ------------------------------------------------------------------
+
+    def poll_task(
+        self,
+        id: str,
+        *,
+        poll_interval: float = 1.0,
+        timeout_seconds: int = 15 * 60,
+    ) -> "ResearchTask":
+        """Blocking helper that polls until task completes or fails.
 
         Parameters
         ----------
-        input_instructions:
-            Natural-language instructions that describe *what* should be
-            researched or extracted.
-        output_schema:
-            JSON-schema describing the desired structured output format.
+        id:
+            The ID of the research task to poll.
+        poll_interval:
+            Seconds to wait between successive polls (default 1s).
+        timeout_seconds:
+            Maximum time to wait before raising :class:`TimeoutError` (default 15 min).
         """
 
-        import json
+        import time
 
-        payload = {
-            "input": {"instructions": input_instructions},
-            "output": {"schema": output_schema},
-        }
+        deadline = time.monotonic() + timeout_seconds
 
-        raw_response = self._client.request(
-            "/research/tasks", payload, force_stream=True
-        )
+        while True:
+            task = self.get_task(id)
+            status = task.status.lower() if isinstance(task.status, str) else ""
 
-        def _handle_payload(tag: Optional[str], payload_dict: Dict[str, Any]):
-            """Inner helper handling decoded JSON chunks."""
-            if tag is None:
-                tag_local = payload_dict.get("tag")
-            else:
-                tag_local = tag
-
-            if tag_local == "progress":
-                return None  # ignore
-            if tag_local == "error":
-                msg = payload_dict.get("error", {}).get("message", "Unknown error")
-                raise RuntimeError(f"Research task failed: {msg}")
-            if tag_local == "complete":
-                data_obj = payload_dict.get("data")
-                if data_obj is None:
-                    raise RuntimeError("Malformed 'complete' chunk with no data")
-                return _parse_research_response(data_obj)
-
-            # Fallback: if looks like final object
-            if {"id", "status"}.issubset(payload_dict.keys()):
-                return _parse_research_response(payload_dict)
-            return None
-
-        # ------------------------------------------------------------------
-        # Minimal SSE parser (sync)
-        # ------------------------------------------------------------------
-        event_name: Optional[str] = None
-        data_buf: str = ""
-
-        for raw_line in raw_response.iter_lines(decode_unicode=True):
-            line = raw_line
-            if line == "":
-                if data_buf:
-                    try:
-                        payload_dict = json.loads(data_buf)
-                    except json.JSONDecodeError:
-                        data_buf = ""
-                        event_name = None
-                        continue
-                    maybe_resp = _handle_payload(event_name, payload_dict)
-                    if maybe_resp is not None:
-                        raw_response.close()
-                        return maybe_resp
-                # reset after event
-                data_buf = ""
-                event_name = None
-                continue
-
-            if line.startswith("event:"):
-                event_name = line[len("event:") :].strip()
-            elif line.startswith("data:"):
-                data_buf += line[len("data:") :].strip()
-
-        # Process any remaining buffer (in case stream closed without blank line)
-        if data_buf:
-            try:
-                payload_dict = json.loads(data_buf)
-                maybe_resp = _handle_payload(event_name, payload_dict)
-                if maybe_resp is not None:
-                    raw_response.close()
-                    return maybe_resp
-            except json.JSONDecodeError:
-                pass
-
-        raise RuntimeError("Stream ended before completion of research task")
-
-    def get_task(self, id: str):  # noqa: D401 – imperative mood is fine
-        """Placeholder endpoint – not yet implemented on the server side."""
-        raise NotImplementedError(
-            "`exa.research.get_task` is not available yet. Please open an "
-            "issue if you need this sooner."
-        )
+            if status in {"completed", "failed", "complete", "finished", "done"}:
+                return task
+
+            if time.monotonic() > deadline:
+                raise TimeoutError(
+                    f"Research task {id} did not finish within {timeout_seconds} seconds"
+                )
+
+            time.sleep(poll_interval)
 
 
 class AsyncResearchClient:
@@ -156,102 +120,113 @@ class AsyncResearchClient:
         *,
         input_instructions: str,
         output_schema: Dict[str, Any],
-    ) -> "ResearchTaskResponse":
-        """Async variant mirroring the synchronous implementation above."""
-
-        import json
+    ) -> "ResearchTaskId":
+        """Submit a research request and return the *task identifier* (async)."""
 
         payload = {
             "input": {"instructions": input_instructions},
             "output": {"schema": output_schema},
         }
 
-        raw_response = await self._client.async_request(
-            "/research/tasks", payload, force_stream=True
+        raw_response: Dict[str, Any] = await self._client.async_request(
+            "/research/tasks", payload
         )
 
-        async def _handle_payload_async(
-            tag: Optional[str], payload_dict: Dict[str, Any]
-        ):
-            if tag is None:
-                tag_local = payload_dict.get("tag")
-            else:
-                tag_local = tag
-
-            if tag_local == "progress":
-                return None
-            if tag_local == "error":
-                msg = payload_dict.get("error", {}).get("message", "Unknown error")
-                raise RuntimeError(f"Research task failed: {msg}")
-            if tag_local == "complete":
-                data_obj = payload_dict.get("data")
-                if data_obj is None:
-                    raise RuntimeError("Malformed 'complete' chunk with no data")
-                return _parse_research_response(data_obj)
-            if {"id", "status"}.issubset(payload_dict.keys()):
-                return _parse_research_response(payload_dict)
-            return None
-
-        event_name: Optional[str] = None
-        data_buf: str = ""
-
-        async for line in raw_response.aiter_lines():
-            if line == "":
-                if data_buf:
-                    try:
-                        payload_dict = json.loads(data_buf)
-                    except json.JSONDecodeError:
-                        data_buf = ""
-                        event_name = None
-                        continue
-                    maybe_resp = await _handle_payload_async(event_name, payload_dict)
-                    if maybe_resp is not None:
-                        await raw_response.aclose()
-                        return maybe_resp
-                data_buf = ""
-                event_name = None
-                continue
-
-            if line.startswith("event:"):
-                event_name = line[len("event:") :].strip()
-            elif line.startswith("data:"):
-                data_buf += line[len("data:") :].strip()
-
-        if data_buf:
-            try:
-                payload_dict = json.loads(data_buf)
-                maybe_resp = await _handle_payload_async(event_name, payload_dict)
-                if maybe_resp is not None:
-                    await raw_response.aclose()
-                    return maybe_resp
-            except json.JSONDecodeError:
-                pass
-
-        raise RuntimeError("Stream ended before completion of research task")
-
-    async def get_task(self, id: str):  # noqa: D401
-        raise NotImplementedError(
-            "`exa.research.get_task` is not available yet. Please open an "
-            "issue if you need this sooner."
+        # Defensive checks so that we fail loudly if the contract changes.
+        if not isinstance(raw_response, dict) or "id" not in raw_response:
+            raise RuntimeError(
+                f"Unexpected response while creating research task: {raw_response}"
+            )
+
+        # Lazily import to avoid circular deps at runtime.
+        from .models import ResearchTaskId  # noqa: WPS433 – runtime import
+
+        return ResearchTaskId(id=raw_response["id"])
+
+    async def get_task(self, id: str) -> "ResearchTask":  # noqa: D401
+        """Fetch the current status / result for a research task (async)."""
+
+        endpoint = f"/research/tasks/{id}"
+
+        # Perform GET using the underlying HTTP client because `async_request`
+        # only supports POST semantics.
+        resp = await self._client.client.get(
+            self._client.base_url + endpoint, headers=self._client.headers
         )
 
+        if resp.status_code >= 400:
+            raise RuntimeError(
+                f"Request failed with status code {resp.status_code}: {resp.text}"
+            )
+
+        raw_response: Dict[str, Any] = resp.json()
+
+        return _build_research_task(raw_response)
+
+    # ------------------------------------------------------------------
+    # Convenience helpers
+    # ------------------------------------------------------------------
+
+    async def poll_task(
+        self,
+        id: str,
+        *,
+        poll_interval: float = 1.0,
+        timeout_seconds: int = 15 * 60,
+    ) -> "ResearchTask":
+        """Async helper that polls until task completes or fails.
+
+        Mirrors :py:meth:`ResearchClient.poll_task` but uses ``await`` and
+        :pyfunc:`asyncio.sleep`.  Raises :class:`TimeoutError` on timeout.
+        """
+
+        import asyncio
+        import time
+
+        deadline = time.monotonic() + timeout_seconds
+
+        while True:
+            task = await self.get_task(id)
+            status = task.status.lower() if isinstance(task.status, str) else ""
+
+            if status in {"completed", "failed", "complete", "finished", "done"}:
+                return task
+
+            if time.monotonic() > deadline:
+                raise TimeoutError(
+                    f"Research task {id} did not finish within {timeout_seconds} seconds"
+                )
+
+            await asyncio.sleep(poll_interval)
+
 
 # ---------------------------------------------------------------------------
 # Internal helpers (lazy imports to avoid cycles)
 # ---------------------------------------------------------------------------
 
 
-def _parse_research_response(raw: Dict[str, Any]):
-    """Transform camel-case API payload into rich Python objects."""
-    from .models import ResearchTaskResponse
-    from ..api import _Result, to_snake_case
+def _build_research_task(raw: Dict[str, Any]):
+    """Convert raw API response into a :class:`ResearchTask` instance."""
+
+    # Defensive check – fail loudly if the API contract changes.
+    if not isinstance(raw, dict) or "id" not in raw:
+        raise RuntimeError(f"Unexpected response while fetching research task: {raw}")
+
+    # Lazily import heavy deps to avoid cycles and unnecessary startup cost.
+    from .models import ResearchTask  # noqa: WPS433 – runtime import
+    from ..api import _Result, to_snake_case  # noqa: WPS433 – runtime import
+
+    citations_raw = raw.get("citations", {}) or {}
+    citations_parsed = {
+        key: [_Result(**to_snake_case(c)) for c in cites]
+        for key, cites in citations_raw.items()
+    }
 
-    return ResearchTaskResponse(
+    return ResearchTask(
         id=raw["id"],
         status=raw["status"],
-        output=raw.get("output"),
-        citations={
-            key: [_Result(**to_snake_case(c)) for c in citations]
-            for key, citations in raw.get("citations", {}).items()
-        },
+        instructions=raw.get("instructions", ""),
+        schema=raw.get("schema", {}),
+        data=raw.get("data"),
+        citations=citations_parsed,
     )

exa_py/research/models.py

@@ -12,8 +12,27 @@ if TYPE_CHECKING:  # pragma: no cover – for static analysers only
 
 
 @dataclass
-class ResearchTaskResponse:
-    """Structured response returned from the /research/tasks endpoint.
+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
     ----------
@@ -21,7 +40,11 @@ class ResearchTaskResponse:
         Unique identifier for the research task.
     status:
         Current task status
-    output:
+    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:
@@ -31,27 +54,45 @@ class ResearchTaskResponse:
 
     id: str
     status: str
-    output: Optional[Dict[str, Any]]
+    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
-        output_repr = (
-            json.dumps(self.output, indent=2, ensure_ascii=False)
-            if self.output is not None
+        """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"
         )
-        citations_str = "\n\n".join(str(src) for src in self.citations)
+
+        # 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"Output: {output_repr}\n\n"
+            f"Instructions: {self.instructions}\n"
+            f"Schema:\n{schema_repr}\n"
+            f"Data:\n{data_repr}\n\n"
             f"Citations:\n{citations_str}"
         )
 
 
 __all__ = [
-    "ResearchTaskResponse",
+    "ResearchTaskId",
+    "ResearchTask",
 ]

{exa_py-1.12.3.dist-info → exa_py-1.12.4.dist-info}/METADATA

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

{exa_py-1.12.3.dist-info → exa_py-1.12.4.dist-info}/RECORD

@@ -1,9 +1,9 @@
 exa_py/__init__.py,sha256=M2GC9oSdoV6m2msboW0vMWWl8wrth4o6gmEV4MYLGG8,66
-exa_py/api.py,sha256=CNDOERNAWjzp_XMmSCjH59mPe_8ePOXSJ82qIAWf60g,86106
+exa_py/api.py,sha256=Bn7h_eRvXmwBUmJi2B2JpHAQPrHfbwKf0A-XVXLjqa0,84876
 exa_py/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-exa_py/research/__init__.py,sha256=tFahd_P8Gqd9IkPQPk8GXag4AbxulRPpe8u44RhnB3w,189
-exa_py/research/client.py,sha256=rJZf7iFnIAUBYwx0xqx76TbWx-yX1CKIwCr-t2P7RvA,9930
-exa_py/research/models.py,sha256=n6_E3Kog6Pg-sLnXTsHocyWqz_qdzTckJW3aCAR71Vk,1698
+exa_py/research/__init__.py,sha256=D1xgm4VlbWtRb1cMgshcW4dIyR7IXhOW2s7ihxcE1Jc,195
+exa_py/research/client.py,sha256=Zno5xblfwhX8gWgc4OvI24a-ZS7_g1b32_tr6j7C7Jg,8217
+exa_py/research/models.py,sha256=WXTnALhM9FcVQ95Tzzc5EDKU48hyPhu8RSMmipqCjOk,2982
 exa_py/utils.py,sha256=Rc1FJjoR9LQ7L_OJM91Sd1GNkbHjcLyEvJENhRix6gc,2405
 exa_py/websets/__init__.py,sha256=uOBAb9VrIHrPKoddGOp2ai2KgWlyUVCLMZqfbGOlboA,70
 exa_py/websets/_generator/pydantic/BaseModel.jinja2,sha256=RUDCmPZVamoVx1WudylscYFfDhGoNNtRYlpTvKjAiuA,1276
@@ -19,6 +19,6 @@ exa_py/websets/searches/client.py,sha256=X3f7axWGfecmxf-2tBTX0Yf_--xToz1X8ZHbbud
 exa_py/websets/types.py,sha256=jKnJFAHTFN55EzsusgDce-yux71zVbdSJ1m8utR4EjU,28096
 exa_py/websets/webhooks/__init__.py,sha256=iTPBCxFd73z4RifLQMX6iRECx_6pwlI5qscLNjMOUHE,77
 exa_py/websets/webhooks/client.py,sha256=zsIRMTeJU65yj-zo7Zz-gG02Prtzgcx6utGFSoY4HQQ,4222
-exa_py-1.12.3.dist-info/METADATA,sha256=-OzwAZigBPwXnt9uRSV9XYsGioPwr55DR8WrnASkymk,4098
-exa_py-1.12.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-exa_py-1.12.3.dist-info/RECORD,,
+exa_py-1.12.4.dist-info/METADATA,sha256=DYYA35UrWmW9ND3x5L5YcDr1tPJN05x64UV4nJsRg1k,4098
+exa_py-1.12.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+exa_py-1.12.4.dist-info/RECORD,,