lmnr 0.5.1a0__py3-none-any.whl → 0.5.2__py3-none-any.whl

lmnr/sdk/browser/playwright_otel.py

@@ -2,7 +2,9 @@ import logging
 import uuid
 
 from lmnr.sdk.browser.pw_utils import handle_navigation_async, handle_navigation_sync
-from lmnr.sdk.browser.utils import with_tracer_wrapper
+from lmnr.sdk.browser.utils import with_tracer_and_client_wrapper
+from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
+from lmnr.sdk.client.synchronous.sync_client import LaminarClient
 from lmnr.version import __version__
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
@@ -16,11 +18,11 @@ from opentelemetry.trace import (
     set_span_in_context,
 )
 from opentelemetry.context import get_current
-from typing import Collection, Optional
+from typing import Collection
 from wrapt import wrap_function_wrapper
 
 try:
-    from playwright.async_api import Browser
+    from playwright.async_api import Browser, BrowserContext
     from playwright.sync_api import (
         Browser as SyncBrowser,
         BrowserContext as SyncBrowserContext,
@@ -37,12 +39,12 @@ _instruments = ("playwright >= 1.9.0",)
 logger = logging.getLogger(__name__)
 
 _context_spans: dict[str, Span] = {}
-_project_api_key: Optional[str] = None
-_base_http_url: Optional[str] = None
 
 
-@with_tracer_wrapper
-def _wrap_new_page(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+@with_tracer_and_client_wrapper
+def _wrap_new_page(
+    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
     with tracer.start_as_current_span(
         f"{to_wrap.get('object')}.{to_wrap.get('method')}"
     ) as span:
@@ -50,15 +52,13 @@ def _wrap_new_page(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
         session_id = str(uuid.uuid4().hex)
         trace_id = format(get_current_span().get_span_context().trace_id, "032x")
         span.set_attribute("lmnr.internal.has_browser_session", True)
-        handle_navigation_sync(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        )
+        handle_navigation_sync(page, session_id, trace_id, client)
         return page
 
 
-@with_tracer_wrapper
+@with_tracer_and_client_wrapper
 async def _wrap_new_page_async(
-    tracer: Tracer, to_wrap, wrapped, instance, args, kwargs
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
     with tracer.start_as_current_span(
         f"{to_wrap.get('object')}.{to_wrap.get('method')}"
@@ -67,14 +67,14 @@ async def _wrap_new_page_async(
         session_id = str(uuid.uuid4().hex)
         trace_id = format(span.get_span_context().trace_id, "032x")
         span.set_attribute("lmnr.internal.has_browser_session", True)
-        await handle_navigation_async(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        )
+        await handle_navigation_async(page, session_id, trace_id, client)
         return page
 
 
-@with_tracer_wrapper
-def _wrap_new_browser_sync(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+@with_tracer_and_client_wrapper
+def _wrap_new_browser_sync(
+    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
     global _context_spans
     browser: SyncBrowser = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
@@ -90,20 +90,16 @@ def _wrap_new_browser_sync(tracer: Tracer, to_wrap, wrapped, instance, args, kwa
         trace_id = format(span.get_span_context().trace_id, "032x")
         context.on(
             "page",
-            lambda page: handle_navigation_sync(
-                page, session_id, trace_id, _project_api_key, _base_http_url
-            ),
+            lambda page: handle_navigation_sync(page, session_id, trace_id, client),
         )
         for page in context.pages:
-            handle_navigation_sync(
-                page, session_id, trace_id, _project_api_key, _base_http_url
-            )
+            handle_navigation_sync(page, session_id, trace_id, client)
     return browser
 
 
-@with_tracer_wrapper
+@with_tracer_and_client_wrapper
 async def _wrap_new_browser_async(
-    tracer: Tracer, to_wrap, wrapped, instance, args, kwargs
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
     global _context_spans
     browser: Browser = await wrapped(*args, **kwargs)
@@ -120,20 +116,18 @@ async def _wrap_new_browser_async(
         trace_id = format(span.get_span_context().trace_id, "032x")
 
         async def handle_page_navigation(page):
-            return await handle_navigation_async(
-                page, session_id, trace_id, _project_api_key, _base_http_url
-            )
+            return await handle_navigation_async(page, session_id, trace_id, client)
 
         context.on("page", handle_page_navigation)
         for page in context.pages:
-            await handle_navigation_async(
-                page, session_id, trace_id, _project_api_key, _base_http_url
-            )
+            await handle_navigation_async(page, session_id, trace_id, client)
     return browser
 
 
-@with_tracer_wrapper
-def _wrap_new_context_sync(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+@with_tracer_and_client_wrapper
+def _wrap_new_context_sync(
+    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
     context: SyncBrowserContext = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
     span = get_current_span()
@@ -148,20 +142,16 @@ def _wrap_new_context_sync(tracer: Tracer, to_wrap, wrapped, instance, args, kwa
 
     context.on(
         "page",
-        lambda page: handle_navigation_sync(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        ),
+        lambda page: handle_navigation_sync(page, session_id, trace_id, client),
     )
     for page in context.pages:
-        handle_navigation_sync(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        )
+        handle_navigation_sync(page, session_id, trace_id, client)
     return context
 
 
-@with_tracer_wrapper
+@with_tracer_and_client_wrapper
 async def _wrap_new_context_async(
-    tracer: Tracer, to_wrap, wrapped, instance, args, kwargs
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
     context: SyncBrowserContext = await wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
@@ -176,21 +166,18 @@ async def _wrap_new_context_async(
     trace_id = format(span.get_span_context().trace_id, "032x")
 
     async def handle_page_navigation(page):
-        return await handle_navigation_async(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        )
+        return await handle_navigation_async(page, session_id, trace_id, client)
 
     context.on("page", handle_page_navigation)
     for page in context.pages:
-        await handle_navigation_async(
-            page, session_id, trace_id, _project_api_key, _base_http_url
-        )
+        await handle_navigation_async(page, session_id, trace_id, client)
     return context
 
 
-@with_tracer_wrapper
+@with_tracer_and_client_wrapper
 def _wrap_close_browser_sync(
     tracer: Tracer,
+    client: LaminarClient,
     to_wrap,
     wrapped,
     instance: SyncBrowser,
@@ -208,9 +195,10 @@ def _wrap_close_browser_sync(
     return wrapped(*args, **kwargs)
 
 
-@with_tracer_wrapper
+@with_tracer_and_client_wrapper
 async def _wrap_close_browser_async(
     tracer: Tracer,
+    client: AsyncLaminarClient,
     to_wrap,
     wrapped,
     instance: Browser,
@@ -332,24 +320,18 @@ WRAPPED_METHODS_ASYNC = [
 
 
 class PlaywrightInstrumentor(BaseInstrumentor):
-    def __init__(self):
+    def __init__(self, client: LaminarClient, async_client: AsyncLaminarClient):
         super().__init__()
+        self.client = client
+        self.async_client = async_client
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        global _project_api_key, _base_http_url
-
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, __version__, tracer_provider)
 
-        if kwargs.get("project_api_key"):
-            _project_api_key = kwargs.get("project_api_key")
-
-        if kwargs.get("base_http_url"):
-            _base_http_url = kwargs.get("base_http_url")
-
         for wrapped_method in WRAPPED_METHODS:
             wrap_package = wrapped_method.get("package")
             wrap_object = wrapped_method.get("object")
@@ -360,6 +342,7 @@ class PlaywrightInstrumentor(BaseInstrumentor):
                     f"{wrap_object}.{wrap_method}",
                     wrapped_method.get("wrapper")(
                         tracer,
+                        self.client,
                         wrapped_method,
                     ),
                 )
@@ -377,6 +360,7 @@ class PlaywrightInstrumentor(BaseInstrumentor):
                     f"{wrap_object}.{wrap_method}",
                     wrapped_method.get("wrapper")(
                         tracer,
+                        self.async_client,
                         wrapped_method,
                     ),
                 )

lmnr/sdk/browser/pw_utils.py

@@ -3,7 +3,6 @@ import logging
 import os
 import time
 import threading
-from typing import Optional
 
 from opentelemetry import trace
 
@@ -27,8 +26,6 @@ logger = logging.getLogger(__name__)
 # Track pages we've already instrumented to avoid double-instrumentation
 instrumented_pages = set()
 async_instrumented_pages = set()
-client: Optional[LaminarClient] = None
-async_client: Optional[AsyncLaminarClient] = None
 
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -223,21 +220,8 @@ async def inject_rrweb_async(page: Page):
 
 @observe(name="playwright.page", ignore_input=True, ignore_output=True)
 def handle_navigation_sync(
-    page: SyncPage,
-    session_id: str,
-    trace_id: str,
-    project_api_key: Optional[str] = None,
-    base_http_url: Optional[str] = None,
+    page: SyncPage, session_id: str, trace_id: str, client: LaminarClient
 ):
-    global client
-    if client is None:
-        client = LaminarClient(base_url=base_http_url, project_api_key=project_api_key)
-    if (base_http_url is not None and base_http_url != client.base_url) or (
-        project_api_key is not None and project_api_key != client.project_api_key
-    ):
-        if client is not None:
-            client.close()
-        client = LaminarClient(base_url=base_http_url, project_api_key=project_api_key)
     trace.get_current_span().set_attribute("lmnr.internal.has_browser_session", True)
     # Check if we've already instrumented this page
     page_id = id(page)
@@ -255,10 +239,11 @@ def handle_navigation_sync(
     inject_rrweb_sync(page)
 
     def collection_loop():
-        while not page.is_closed():
+        while not page.is_closed():  # Stop when page closes
             send_events_sync(page, session_id, trace_id, client)
             time.sleep(2)
 
+        # Clean up when page closes
         if page_id in instrumented_pages:
             instrumented_pages.remove(page_id)
 
@@ -268,25 +253,8 @@ def handle_navigation_sync(
 
 @observe(name="playwright.page", ignore_input=True, ignore_output=True)
 async def handle_navigation_async(
-    page: Page,
-    session_id: str,
-    trace_id: str,
-    project_api_key: Optional[str] = None,
-    base_http_url: Optional[str] = None,
+    page: Page, session_id: str, trace_id: str, client: AsyncLaminarClient
 ):
-    global async_client
-    if async_client is None:
-        async_client = AsyncLaminarClient(
-            base_url=base_http_url, project_api_key=project_api_key
-        )
-    if (base_http_url is not None and base_http_url != async_client.base_url) or (
-        project_api_key is not None and project_api_key != async_client.project_api_key
-    ):
-        if async_client is not None:
-            await async_client.close()
-        async_client = AsyncLaminarClient(
-            base_url=base_http_url, project_api_key=project_api_key
-        )
     trace.get_current_span().set_attribute("lmnr.internal.has_browser_session", True)
     # Check if we've already instrumented this page
     page_id = id(page)
@@ -303,19 +271,19 @@ async def handle_navigation_async(
     page.on("load", lambda: asyncio.create_task(on_load()))
     await inject_rrweb_async(page)
 
-    async def collection_loop(client: AsyncLaminarClient):
+    async def collection_loop():
         try:
-            while not page.is_closed():
+            while not page.is_closed():  # Stop when page closes
                 await send_events_async(page, session_id, trace_id, client)
                 await asyncio.sleep(2)
-
+            # Clean up when page closes
             async_instrumented_pages.remove(page_id)
             logger.info("Event collection stopped")
         except Exception as e:
             logger.error(f"Event collection stopped: {e}")
 
     # Create and store task
-    task = asyncio.create_task(collection_loop(async_client))
+    task = asyncio.create_task(collection_loop())
 
     # Clean up task when page closes
     page.on("close", lambda: task.cancel())

lmnr/sdk/client/asynchronous/async_client.py

@@ -11,8 +11,6 @@ from lmnr.sdk.client.asynchronous.resources import (
     AsyncAgent,
     AsyncBrowserEvents,
     AsyncEvals,
-    AsyncPipeline,
-    AsyncSemanticSearch,
 )
 from lmnr.sdk.utils import from_env
 
@@ -66,12 +64,6 @@ class AsyncLaminarClient:
         )
 
         # Initialize resource objects
-        self.__pipeline = AsyncPipeline(
-            self.__client, self.__base_url, self.__project_api_key
-        )
-        self.__semantic_search = AsyncSemanticSearch(
-            self.__client, self.__base_url, self.__project_api_key
-        )
         self.__agent = AsyncAgent(
             self.__client, self.__base_url, self.__project_api_key
         )
@@ -82,24 +74,6 @@ class AsyncLaminarClient:
             self.__client, self.__base_url, self.__project_api_key
         )
 
-    @property
-    def pipeline(self) -> AsyncPipeline:
-        """Get the Pipeline resource.
-
-        Returns:
-            Pipeline: The Pipeline resource instance.
-        """
-        return self.__pipeline
-
-    @property
-    def semantic_search(self) -> AsyncSemanticSearch:
-        """Get the SemanticSearch resource.
-
-        Returns:
-            SemanticSearch: The SemanticSearch resource instance.
-        """
-        return self.__semantic_search
-
     @property
     def agent(self) -> AsyncAgent:
         """Get the Agent resource.
@@ -137,14 +111,6 @@ class AsyncLaminarClient:
         """
         await self.__client.aclose()
 
-    @property
-    def base_url(self) -> str:
-        return self.__base_url
-
-    @property
-    def project_api_key(self) -> str:
-        return self.__project_api_key
-
     async def __aenter__(self: _T) -> _T:
         return self
 

lmnr/sdk/client/asynchronous/resources/__init__.py

@@ -1,12 +1,8 @@
 from lmnr.sdk.client.asynchronous.resources.agent import AsyncAgent
 from lmnr.sdk.client.asynchronous.resources.browser_events import AsyncBrowserEvents
 from lmnr.sdk.client.asynchronous.resources.evals import AsyncEvals
-from lmnr.sdk.client.asynchronous.resources.pipeline import AsyncPipeline
-from lmnr.sdk.client.asynchronous.resources.semantic_search import AsyncSemanticSearch
 
 __all__ = [
-    "AsyncPipeline",
-    "AsyncSemanticSearch",
     "AsyncAgent",
     "AsyncEvals",
     "AsyncBrowserEvents",

lmnr/sdk/client/asynchronous/resources/agent.py

@@ -35,7 +35,16 @@ class AsyncAgent(BaseAsyncResource):
         model_provider: Optional[ModelProvider] = None,
         model: Optional[str] = None,
         enable_thinking: bool = True,
+        agent_state: Optional[str] = None,
+        storage_state: Optional[str] = None,
         return_screenshots: bool = False,
+        return_agent_state: bool = False,
+        return_storage_state: bool = False,
+        timeout: Optional[int] = None,
+        cdp_url: Optional[str] = None,
+        max_steps: Optional[int] = None,
+        thinking_token_budget: Optional[int] = None,
+        start_url: Optional[str] = None,
     ) -> AsyncIterator[RunAgentResponseChunk]:
         """Run Laminar index agent in streaming mode.
 
@@ -46,7 +55,17 @@ class AsyncAgent(BaseAsyncResource):
             model_provider (Optional[ModelProvider], optional): LLM model provider
             model (Optional[str], optional): LLM model name
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
+            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
+            return_agent_state (bool, optional): whether to return the agent's state in the final chunk. Default to False.
+            return_storage_state (bool, optional): whether to return the storage state in the final chunk. Default to False.
+            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+
         Returns:
             AsyncIterator[RunAgentResponseChunk]: a generator of response chunks
         """
@@ -60,7 +79,16 @@ class AsyncAgent(BaseAsyncResource):
         model_provider: Optional[ModelProvider] = None,
         model: Optional[str] = None,
         enable_thinking: bool = True,
+        agent_state: Optional[str] = None,
+        storage_state: Optional[str] = None,
         return_screenshots: bool = False,
+        return_agent_state: bool = False,
+        return_storage_state: bool = False,
+        timeout: Optional[int] = None,
+        cdp_url: Optional[str] = None,
+        max_steps: Optional[int] = None,
+        thinking_token_budget: Optional[int] = None,
+        start_url: Optional[str] = None,
     ) -> AgentOutput:
         """Run Laminar index agent.
 
@@ -70,7 +98,17 @@ class AsyncAgent(BaseAsyncResource):
             model_provider (Optional[ModelProvider], optional): LLM model provider
             model (Optional[str], optional): LLM model name
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
+            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
+            return_agent_state (bool, optional): whether to return the agent's state. Default to False.
+            return_storage_state (bool, optional): whether to return the storage state. Default to False.
+            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+
         Returns:
             AgentOutput: agent output
         """
@@ -85,7 +123,15 @@ class AsyncAgent(BaseAsyncResource):
         model: Optional[str] = None,
         stream: Literal[False] = False,
         enable_thinking: bool = True,
+        agent_state: Optional[str] = None,
+        storage_state: Optional[str] = None,
         return_screenshots: bool = False,
+        return_agent_state: bool = False,
+        return_storage_state: bool = False,
+        timeout: Optional[int] = None,
+        max_steps: Optional[int] = None,
+        thinking_token_budget: Optional[int] = None,
+        start_url: Optional[str] = None,
     ) -> AgentOutput:
         """Run Laminar index agent.
 
@@ -96,7 +142,16 @@ class AsyncAgent(BaseAsyncResource):
             model (Optional[str], optional): LLM model name
             stream (Literal[False], optional): whether to stream the agent's response
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
+            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
+            return_agent_state (bool, optional): whether to return the agent's state. Default to False.
+            return_storage_state (bool, optional): whether to return the storage state. Default to False.
+            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
         Returns:
             AgentOutput: agent output
         """
@@ -110,7 +165,16 @@ class AsyncAgent(BaseAsyncResource):
         model: Optional[str] = None,
         stream: bool = False,
         enable_thinking: bool = True,
+        agent_state: Optional[str] = None,
+        storage_state: Optional[str] = None,
         return_screenshots: bool = False,
+        return_agent_state: bool = False,
+        return_storage_state: bool = False,
+        timeout: Optional[int] = None,
+        cdp_url: Optional[str] = None,
+        max_steps: Optional[int] = None,
+        thinking_token_budget: Optional[int] = None,
+        start_url: Optional[str] = None,
     ) -> Union[AgentOutput, Awaitable[AsyncIterator[RunAgentResponseChunk]]]:
         """Run Laminar index agent.
 
@@ -121,7 +185,17 @@ class AsyncAgent(BaseAsyncResource):
             model (Optional[str], optional): LLM model name
             stream (bool, optional): whether to stream the agent's response
             enable_thinking (bool, optional): whether to enable thinking on the underlying LLM. Default to True.
+            agent_state (Optional[str], optional): the agent's state as returned by the previous agent run. Default to None.
+            storage_state (Optional[str], optional): the browser's storage state as returned by the previous agent run. Default to None.
             return_screenshots (bool, optional): whether to return screenshots of the agent's states at every step. Default to False.
+            return_agent_state (bool, optional): whether to return the agent's state. Default to False.
+            return_storage_state (bool, optional): whether to return the storage state. Default to False.
+            timeout (Optional[int], optional): timeout seconds for the agent's response. Default to None.
+            cdp_url (Optional[str], optional): Chrome DevTools Protocol URL of an existing browser session. Default to None.
+            max_steps (Optional[int], optional): maximum number of steps the agent can take. If not set, the backend will use a default value (currently 100). Default to None.
+            thinking_token_budget (Optional[int], optional): maximum number of tokens the underlying LLM can spend on thinking in each step, if supported by the model. Default to None.
+            start_url (Optional[str], optional): the URL to start the agent on. Must be a valid URL - refer to https://playwright.dev/docs/api/class-page#page-goto. If not specified, the agent infers this from the prompt. Default to None.
+
         Returns:
             Union[AgentOutput, AsyncIterator[RunAgentResponseChunk]]: agent output or a generator of response chunks
         """
@@ -142,15 +216,23 @@ class AsyncAgent(BaseAsyncResource):
             parent_span_context=parent_span_context,
             model_provider=model_provider,
             model=model,
-            # We always connect to stream, because our TLS listeners on AWS
-            # Network load balancers have a hard fixed idle timeout of 350 seconds.
+            agent_state=agent_state,
+            storage_state=storage_state,
+            # We always connect to stream, because our network configuration
+            # has a hard fixed idle timeout of 350 seconds.
             # This means that if we don't stream, the connection will be closed.
             # For now, we just return the content of the final chunk if `stream` is
             # `False`.
-            # https://aws.amazon.com/blogs/networking-and-content-delivery/introducing-nlb-tcp-configurable-idle-timeout/
             stream=True,
             enable_thinking=enable_thinking,
             return_screenshots=return_screenshots,
+            return_agent_state=return_agent_state,
+            return_storage_state=return_storage_state,
+            timeout=timeout,
+            cdp_url=cdp_url,
+            max_steps=max_steps,
+            thinking_token_budget=thinking_token_budget,
+            start_url=start_url,
         )
 
         # For streaming case, use a generator function
@@ -174,7 +256,7 @@ class AsyncAgent(BaseAsyncResource):
         async with self._client.stream(
             "POST",
             self._base_url + "/v1/agent/run",
-            json=request.to_dict(),
+            json=request.model_dump(by_alias=True),
             headers=self._headers(),
         ) as response:
             async for line in response.aiter_lines():
@@ -187,6 +269,8 @@ class AsyncAgent(BaseAsyncResource):
                 if line:
                     chunk = RunAgentResponseChunk.model_validate_json(line)
                     yield chunk.root
+                    if chunk.root.chunk_type in ["finalOutput", "error"]:
+                        break
 
     async def __run_non_streaming(self, request: RunAgentRequest) -> AgentOutput:
         """Run agent in non-streaming mode.
@@ -202,9 +286,11 @@ class AsyncAgent(BaseAsyncResource):
         async with self._client.stream(
             "POST",
             self._base_url + "/v1/agent/run",
-            json=request.to_dict(),
+            json=request.model_dump(by_alias=True),
             headers=self._headers(),
         ) as response:
+            if response.status_code != 200:
+                raise RuntimeError(await response.read())
             async for line in response.aiter_lines():
                 line = str(line)
                 if line.startswith("[DONE]"):
@@ -214,7 +300,11 @@ class AsyncAgent(BaseAsyncResource):
                 line = line[6:]
                 if line:
                     chunk = RunAgentResponseChunk.model_validate_json(line)
-                    if chunk.root.chunkType == "finalOutput":
+                    if chunk.root.chunk_type == "finalOutput":
                         final_chunk = chunk.root
+                    elif chunk.root.chunk_type == "error":
+                        raise RuntimeError(chunk.root.error)
+                    elif chunk.root.chunk_type == "timeout":
+                        raise TimeoutError("Agent timed out")
 
         return final_chunk.content if final_chunk is not None else AgentOutput()

lmnr/sdk/client/synchronous/resources/__init__.py

@@ -1,7 +1,5 @@
 from lmnr.sdk.client.synchronous.resources.agent import Agent
 from lmnr.sdk.client.synchronous.resources.browser_events import BrowserEvents
 from lmnr.sdk.client.synchronous.resources.evals import Evals
-from lmnr.sdk.client.synchronous.resources.pipeline import Pipeline
-from lmnr.sdk.client.synchronous.resources.semantic_search import SemanticSearch
 
-__all__ = ["Pipeline", "SemanticSearch", "Agent", "Evals", "BrowserEvents"]
+__all__ = ["Agent", "Evals", "BrowserEvents"]