lmnr 0.4.37__py3-none-any.whl → 0.4.45__py3-none-any.whl

lmnr/sdk/laminar.py

@@ -16,8 +16,10 @@ from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExport
 from opentelemetry.util.types import AttributeValue
 
 from pydantic.alias_generators import to_snake
-from typing import Any, Literal, Optional, Set, Union
+from typing import Any, Awaitable, Literal, Optional, Set, Union
 
+import aiohttp
+import asyncio
 import copy
 import datetime
 import dotenv
@@ -25,8 +27,8 @@ import json
 import logging
 import os
 import random
-import re
 import requests
+import re
 import urllib.parse
 import uuid
 
@@ -36,7 +38,6 @@ from lmnr.openllmetry_sdk.tracing.attributes import (
     SPAN_OUTPUT,
     SPAN_PATH,
     TRACE_TYPE,
-    USER_ID,
 )
 from lmnr.openllmetry_sdk.tracing.tracing import (
     get_span_path,
@@ -55,7 +56,10 @@ from .types import (
     PipelineRunResponse,
     NodeInput,
     PipelineRunRequest,
+    SemanticSearchRequest,
+    SemanticSearchResponse,
     TraceType,
+    TracingLevel,
 )
 
 
@@ -65,7 +69,6 @@ class Laminar:
     __project_api_key: Optional[str] = None
     __env: dict[str, str] = {}
     __initialized: bool = False
-    __http_session: Optional[requests.Session] = None
 
     @classmethod
     def initialize(
@@ -130,7 +133,6 @@ class Laminar:
         cls.__env = env
         cls.__initialized = True
         cls._initialize_logger()
-        cls.__http_session = requests.Session()
         Traceloop.init(
             exporter=OTLPSpanExporter(
                 endpoint=cls.__base_grpc_url,
@@ -165,8 +167,9 @@ class Laminar:
         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
+    ) -> Union[PipelineRunResponse, Awaitable[PipelineRunResponse]]:
+        """Runs the pipeline with the given inputs. If called from an async
+        function, must be awaited.
 
         Args:
             pipeline (str): name of the Laminar pipeline.\
@@ -216,34 +219,47 @@ class Laminar:
                 parent_span_id=parent_span_id,
                 trace_id=trace_id,
             )
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                return cls.__run(request)
+            else:
+                return asyncio.run(cls.__run(request))
         except Exception as e:
             raise ValueError(f"Invalid request: {e}")
 
-        response = (
-            cls.__http_session.post(
-                cls.__base_http_url + "/v1/pipeline/run",
-                data=json.dumps(request.to_dict()),
-                headers=cls._headers(),
-            )
-            if cls.__http_session
-            else requests.post(
-                cls.__base_http_url + "/v1/pipeline/run",
-                data=json.dumps(request.to_dict()),
-                headers=cls._headers(),
-            )
+    @classmethod
+    def semantic_search(
+        cls,
+        query: str,
+        dataset_id: uuid.UUID,
+        limit: Optional[int] = None,
+        threshold: Optional[float] = None,
+    ) -> SemanticSearchResponse:
+        """Perform a semantic search on a dataset. If called from an async
+        function, must be awaited.
+
+        Args:
+            query (str): query string to search by
+            dataset_id (uuid.UUID): id of the dataset to search in
+            limit (Optional[int], optional): maximum number of results to\
+                return. Defaults to None.
+            threshold (Optional[float], optional): minimum score for a result\
+                to be returned. Defaults to None.
+
+        Returns:
+            SemanticSearchResponse: response object containing the search results sorted by score in descending order
+        """
+        request = SemanticSearchRequest(
+            query=query,
+            dataset_id=dataset_id,
+            limit=limit,
+            threshold=threshold,
         )
-        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)
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            return cls.__semantic_search(request)
+        else:
+            return asyncio.run(cls.__semantic_search(request))
 
     @classmethod
     def event(
@@ -330,6 +346,10 @@ class Laminar:
                 span. Defaults to None.
         """
 
+        if not cls.is_initialized():
+            yield
+            return
+
         with get_tracer() as tracer:
             span_path = get_span_path(name)
             ctx = set_value("span_path", span_path, context)
@@ -539,6 +559,39 @@ class Laminar:
         if output is not None and span != trace.INVALID_SPAN:
             span.set_attribute(SPAN_OUTPUT, json_dumps(output))
 
+    @classmethod
+    @contextmanager
+    def set_tracing_level(self, level: TracingLevel):
+        """Set the tracing level for the current span and the context
+        (i.e. any children spans created from the current span in the current
+        thread).
+
+        Tracing level can be one of:
+        - `TracingLevel.ALL`: Enable tracing for the current span and all
+            children spans.
+        - `TracingLevel.META_ONLY`: Enable tracing for the current span and all
+            children spans, but only record metadata, e.g. tokens, costs.
+        - `TracingLevel.OFF`: Disable recording any spans.
+
+        Example:
+        ```python
+        from lmnr import Laminar, TracingLevel
+
+        with Laminar.set_tracing_level(TracingLevel.META_ONLY):
+            openai_client.chat.completions.create()
+        ```
+        """
+        if level == TracingLevel.ALL:
+            yield
+        else:
+            level = "meta_only" if level == TracingLevel.META_ONLY else "off"
+            update_association_properties({"tracing_level": level})
+            yield
+            try:
+                remove_association_properties({"tracing_level": level})
+            except Exception:
+                pass
+
     @classmethod
     def set_span_attributes(
         cls,
@@ -588,7 +641,6 @@ class Laminar:
     def set_session(
         cls,
         session_id: Optional[str] = None,
-        user_id: Optional[str] = None,
     ):
         """Set the session and user id for the current span and the context
         (i.e. any children spans created from the current span in the current
@@ -599,29 +651,18 @@ class Laminar:
                             Useful to debug and group long-running\
                             sessions/conversations.
                             Defaults to None.
-            user_id (Optional[str], optional). Deprecated.\
-                            Use `Laminar.set_metadata` instead.\
-                            Custom user id.\
-                            Useful for grouping spans or traces by user.\
-                            Defaults to None.
         """
         association_properties = {}
         if session_id is not None:
             association_properties[SESSION_ID] = session_id
-        if user_id is not None:
-            cls.__logger.warning(
-                "User ID in set_session is deprecated and will be removed soon. "
-                "Please use `Laminar.set_metadata` instead."
-            )
-            association_properties["metadata." + USER_ID] = user_id
         update_association_properties(association_properties)
 
     @classmethod
-    def set_metadata(cls, metadata: dict[str, Any]):
+    def set_metadata(cls, metadata: dict[str, str]):
         """Set the metadata for the current trace.
 
         Args:
-            metadata (dict[str, Any]): Metadata to set for the trace. Willl be\
+            metadata (dict[str, str]): Metadata to set for the trace. Willl be\
                 sent as attributes, so must be json serializable.
         """
         props = {f"metadata.{k}": json_dumps(v) for k, v in metadata.items()}
@@ -636,20 +677,6 @@ class Laminar:
                 props.pop(k)
         set_association_properties(props)
 
-    @classmethod
-    def _set_trace_type(
-        cls,
-        trace_type: TraceType,
-    ):
-        """Set the trace_type for the current span and the context
-        Args:
-            trace_type (TraceType): Type of the trace
-        """
-        association_properties = {
-            TRACE_TYPE: trace_type.value,
-        }
-        update_association_properties(association_properties)
-
     @classmethod
     def clear_session(cls):
         """Clear the session and user id from  the context"""
@@ -659,30 +686,33 @@ class Laminar:
         set_association_properties(props)
 
     @classmethod
-    def create_evaluation(
+    async def create_evaluation(
         cls,
         data: list[EvaluationResultDatapoint],
         group_id: Optional[str] = None,
         name: Optional[str] = None,
     ) -> CreateEvaluationResponse:
-        response = requests.post(
-            cls.__base_http_url + "/v1/evaluations",
-            data=json.dumps(
-                {
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                cls.__base_http_url + "/v1/evaluations",
+                json={
                     "groupId": group_id,
                     "name": name,
                     "points": [datapoint.to_dict() for datapoint in data],
-                }
-            ),
-            headers=cls._headers(),
-        )
-        if response.status_code != 200:
-            try:
-                resp_json = response.json()
-                raise ValueError(f"Error creating evaluation {json.dumps(resp_json)}")
-            except requests.exceptions.RequestException:
-                raise ValueError(f"Error creating evaluation {response.text}")
-        return CreateEvaluationResponse.model_validate(response.json())
+                },
+                headers=cls._headers(),
+            ) as response:
+                if response.status != 200:
+                    try:
+                        resp_json = await response.json()
+                        raise ValueError(
+                            f"Error creating evaluation {json.dumps(resp_json)}"
+                        )
+                    except aiohttp.ClientError:
+                        text = await response.text()
+                        raise ValueError(f"Error creating evaluation {text}")
+                resp_json = await response.json()
+                return CreateEvaluationResponse.model_validate(resp_json)
 
     @classmethod
     def get_datapoints(
@@ -691,6 +721,10 @@ class Laminar:
         offset: int,
         limit: int,
     ) -> GetDatapointsResponse:
+        # TODO: Use aiohttp. Currently, this function is called from within
+        # `LaminarDataset.__len__`, which is sync, but can be called from
+        # both sync and async. Python does not make it easy to mix things this
+        # way, so we should probably refactor `LaminarDataset`.
         params = {"name": dataset_name, "offset": offset, "limit": limit}
         url = (
             cls.__base_http_url
@@ -717,3 +751,66 @@ class Laminar:
             "Authorization": "Bearer " + cls.__project_api_key,
             "Content-Type": "application/json",
         }
+
+    @classmethod
+    def _set_trace_type(
+        cls,
+        trace_type: TraceType,
+    ):
+        """Set the trace_type for the current span and the context
+        Args:
+            trace_type (TraceType): Type of the trace
+        """
+        association_properties = {
+            TRACE_TYPE: trace_type.value,
+        }
+        update_association_properties(association_properties)
+
+    @classmethod
+    async def __run(
+        cls,
+        request: PipelineRunRequest,
+    ) -> PipelineRunResponse:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                cls.__base_http_url + "/v1/pipeline/run",
+                data=json.dumps(request.to_dict()),
+                headers=cls._headers(),
+            ) as response:
+                if response.status != 200:
+                    raise PipelineRunError(response)
+                try:
+                    resp_json = await 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)
+
+    @classmethod
+    async def __semantic_search(
+        cls,
+        request: SemanticSearchRequest,
+    ) -> SemanticSearchResponse:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                cls.__base_http_url + "/v1/semantic-search",
+                data=json.dumps(request.to_dict()),
+                headers=cls._headers(),
+            ) as response:
+                if response.status != 200:
+                    raise ValueError(
+                        f"Error performing semantic search: [{response.status}] {response.text}"
+                    )
+                try:
+                    resp_json = await 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} error={e}"
+                    )

lmnr/sdk/types.py

@@ -1,7 +1,7 @@
+import aiohttp
 import datetime
 from enum import Enum
 import pydantic
-import requests
 from typing import Any, Awaitable, Callable, Optional, Union
 import uuid
 
@@ -55,11 +55,40 @@ class PipelineRunResponse(pydantic.BaseModel):
     run_id: str
 
 
+class SemanticSearchRequest(pydantic.BaseModel):
+    query: str
+    dataset_id: uuid.UUID
+    limit: Optional[int] = pydantic.Field(default=None)
+    threshold: Optional[float] = pydantic.Field(default=None, ge=0.0, le=1.0)
+
+    def to_dict(self):
+        res = {
+            "query": self.query,
+            "datasetId": str(self.dataset_id),
+        }
+        if self.limit is not None:
+            res["limit"] = self.limit
+        if self.threshold is not None:
+            res["threshold"] = self.threshold
+        return res
+
+
+class SemanticSearchResult(pydantic.BaseModel):
+    dataset_id: uuid.UUID
+    score: float
+    data: dict[str, Any]
+    content: str
+
+
+class SemanticSearchResponse(pydantic.BaseModel):
+    results: list[SemanticSearchResult]
+
+
 class PipelineRunError(Exception):
     error_code: str
     error_message: str
 
-    def __init__(self, response: requests.Response):
+    def __init__(self, response: aiohttp.ClientResponse):
         try:
             resp_json = response.json()
             self.error_code = resp_json["error_code"]
@@ -95,7 +124,7 @@ ExecutorFunctionReturnType = Any
 EvaluatorFunctionReturnType = Union[Numeric, dict[str, Numeric]]
 
 ExecutorFunction = Callable[
-    [EvaluationDatapointData, Any, dict[str, Any]],
+    [EvaluationDatapointData, Any],
     Union[ExecutorFunctionReturnType, Awaitable[ExecutorFunctionReturnType]],
 ]
 
@@ -104,7 +133,7 @@ ExecutorFunction = Callable[
 # record of string keys and number values. The latter is useful for evaluating
 # multiple criteria in one go instead of running multiple evaluators.
 EvaluatorFunction = Callable[
-    [ExecutorFunctionReturnType, Any, dict[str, Any]],
+    [ExecutorFunctionReturnType, Any],
     Union[EvaluatorFunctionReturnType, Awaitable[EvaluatorFunctionReturnType]],
 ]
 
@@ -174,3 +203,9 @@ class TraceType(Enum):
 class GetDatapointsResponse(pydantic.BaseModel):
     items: list[Datapoint]
     totalCount: int
+
+
+class TracingLevel(Enum):
+    OFF = 0
+    META_ONLY = 1
+    ALL = 2