lmnr 0.2.15__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

lmnr/__init__.py

@@ -1,4 +1,4 @@
-from .sdk.endpoint import Laminar
-from .types import ChatMessage, EndpointRunError, EndpointRunResponse, NodeInput
-from .sdk.remote_debugger import RemoteDebugger as LaminarRemoteDebugger
-from .sdk.registry import Registry as Pipeline
+from .sdk.client import Laminar
+from .sdk.decorators import observe, lmnr_context, wrap_llm_call
+from .sdk.interface import trace, TraceContext, SpanContext
+from .sdk.types import ChatMessage, PipelineRunError, PipelineRunResponse, NodeInput

lmnr/sdk/client.py

@@ -0,0 +1,156 @@
+from .tracing_types import Span, Trace
+
+from pydantic.alias_generators import to_snake
+from typing import Any, Optional, Union
+import dotenv
+import json
+import logging
+import os
+import requests
+import uuid
+
+from .types import (
+    PipelineRunError,
+    PipelineRunResponse,
+    NodeInput,
+    PipelineRunRequest,
+)
+
+
+class APIError(Exception):
+    def __init__(self, status: Union[int, str], message: str, details: Any = None):
+        self.message = message
+        self.status = status
+        self.details = details
+
+    def __str__(self):
+        msg = "{0} ({1}): {2}"
+        return msg.format(self.message, self.status, self.details)
+
+
+class Laminar:
+    _base_url = "https://api.lmnr.ai"
+
+    def __init__(self, project_api_key: Optional[str] = None):
+        self.project_api_key = project_api_key or os.environ.get("LMNR_PROJECT_API_KEY")
+        if not self.project_api_key:
+            dotenv_path = dotenv.find_dotenv(usecwd=True)
+            self.project_api_key = dotenv.get_key(
+                dotenv_path=dotenv_path, key_to_get="LMNR_PROJECT_API_KEY"
+            )
+
+    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:
+        """Runs the pipeline with the given inputs
+
+        Args:
+            pipeline (str): name of the Laminar pipeline
+            inputs (dict[str, NodeInput]):
+                inputs to the endpoint's target pipeline.
+                Keys in the dictionary must match input node names
+            env (dict[str, str], optional):
+                Environment variables for the pipeline execution.
+                Defaults to {}.
+            metadata (dict[str, str], optional):
+                any custom metadata to be stored
+                with execution trace. Defaults to {}.
+            parent_span_id (Optional[uuid.UUID], optional):
+                parent span id for the resulting span.
+                Must usually be SpanContext.id()
+                Defaults to None.
+            trace_id (Optional[uuid.UUID], optional):
+                trace id for the resulting trace.
+                Must usually be TraceContext.id()
+                Defaults to None.
+
+        Returns:
+            PipelineRunResponse: response object containing the outputs
+
+        Raises:
+            ValueError: if project API key is not set
+            PipelineRunError: if the endpoint run fails
+        """
+        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"
+            )
+        try:
+            request = PipelineRunRequest(
+                inputs=inputs,
+                pipeline=pipeline,
+                env=env,
+                metadata=metadata,
+                parent_span_id=parent_span_id,
+                trace_id=trace_id,
+            )
+        except Exception as e:
+            raise ValueError(f"Invalid request: {e}")
+
+        response = requests.post(
+            self._base_url + "/v1/pipeline/run",
+            data=json.dumps(request.to_dict()),
+            headers=self._headers(),
+        )
+        if response.status_code != 200:
+            raise PipelineRunError(response)
+        try:
+            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)
+
+    def batch_post_traces(self, batch: list[Union[Span, Trace]]):
+        log = logging.getLogger("laminar.client")
+        url = self._base_url + "/v1/traces"
+        data = json.dumps({"traces": [item.to_dict() for item in batch]})
+        log.debug(f"making request to {url}")
+        headers = self._headers()
+        res = requests.post(url, data=data, headers=headers)
+
+        if res.status_code == 200:
+            log.debug("data uploaded successfully")
+
+        return self._process_response(
+            res, success_message="data uploaded successfully", return_json=False
+        )
+
+    def _process_response(
+        self, res: requests.Response, success_message: str, return_json: bool = True
+    ) -> Union[requests.Response, Any]:
+        log = logging.getLogger("laminar.client")
+        log.debug("received response: %s", res.text)
+        if res.status_code in (200, 201):
+            log.debug(success_message)
+            if return_json:
+                try:
+                    return res.json()
+                except json.JSONDecodeError:
+                    log.error("Response is not valid JSON.")
+                    raise APIError(res.status_code, "Invalid JSON response received")
+            else:
+                return res
+        try:
+            payload = res.json()
+            log.error("received error response: %s", payload)
+            raise APIError(res.status_code, payload)
+        except (KeyError, ValueError):
+            raise APIError(res.status_code, res.text)
+
+    def _headers(self):
+        return {
+            "Authorization": "Bearer " + self.project_api_key,
+            "Content-Type": "application/json",
+        }

lmnr/sdk/collector.py

@@ -0,0 +1,177 @@
+from .client import APIError, Laminar
+from .tracing_types import Span, Trace
+
+from queue import Queue, Empty, Full
+from typing import Union
+
+import atexit
+import backoff
+import logging
+import time
+import threading
+
+
+class Collector(threading.Thread):
+    _log = logging.getLogger("laminar.collector")
+    _queue: Queue[Union[Span, Trace]]
+    _client: Laminar
+    _flush_interval: float
+
+    def __init__(
+        self,
+        queue: Queue[Union[Span, Trace]],
+        client: Laminar,
+        flush_interval: float = 5.0,
+    ):
+        super().__init__()
+        self.daemon = True
+        self._queue = queue
+        self.running = True
+        self._flush_interval = flush_interval
+        self._client = client
+
+    def run(self):
+        """Runs the collector."""
+        self._log.debug("collector is running...")
+        while self.running:
+            self.upload()
+
+    def upload(self):
+        """Upload the next batch of items, return whether successful."""
+        batch = self._next()
+        if len(batch) == 0:
+            return
+
+        try:
+            self._upload_batch(batch)
+        except Exception as e:
+            self._log.exception("error uploading: %s", e)
+        finally:
+            # mark items as acknowledged from queue
+            for _ in batch:
+                self._queue.task_done()
+
+    def pause(self):
+        self.running = False
+
+    def _next(self):
+        items = []
+        start_time = time.monotonic()
+
+        while True:
+            elapsed = time.monotonic() - start_time
+            if elapsed >= self._flush_interval:
+                break
+            try:
+                item = self._queue.get(
+                    block=True, timeout=self._flush_interval - elapsed
+                )
+                items.append(item)
+            except Empty:
+                break
+        return items
+
+    def _upload_batch(self, batch: list[Union[Trace, Span]]):
+        self._log.debug("uploading batch of %d items", len(batch))
+
+        @backoff.on_exception(backoff.expo, Exception, max_tries=5)
+        def execute_task_with_backoff(batch: list[Union[Trace, Span]]):
+            try:
+                self._client.batch_post_traces(batch=batch)
+            except Exception as e:
+                if (
+                    isinstance(e, APIError)
+                    and 400 <= int(e.status) < 500
+                    and int(e.status) != 429
+                ):
+                    self._log.warn(
+                        f"Received {e.status} error by Laminar server, not retrying: {e.message}"
+                    )
+                    return
+
+                raise e
+
+        execute_task_with_backoff(batch)
+        self._log.debug("successfully uploaded batch of %d items", len(batch))
+
+
+class ThreadManager:
+    _log = logging.getLogger("laminar.task_manager")
+    _queue: Queue[Union[Span, Trace]]
+    _client: Laminar
+    _max_task_queue_size: int
+    _flush_interval: float
+    _collectors: list[Collector] = []
+    _threads: int
+
+    def __init__(
+        self,
+        client: Laminar,
+        flush_interval: float = 2.0,
+        max_task_queue_size: int = 1000,
+        threads: int = 1,
+    ):
+        self._max_task_queue_size = max_task_queue_size
+        self._threads = threads
+        self._queue = Queue(maxsize=self._max_task_queue_size)
+        self._flush_interval = flush_interval
+        self._client = client
+        for _ in range(self._threads):
+            collector = Collector(
+                queue=self._queue,
+                client=self._client,
+                flush_interval=flush_interval,
+            )
+            self._collectors.append(collector)
+            collector.start()
+        atexit.register(self.join)
+
+    def add_task(self, event: Union[Span, Trace]) -> bool:
+        try:
+            self._queue.put(event, block=False)
+            return True
+        except Full:
+            self._log.warning("queue is full")
+            return False
+        except Exception as e:
+            self._log.exception(f"Exception in adding task {e}")
+
+            return False
+
+    def flush(self):
+        """Forces a flush from the internal queue to the server"""
+        self._log.debug("flushing queue")
+        queue = self._queue
+        size = queue.qsize()
+        queue.join()
+        # Note that this message may not be precise, because of threading.
+        self._log.debug("successfully flushed about %s items.", size)
+
+    def join(self):
+        """Ends the collector threads once the queue is empty.
+        Blocks execution until finished
+        """
+        self._log.debug(f"joining {len(self._collectors)} collector threads")
+
+        # pause all collectors before joining them so we don't have to wait for multiple
+        # flush intervals to join them all.
+        for collector in self._collectors:
+            collector.pause()
+
+        for i, collector in enumerate(self._collectors):
+            try:
+                collector.join()
+            except RuntimeError:
+                # collector thread has not started
+                pass
+
+            self._log.debug(f"collector thread {i} joined")
+
+    def shutdown(self):
+        """Flush all messages and cleanly shutdown the client"""
+        self._log.debug("shutdown initiated")
+
+        self.flush()
+        self.join()
+
+        self._log.debug("shutdown completed")

lmnr/sdk/constants.py

@@ -0,0 +1 @@
+CURRENT_TRACING_VERSION = "0.1.0"