lmnr 0.4.65__py3-none-any.whl → 0.5.0__py3-none-any.whl

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

@@ -0,0 +1,215 @@
+"""Agent resource for interacting with Laminar agents."""
+
+from typing import (
+    AsyncGenerator,
+    AsyncIterator,
+    Awaitable,
+    Literal,
+    Optional,
+    Union,
+    overload,
+)
+import uuid
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+from lmnr.sdk.types import (
+    AgentOutput,
+    LaminarSpanContext,
+    ModelProvider,
+    RunAgentRequest,
+    RunAgentResponseChunk,
+)
+
+from opentelemetry import trace
+
+
+class AsyncAgent(BaseAsyncResource):
+    """Resource for interacting with Laminar agents."""
+
+    @overload
+    async def run(
+        self,
+        prompt: str,
+        stream: Literal[True],
+        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
+        model_provider: Optional[ModelProvider] = None,
+        model: Optional[str] = None,
+        enable_thinking: bool = True,
+    ) -> AsyncIterator[RunAgentResponseChunk]:
+        """Run Laminar index agent in streaming mode.
+
+        Args:
+            prompt (str): prompt for the agent
+            stream (Literal[True]): whether to stream the agent's response
+            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
+            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.
+
+        Returns:
+            AsyncIterator[RunAgentResponseChunk]: a generator of response chunks
+        """
+        pass
+
+    @overload
+    async def run(
+        self,
+        prompt: str,
+        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
+        model_provider: Optional[ModelProvider] = None,
+        model: Optional[str] = None,
+        enable_thinking: bool = True,
+    ) -> AgentOutput:
+        """Run Laminar index agent.
+
+        Args:
+            prompt (str): prompt for the agent
+            parent_span_context (Optional[LaminarSpanContext], optional): span context if the agent is part of a trace
+            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.
+
+        Returns:
+            AgentOutput: agent output
+        """
+        pass
+
+    @overload
+    async def run(
+        self,
+        prompt: str,
+        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
+        model_provider: Optional[ModelProvider] = None,
+        model: Optional[str] = None,
+        stream: Literal[False] = False,
+        enable_thinking: bool = True,
+    ) -> AgentOutput:
+        """Run Laminar index agent.
+
+        Args:
+            prompt (str): prompt for the agent
+            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
+            model_provider (Optional[ModelProvider], optional): LLM model provider
+            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.
+
+        Returns:
+            AgentOutput: agent output
+        """
+        pass
+
+    async def run(
+        self,
+        prompt: str,
+        parent_span_context: Optional[Union[LaminarSpanContext, str]] = None,
+        model_provider: Optional[ModelProvider] = None,
+        model: Optional[str] = None,
+        stream: bool = False,
+        enable_thinking: bool = True,
+    ) -> Union[AgentOutput, Awaitable[AsyncIterator[RunAgentResponseChunk]]]:
+        """Run Laminar index agent.
+
+        Args:
+            prompt (str): prompt for the agent
+            parent_span_context (Optional[Union[LaminarSpanContext, str]], optional): span context if the agent is part of a trace
+            model_provider (Optional[ModelProvider], optional): LLM model provider
+            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.
+
+        Returns:
+            Union[AgentOutput, AsyncIterator[RunAgentResponseChunk]]: agent output or a generator of response chunks
+        """
+        if parent_span_context is None:
+            span = trace.get_current_span()
+            if span != trace.INVALID_SPAN:
+                parent_span_context = LaminarSpanContext(
+                    trace_id=uuid.UUID(int=span.get_span_context().trace_id),
+                    span_id=uuid.UUID(int=span.get_span_context().span_id),
+                    is_remote=span.get_span_context().is_remote,
+                )
+        if parent_span_context is not None and isinstance(
+            parent_span_context, LaminarSpanContext
+        ):
+            parent_span_context = str(parent_span_context)
+        request = RunAgentRequest(
+            prompt=prompt,
+            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.
+            # 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,
+        )
+
+        # For streaming case, use a generator function
+        if stream:
+            return self.__run_streaming(request)
+        else:
+            # For non-streaming case, process all chunks and return the final result
+            return await self.__run_non_streaming(request)
+
+    async def __run_streaming(
+        self, request: RunAgentRequest
+    ) -> AsyncGenerator[RunAgentResponseChunk, None]:
+        """Run agent in streaming mode.
+
+        Args:
+            request (RunAgentRequest): The request to run the agent with.
+
+        Yields:
+            RunAgentResponseChunk: Chunks of the agent's response.
+        """
+        async with self._client.stream(
+            "POST",
+            self._base_url + "/v1/agent/run",
+            json=request.to_dict(),
+            headers=self._headers(),
+        ) as response:
+            async for line in response.aiter_lines():
+                line = str(line)
+                if line.startswith("[DONE]"):
+                    break
+                if not line.startswith("data: "):
+                    continue
+                line = line[6:]
+                if line:
+                    chunk = RunAgentResponseChunk.model_validate_json(line)
+                    yield chunk.root
+
+    async def __run_non_streaming(self, request: RunAgentRequest) -> AgentOutput:
+        """Run agent in non-streaming mode.
+
+        Args:
+            request (RunAgentRequest): The request to run the agent with.
+
+        Returns:
+            AgentOutput: The agent's output.
+        """
+        final_chunk = None
+
+        async with self._client.stream(
+            "POST",
+            self._base_url + "/v1/agent/run",
+            json=request.to_dict(),
+            headers=self._headers(),
+        ) as response:
+            async for line in response.aiter_lines():
+                line = str(line)
+                if line.startswith("[DONE]"):
+                    break
+                if not line.startswith("data: "):
+                    continue
+                line = line[6:]
+                if line:
+                    chunk = RunAgentResponseChunk.model_validate_json(line)
+                    if chunk.root.chunkType == "finalOutput":
+                        final_chunk = chunk.root
+
+        return final_chunk.content if final_chunk is not None else AgentOutput()

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

@@ -0,0 +1,32 @@
+"""Base class for resource objects."""
+
+import httpx
+
+
+class BaseAsyncResource:
+    """Base class for all API resources."""
+
+    def __init__(self, client: httpx.AsyncClient, base_url: str, project_api_key: str):
+        """Initialize the resource.
+
+        Args:
+            client (httpx.AsyncClient): HTTP client instance
+            base_url (str): Base URL for the API
+            project_api_key (str): Project API key
+        """
+        self._client = client
+        self._base_url = base_url
+        self._project_api_key = project_api_key
+
+    def _headers(self) -> dict[str, str]:
+        """Generate request headers with authentication.
+
+        Returns:
+            dict[str, str]: Headers dictionary
+        """
+        assert self._project_api_key is not None, "Project API key is not set"
+        return {
+            "Authorization": "Bearer " + self._project_api_key,
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }

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

@@ -0,0 +1,40 @@
+"""Resource for sending browser events."""
+
+import gzip
+import json
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+
+from lmnr.version import PYTHON_VERSION, __version__
+
+
+class AsyncBrowserEvents(BaseAsyncResource):
+    """Resource for sending browser events."""
+
+    async def send(
+        self,
+        session_id: str,
+        trace_id: str,
+        events: list[dict],
+    ):
+        url = self._base_url + "/v1/browser-sessions/events"
+        payload = {
+            "sessionId": session_id,
+            "traceId": trace_id,
+            "events": events,
+            "source": f"python@{PYTHON_VERSION}",
+            "sdkVersion": __version__,
+        }
+        compressed_payload = gzip.compress(json.dumps(payload).encode("utf-8"))
+        response = await self._client.post(
+            url,
+            content=compressed_payload,
+            headers={
+                **self._headers(),
+                "Content-Encoding": "gzip",
+            },
+        )
+        if response.status_code != 200:
+            raise ValueError(
+                f"Failed to send events: [{response.status_code}] {response.text}"
+            )

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

@@ -0,0 +1,64 @@
+"""Evals resource for interacting with Laminar evaluations API."""
+
+import uuid
+from typing import Optional
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+from lmnr.sdk.types import (
+    InitEvaluationResponse,
+    EvaluationResultDatapoint,
+)
+
+
+class AsyncEvals(BaseAsyncResource):
+    """Resource for interacting with Laminar evaluations API."""
+
+    async def init(
+        self, name: Optional[str] = None, group_name: Optional[str] = None
+    ) -> InitEvaluationResponse:
+        """Initialize a new evaluation.
+
+        Args:
+            name (Optional[str], optional): Name of the evaluation. Defaults to None.
+            group_name (Optional[str], optional): Group name for the evaluation. Defaults to None.
+
+        Returns:
+            InitEvaluationResponse: The response from the initialization request.
+        """
+        response = await self._client.post(
+            self._base_url + "/v1/evals",
+            json={
+                "name": name,
+                "groupName": group_name,
+            },
+            headers=self._headers(),
+        )
+        resp_json = response.json()
+        return InitEvaluationResponse.model_validate(resp_json)
+
+    async def save_datapoints(
+        self,
+        eval_id: uuid.UUID,
+        datapoints: list[EvaluationResultDatapoint],
+        group_name: Optional[str] = None,
+    ):
+        """Save evaluation datapoints.
+
+        Args:
+            eval_id (uuid.UUID): The evaluation ID.
+            datapoints (list[EvaluationResultDatapoint]): The datapoints to save.
+            group_name (Optional[str], optional): Group name for the datapoints. Defaults to None.
+
+        Raises:
+            ValueError: If there's an error saving the datapoints.
+        """
+        response = await self._client.post(
+            self._base_url + f"/v1/evals/{eval_id}/datapoints",
+            json={
+                "points": [datapoint.to_dict() for datapoint in datapoints],
+                "groupName": group_name,
+            },
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise ValueError(f"Error saving evaluation datapoints: {response.text}")

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

@@ -0,0 +1,89 @@
+"""Pipeline resource for running Laminar pipelines."""
+
+import uuid
+from typing import Optional
+from opentelemetry import trace
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+from lmnr.sdk.types import (
+    NodeInput,
+    PipelineRunError,
+    PipelineRunRequest,
+    PipelineRunResponse,
+)
+
+
+class AsyncPipeline(BaseAsyncResource):
+    """Resource for interacting with Laminar pipelines."""
+
+    async def run(
+        self,
+        pipeline: str,
+        inputs: dict[str, NodeInput],
+        env: dict[str, str] = {},
+        metadata: dict[str, str] = {},
+        parent_span_id: Optional[uuid.UUID] = None,
+        trace_id: Optional[uuid.UUID] = None,
+    ) -> PipelineRunResponse:
+        """Run a pipeline with the given inputs and environment variables.
+
+        Args:
+            pipeline (str): pipeline name
+            inputs (dict[str, NodeInput]): input values for the pipeline
+            env (dict[str, str], optional): environment variables for the pipeline
+            metadata (dict[str, str], optional): metadata for the pipeline run
+            parent_span_id (Optional[uuid.UUID], optional): parent span id for the pipeline
+            trace_id (Optional[uuid.UUID], optional): trace id for the pipeline
+
+        Raises:
+            ValueError: if the project API key is not set
+            PipelineRunError: if the pipeline run fails
+
+        Returns:
+            PipelineRunResponse: response from the pipeline run
+        """
+        if self._project_api_key is None:
+            raise ValueError(
+                "Please initialize the Laminar object with your project "
+                "API key or set the LMNR_PROJECT_API_KEY environment variable"
+            )
+
+        current_span = trace.get_current_span()
+        if current_span != trace.INVALID_SPAN:
+            parent_span_id = parent_span_id or uuid.UUID(
+                int=current_span.get_span_context().span_id
+            )
+            trace_id = trace_id or uuid.UUID(
+                int=current_span.get_span_context().trace_id
+            )
+
+        request = PipelineRunRequest(
+            inputs=inputs,
+            pipeline=pipeline,
+            env=env or {},
+            metadata=metadata,
+            parent_span_id=parent_span_id,
+            trace_id=trace_id,
+        )
+
+        response = await self._client.post(
+            self._base_url + "/v1/pipeline/run",
+            json=request.to_dict(),
+            headers=self._headers(),
+        )
+
+        if response.status_code != 200:
+            raise PipelineRunError(response)
+
+        try:
+            from pydantic.alias_generators import to_snake
+
+            resp_json = response.json()
+            keys = list(resp_json.keys())
+            for key in keys:
+                value = resp_json[key]
+                del resp_json[key]
+                resp_json[to_snake(key)] = value
+            return PipelineRunResponse(**resp_json)
+        except Exception:
+            raise PipelineRunError(response)

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

@@ -0,0 +1,60 @@
+"""SemanticSearch resource for interacting with Laminar semantic search API."""
+
+import uuid
+from typing import Optional
+
+from lmnr.sdk.client.asynchronous.resources.base import BaseAsyncResource
+from lmnr.sdk.types import (
+    SemanticSearchRequest,
+    SemanticSearchResponse,
+)
+
+
+class AsyncSemanticSearch(BaseAsyncResource):
+    """Resource for interacting with Laminar semantic search API."""
+
+    async def search(
+        self,
+        query: str,
+        dataset_id: uuid.UUID,
+        limit: Optional[int] = None,
+        threshold: Optional[float] = None,
+    ) -> SemanticSearchResponse:
+        """Perform a semantic search on the given dataset.
+
+        Args:
+            query (str): query to search for
+            dataset_id (uuid.UUID): dataset ID created in the UI
+            limit (Optional[int], optional): maximum number of results to return
+            threshold (Optional[float], optional): lowest similarity score to return
+
+        Raises:
+            ValueError: if an error happens while performing the semantic search
+
+        Returns:
+            SemanticSearchResponse: response from the semantic search
+        """
+        request = SemanticSearchRequest(
+            query=query,
+            dataset_id=dataset_id,
+            limit=limit,
+            threshold=threshold,
+        )
+        response = await self._client.post(
+            self._base_url + "/v1/semantic-search",
+            json=request.to_dict(),
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise ValueError(
+                f"Error performing semantic search: [{response.status_code}] {response.text}"
+            )
+        try:
+            resp_json = response.json()
+            for result in resp_json["results"]:
+                result["dataset_id"] = uuid.UUID(result["datasetId"])
+            return SemanticSearchResponse(**resp_json)
+        except Exception as e:
+            raise ValueError(
+                f"Error parsing semantic search response: status={response.status_code} error={e}"
+            )

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

@@ -0,0 +1,7 @@
+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"]