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

lmnr/openllmetry_sdk/opentelemetry/instrumentation/google_genai/utils.py

@@ -0,0 +1,216 @@
+import logging
+import traceback
+import json
+
+from .config import (
+    Config,
+)
+from google.genai import types
+from google.genai._common import BaseModel
+import pydantic
+from opentelemetry.trace import Span
+from typing import Any, Optional, Union
+
+
+def set_span_attribute(span: Span, name: str, value: str):
+    if value is not None:
+        if value != "":
+            span.set_attribute(name, value)
+    return
+
+
+def dont_throw(func):
+    """
+    A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+
+    @param func: The function to wrap
+    @return: The wrapper function
+    """
+    # Obtain a logger specific to the function's module
+    logger = logging.getLogger(func.__module__)
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.debug(
+                "OpenLLMetry failed to trace in %s, error: %s",
+                func.__name__,
+                traceback.format_exc(),
+            )
+            if Config.exception_logger:
+                Config.exception_logger(e)
+
+    return wrapper
+
+
+def to_dict(obj: Union[BaseModel, pydantic.BaseModel, dict]) -> dict[str, Any]:
+    try:
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+        elif isinstance(obj, pydantic.BaseModel):
+            return obj.model_dump()
+        elif isinstance(obj, dict):
+            return obj
+        else:
+            return dict(obj)
+    except Exception:
+        return dict(obj)
+
+
+def process_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+) -> Optional[str]:
+    parts = _process_content_union(content, trace_id, span_id, message_index)
+    if parts is None:
+        return None
+    if isinstance(parts, str):
+        return parts
+    elif isinstance(parts, list):
+        if len(parts) == 1 and isinstance(parts[0], str):
+            return parts[0]
+        return json.dumps(
+            [
+                {"type": "text", "text": part} if isinstance(part, str) else part
+                for part in parts
+            ]
+        )
+    else:
+        return None
+
+
+def _process_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+) -> Union[str, list[str], None]:
+    if isinstance(content, types.Content):
+        parts = to_dict(content).get("parts", [])
+        return [_process_part(part) for part in parts]
+    elif isinstance(content, list):
+        return [_process_part_union(item) for item in content]
+    elif isinstance(content, (types.Part, types.File, str)):
+        return _process_part_union(content)
+    elif isinstance(content, dict):
+        if "parts" in content:
+            return [
+                _process_part_union(
+                    item, trace_id, span_id, message_index, content_index
+                )
+                for content_index, item in enumerate(content.get("parts", []))
+            ]
+        else:
+            # Assume it's PartDict
+            return _process_part_union(content, trace_id, span_id, message_index)
+    else:
+        return None
+
+
+def _process_part_union(
+    content: Union[types.PartDict, types.File, types.Part, str],
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+    content_index: int = 0,
+) -> Optional[str]:
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, types.File):
+        content_dict = to_dict(content)
+        name = (
+            content_dict.get("name")
+            or content_dict.get("display_name")
+            or content_dict.get("uri")
+        )
+        return f"files/{name}"
+    elif isinstance(content, (types.Part, dict)):
+        return _process_part(content, trace_id, span_id, message_index, content_index)
+    else:
+        return None
+
+
+def _process_part(
+    content: types.Part,
+    trace_id: Optional[str] = None,
+    span_id: Optional[str] = None,
+    message_index: int = 0,
+    content_index: int = 0,
+) -> Optional[str]:
+    part_dict = to_dict(content)
+    if part_dict.get("text") is not None:
+        return part_dict.get("text")
+    elif part_dict.get("inline_data"):
+        blob = to_dict(part_dict.get("inline_data"))
+        if blob.get("mime_type").startswith("image/"):
+            return _process_image_item(
+                blob, trace_id, span_id, message_index, content_index
+            )
+        else:
+            # currently, only images are supported
+            return blob.get("mime_type") or "unknown_media"
+    else:
+        return None
+
+
+def role_from_content_union(
+    content: Union[types.ContentUnion, types.ContentUnionDict],
+) -> Optional[str]:
+    if isinstance(content, types.Content):
+        return to_dict(content).get("role")
+    elif isinstance(content, list) and len(content) > 0:
+        return role_from_content_union(content[0])
+    else:
+        return None
+
+
+def with_tracer_wrapper(func):
+    """Helper for providing tracer for wrapper functions."""
+
+    def _with_tracer(tracer, to_wrap):
+        def wrapper(wrapped, instance, args, kwargs):
+            return func(tracer, to_wrap, wrapped, instance, args, kwargs)
+
+        return wrapper
+
+    return _with_tracer
+
+
+def _run_async(method):
+    import asyncio
+    import threading
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop and loop.is_running():
+        thread = threading.Thread(target=lambda: asyncio.run(method))
+        thread.start()
+        thread.join()
+    else:
+        asyncio.run(method)
+
+
+def _process_image_item(
+    blob: dict[str, Any],
+    trace_id: str,
+    span_id: str,
+    message_index: int,
+    content_index: int,
+):
+    # Convert to openai format, so backends can handle it
+    return (
+        {
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/{blob.get('mime_type').split('/')[1]};base64,{blob.get('data')}",
+            },
+        }
+        if Config.convert_image_to_openai_format
+        else blob
+    )

lmnr/openllmetry_sdk/tracing/tracing.py

@@ -84,6 +84,8 @@ class TracerWrapper(object):
     __span_id_lists: dict[int, list[str]] = {}
     __client: LaminarClient = None
     __async_client: AsyncLaminarClient = None
+    __spans_processor: SpanProcessor = None
+    __spans_exporter: SpanExporter = None
 
     def __new__(
         cls,
@@ -115,10 +117,10 @@ class TracerWrapper(object):
             obj.__resource = Resource(attributes=TracerWrapper.resource_attributes)
             obj.__tracer_provider = init_tracer_provider(resource=obj.__resource)
             if processor:
-                obj.__spans_processor: SpanProcessor = processor
+                obj.__spans_processor = processor
                 obj.__spans_processor_original_on_start = processor.on_start
             else:
-                obj.__spans_exporter: SpanExporter = (
+                obj.__spans_exporter = (
                     exporter
                     if exporter
                     else init_spans_exporter(
@@ -126,11 +128,9 @@ class TracerWrapper(object):
                     )
                 )
                 if disable_batch or is_notebook():
-                    obj.__spans_processor: SpanProcessor = SimpleSpanProcessor(
-                        obj.__spans_exporter
-                    )
+                    obj.__spans_processor = SimpleSpanProcessor(obj.__spans_exporter)
                 else:
-                    obj.__spans_processor: SpanProcessor = BatchSpanProcessor(
+                    obj.__spans_processor = BatchSpanProcessor(
                         obj.__spans_exporter,
                         max_export_batch_size=max_export_batch_size,
                     )
@@ -237,8 +237,6 @@ class TracerWrapper(object):
         cls.__span_id_lists = {}
 
     def shutdown(self):
-        self.__spans_processor.force_flush()
-        self.__spans_processor.shutdown()
         self.__tracer_provider.shutdown()
 
     def flush(self):
@@ -382,6 +380,9 @@ def init_instrumentations(
         elif instrument == Instruments.GOOGLE_GENERATIVEAI:
             if init_google_generativeai_instrumentor():
                 instrument_set = True
+        elif instrument == Instruments.GOOGLE_GENAI:
+            if init_google_genai_instrumentor():
+                instrument_set = True
         elif instrument == Instruments.GROQ:
             if init_groq_instrumentor():
                 instrument_set = True
@@ -622,6 +623,27 @@ def init_google_generativeai_instrumentor():
         return False
 
 
+def init_google_genai_instrumentor():
+    try:
+        if is_package_installed("google-genai"):
+            # TODO: uncomment this once we migrate to the contrib package
+            # and is_package_installed(
+            #     "opentelemetry-instrumentation-google-genai"
+            # ):
+            # from opentelemetry.instrumentation.google_genai import (
+            from ..opentelemetry.instrumentation.google_genai import (
+                GoogleGenAiSdkInstrumentor,
+            )
+
+            instrumentor = GoogleGenAiSdkInstrumentor()
+            if not instrumentor.is_instrumented_by_opentelemetry:
+                instrumentor.instrument()
+        return True
+    except Exception as e:
+        module_logger.error(f"Error initializing Google GenAI instrumentor: {e}")
+        return False
+
+
 def init_haystack_instrumentor():
     try:
         if is_package_installed("haystack") and is_package_installed(

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.

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"]