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

lmnr/sdk/providers/fallback.py

@@ -0,0 +1,131 @@
+from .base import Provider
+from .utils import parse_or_dump_to_dict
+
+from collections import defaultdict
+from typing import Any, Optional, Union
+import logging
+import pydantic
+
+
+class FallbackProvider(Provider):
+    logger = logging.getLogger("lmnr.sdk.tracing.providers.fallback")
+
+    def display_name(self) -> str:
+        return ""
+
+    def stream_list_to_dict(
+        self,
+        response: list[Union[dict[str, Any], str, pydantic.BaseModel]],
+    ) -> dict[str, Any]:
+        model = None
+        output_tokens = 0
+        outputs = defaultdict(lambda: defaultdict(str))
+        try:
+            for chunk in response:
+                chunk = parse_or_dump_to_dict(chunk)
+                if model is None:
+                    model = chunk.get("model")
+                for i, choice in enumerate(chunk.get("choices", [])):
+                    # TODO: not only content, but also tool_calls and refusal
+                    # also handle roles
+                    if choice["delta"] and isinstance(choice["delta"], dict):
+                        for key in choice["delta"]:
+                            if choice["delta"][key] is None:
+                                if key not in outputs[i]:
+                                    outputs[i][key] = None
+                                continue
+                            outputs[i][key] += choice["delta"][key]
+                        output_tokens += 1
+        except Exception as e:
+            self.logger.error(f"Error parsing streamming response: {e}")
+            pass
+
+        output_key_values = [
+            self._message_to_key_and_output(dict(outputs[i]))
+            for i in range(len(outputs))
+        ]
+        return {
+            "model": model,
+            "prompt_tokens": 0,
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": output_tokens,
+                "total_tokens": output_tokens,
+            },
+            "choices": [
+                (
+                    {
+                        "message": {
+                            output[0]: output[1],
+                            "role": "assistant",
+                        }
+                    }
+                    if output
+                    else None
+                )
+                for output in output_key_values
+            ],
+        }
+
+    def extract_llm_attributes_from_response(
+        self, response: Union[str, dict[str, Any], pydantic.BaseModel]
+    ) -> dict[str, Any]:
+        obj = parse_or_dump_to_dict(response)
+
+        choices = obj.get("choices", [])
+        decisions = []
+        for choice in choices:
+            if choice.get("content"):
+                decisions.append("completion")
+            elif choice.get("refusal"):
+                decisions.append("refusal")
+            elif choice.get("tool_calls"):
+                decisions.append("tool_calls")
+            else:
+                decisions.append(None)
+
+        return {
+            "response_model": obj.get("model"),
+            "input_token_count": obj.get("usage", {}).get("prompt_tokens"),
+            "output_token_count": obj.get("usage", {}).get("completion_tokens"),
+            "total_token_count": obj.get("usage", {}).get("total_tokens"),
+            "decision": self._from_singleton_list(decisions),
+        }
+
+    def extract_llm_output(
+        self, result: Union[str, dict[str, Any], pydantic.BaseModel]
+    ) -> Any:
+        result = parse_or_dump_to_dict(result)
+        choices = result.get("choices")
+        if not choices:
+            return None
+        outputs = [choice.get("message") for choice in choices]
+
+        return self._from_singleton_list(outputs)
+
+    def extract_llm_attributes_from_args(
+        self, func_args: list[Any], func_kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        return {
+            "request_model": func_kwargs.get("model"),
+            "temperature": func_kwargs.get("temperature"),
+            "stream": func_kwargs.get("stream", False),
+        }
+
+    def _message_to_key_and_output(
+        self, message: Union[dict[str, Any], pydantic.BaseModel]
+    ) -> Optional[tuple[str, str]]:
+        message = parse_or_dump_to_dict(message)
+
+        for key in ["content", "refusal", "tool_calls"]:
+            if message.get(key) is not None:
+                return (key, message[key])
+        return None
+
+    def _from_singleton_list(self, obj: Any) -> Any:
+        # OpenAI returns list of choices. This will have more than one item
+        # only if request parameter `n` is specified and is grate than 1.
+        # That's a rare case, so we return the [contents of the] choice alone if there is just one.
+        if isinstance(obj, list) and len(obj) == 1:
+            return obj[0]
+        return obj

lmnr/sdk/providers/openai.py

@@ -0,0 +1,140 @@
+from .base import Provider
+from .utils import parse_or_dump_to_dict
+
+from collections import defaultdict
+from typing import Any, Optional, Union
+import logging
+import pydantic
+
+from openai.types.chat.chat_completion import ChatCompletion
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+
+
+class OpenAI(Provider):
+    logger = logging.getLogger("lmnr.sdk.tracing.providers.openai")
+
+    def display_name(self) -> str:
+        return "OpenAI"
+
+    def stream_list_to_dict(
+        self, response: list[Union[ChatCompletionChunk, str]]
+    ) -> dict[str, Any]:
+        model = None
+        finish_reasons = []
+        output_tokens = 0
+        outputs = defaultdict(lambda: defaultdict(str))
+        try:
+            for chunk in response:
+                chunk = parse_or_dump_to_dict(chunk)
+                if model is None:
+                    model = chunk["model"]
+                finish_reasons = [
+                    choice.get("finish_reason") for choice in chunk.get("choices", [])
+                ]
+                for i, choice in enumerate(chunk.get("choices", [])):
+                    if choice["delta"] and isinstance(choice["delta"], dict):
+                        for key in choice["delta"]:
+                            if choice["delta"][key] is None:
+                                if key not in outputs[i]:
+                                    outputs[i][key] = None
+                                continue
+                            outputs[i][key] += choice["delta"][key]
+                        output_tokens += 1
+        except Exception as e:
+            self.logger.error(f"Error parsing streamming response: {e}")
+            pass
+
+        output_key_values = [
+            self._message_to_key_and_output(dict(outputs[i]))
+            for i in range(len(outputs))
+        ]
+        return {
+            "model": model,
+            "prompt_tokens": 0,
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": output_tokens,
+                "total_tokens": output_tokens,
+                "finish_reason": self._from_singleton_list(finish_reasons),
+            },
+            "choices": [
+                (
+                    {
+                        "message": {
+                            output[0]: output[1],
+                            "role": "assistant",
+                        }
+                    }
+                    if output
+                    else None
+                )
+                for output in output_key_values
+            ],
+        }
+
+    def extract_llm_attributes_from_response(
+        self, response: Union[str, dict[str, Any], pydantic.BaseModel]
+    ) -> dict[str, Any]:
+        obj = parse_or_dump_to_dict(response)
+
+        choices = obj.get("choices", [])
+        decisions = []
+        for choice in choices:
+            # choice = parse_or_dump_to_dict(choice)
+            if choice.get("content"):
+                decisions.append("completion")
+            elif choice.get("refusal"):
+                decisions.append("refusal")
+            elif choice.get("tool_calls"):
+                decisions.append("tool_calls")
+            else:
+                decisions.append(None)
+
+        return {
+            "response_model": obj.get("model"),
+            "input_token_count": obj.get("usage", {}).get("prompt_tokens"),
+            "output_token_count": obj.get("usage", {}).get("completion_tokens"),
+            "total_token_count": obj.get("usage", {}).get("total_tokens"),
+            "finish_reason": obj.get("finish_reason"),
+            "decision": self._from_singleton_list(decisions),
+        }
+
+    def extract_llm_output(
+        self, result: Union[str, dict[str, Any], ChatCompletion]
+    ) -> Any:
+        result = parse_or_dump_to_dict(result)
+        choices = result.get("choices")
+        if not choices:
+            return None
+        outputs = [choice["message"] for choice in choices]
+
+        return self._from_singleton_list(outputs)
+
+    def extract_llm_attributes_from_args(
+        self, func_args: list[Any], func_kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        return {
+            "request_model": func_kwargs.get("model"),
+            "temperature": func_kwargs.get("temperature"),
+            "top_p": func_kwargs.get("top_p"),
+            "stream": func_kwargs.get("stream", False),
+        }
+
+    def _message_to_key_and_output(
+        self, message: Union[dict[str, Any], ChatCompletionMessage]
+    ) -> Optional[tuple[str, str]]:
+        message = parse_or_dump_to_dict(message)
+
+        for key in ["content", "refusal", "tool_calls"]:
+            if message.get(key) is not None:
+                return (key, message[key])
+        return None
+
+    def _from_singleton_list(self, obj: Any) -> Any:
+        # OpenAI returns list of choices. This will have more than one item
+        # only if request parameter `n` is specified and is grate than 1.
+        # That's a rare case, so we return the [contents of the] choice alone if there is just one.
+        if isinstance(obj, list) and len(obj) == 1:
+            return obj[0]
+        return obj

lmnr/sdk/providers/utils.py

@@ -0,0 +1,33 @@
+import logging
+import json
+import pydantic
+import typing
+
+logger = logging.getLogger("lmnr.sdk.tracing.providers.utils")
+
+
+def parse_or_dump_to_dict(
+    obj: typing.Union[pydantic.BaseModel, dict[str, typing.Any], str]
+) -> dict[str, typing.Any]:
+    if isinstance(obj, pydantic.BaseModel):
+        return obj.model_dump()
+    if isinstance(obj, str):
+        try:
+            parsed = json.loads(obj)
+            if isinstance(parsed, dict):
+                return parsed
+            else:
+                logging.warning(
+                    f"Expected a dict, but got: {type(parsed)}. Returning empty dict."
+                )
+                return {}
+        except Exception as e:
+            logging.error(f"Error parsing string: {e}")
+            return {}
+
+    if isinstance(obj, dict):
+        return obj
+    logging.warning(
+        f"Expected a dict, BaseModel, or str, but got {type(obj)}. Returning empty dict."
+    )
+    return {}

lmnr/sdk/tracing_types.py

@@ -0,0 +1,197 @@
+from typing import Any, Literal, Optional, Union
+import datetime
+import pydantic
+import uuid
+
+from .constants import CURRENT_TRACING_VERSION
+from .utils import to_dict
+
+
+class EvaluateEvent(pydantic.BaseModel):
+    name: str
+    data: str
+
+
+class Span(pydantic.BaseModel):
+    version: str = CURRENT_TRACING_VERSION
+    spanType: Literal["DEFAULT", "LLM"] = "DEFAULT"
+    id: uuid.UUID
+    traceId: uuid.UUID
+    parentSpanId: Optional[uuid.UUID] = None
+    name: str
+    # generated at start of span, so required
+    startTime: datetime.datetime
+    # generated at end of span, optional when span is still running
+    endTime: Optional[datetime.datetime] = None
+    attributes: dict[str, Any] = {}
+    input: Optional[Any] = None
+    output: Optional[Any] = None
+    metadata: Optional[dict[str, Any]] = None
+    evaluateEvents: list[EvaluateEvent] = []
+    events: list["Event"] = None
+
+    def __init__(
+        self,
+        name: str,
+        trace_id: uuid.UUID,
+        start_time: Optional[datetime.datetime] = None,
+        version: str = CURRENT_TRACING_VERSION,
+        span_type: Literal["DEFAULT", "LLM"] = "DEFAULT",
+        id: Optional[uuid.UUID] = None,
+        parent_span_id: Optional[uuid.UUID] = None,
+        input: Optional[Any] = None,
+        metadata: Optional[dict[str, Any]] = {},
+        attributes: Optional[dict[str, Any]] = {},
+        evaluate_events: list[EvaluateEvent] = [],
+    ):
+        super().__init__(
+            version=version,
+            spanType=span_type,
+            id=id or uuid.uuid4(),
+            traceId=trace_id,
+            parentSpanId=parent_span_id,
+            name=name,
+            startTime=start_time or datetime.datetime.now(datetime.timezone.utc),
+            input=input,
+            metadata=metadata or {},
+            attributes=attributes or {},
+            evaluateEvents=evaluate_events,
+            events=[],
+        )
+
+    def update(
+        self,
+        end_time: Optional[datetime.datetime],
+        output: Optional[Any] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        attributes: Optional[dict[str, Any]] = None,
+        evaluate_events: Optional[list[EvaluateEvent]] = None,
+        override: bool = False,
+    ):
+        self.endTime = end_time or datetime.datetime.now(datetime.timezone.utc)
+        self.output = output
+        new_metadata = (
+            metadata if override else {**(self.metadata or {}), **(metadata or {})}
+        )
+        new_attributes = (
+            attributes or {}
+            if override
+            else {**(self.attributes or {}), **(attributes or {})}
+        )
+        new_evaluate_events = (
+            evaluate_events or []
+            if override
+            else self.evaluateEvents + (evaluate_events or [])
+        )
+        self.metadata = new_metadata
+        self.attributes = new_attributes
+        self.evaluateEvents = new_evaluate_events
+
+    def add_event(self, event: "Event"):
+        self.events.append(event)
+
+    def to_dict(self) -> dict[str, Any]:
+        try:
+            obj = self.model_dump()
+        except TypeError:
+            # if inner values are pydantic models, we need to call model_dump on them
+            # see: https://github.com/pydantic/pydantic/issues/7713
+            obj = {}
+            for key, value in self.__dict__.items():
+                obj[key] = (
+                    value.model_dump()
+                    if isinstance(value, pydantic.BaseModel)
+                    else value
+                )
+        obj = to_dict(obj)
+        return obj
+
+
+class Trace(pydantic.BaseModel):
+    id: uuid.UUID
+    version: str = CURRENT_TRACING_VERSION
+    success: bool = True
+    startTime: Optional[datetime.datetime] = None
+    endTime: Optional[datetime.datetime] = None
+    userId: Optional[str] = None  # provided by user or null
+    sessionId: Optional[str] = None  # provided by user or uuid()
+    release: Optional[str] = None
+    metadata: Optional[dict[str, Any]] = None
+
+    def __init__(
+        self,
+        success: bool = True,
+        start_time: Optional[datetime.datetime] = None,
+        end_time: Optional[datetime.datetime] = None,
+        id: Optional[uuid.UUID] = None,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        release: Optional[str] = None,
+        metadata: Optional[dict[str, Any]] = None,
+    ):
+        id_ = id or uuid.uuid4()
+        super().__init__(
+            id=id_,
+            startTime=start_time,
+            success=success,
+            endTime=end_time,
+            userId=user_id,
+            sessionId=session_id,
+            release=release,
+            metadata=metadata or {},
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        try:
+            obj = self.model_dump()
+        except TypeError:
+            # if inner values are pydantic models, we need to call model_dump on them
+            # see: https://github.com/pydantic/pydantic/issues/7713
+            obj = {}
+            for key, value in self.__dict__.items():
+                obj[key] = (
+                    value.model_dump()
+                    if isinstance(value, pydantic.BaseModel)
+                    else value
+                )
+        obj = to_dict(obj)
+        return obj
+
+
+class Event(pydantic.BaseModel):
+    id: uuid.UUID
+    templateName: str
+    timestamp: datetime.datetime
+    spanId: uuid.UUID
+    value: Optional[Union[int, str]] = None
+
+    def __init__(
+        self,
+        name: str,
+        span_id: uuid.UUID,
+        timestamp: Optional[datetime.datetime] = None,
+        value: Optional[Union[int, str]] = None,
+    ):
+        super().__init__(
+            id=uuid.uuid4(),
+            templateName=name,
+            spanId=span_id,
+            timestamp=timestamp or datetime.datetime.now(datetime.timezone.utc),
+            value=value,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        try:
+            obj = self.model_dump()
+        except TypeError:
+            # if inner values are pydantic models, we need to call model_dump on them
+            # see: https://github.com/pydantic/pydantic/issues/7713
+            obj = {}
+            for key, value in self.__dict__.items():
+                obj[key] = (
+                    value.model_dump()
+                    if isinstance(value, pydantic.BaseModel)
+                    else value
+                )
+        obj = to_dict(obj)
+        return obj

lmnr/sdk/types.py

@@ -0,0 +1,69 @@
+import requests
+import pydantic
+import uuid
+from typing import Optional, Union
+
+
+class ChatMessage(pydantic.BaseModel):
+    role: str
+    content: str
+
+
+class ConditionedValue(pydantic.BaseModel):
+    condition: str
+    value: "NodeInput"
+
+
+NodeInput = Union[str, list[ChatMessage], ConditionedValue]  # TypeAlias
+
+
+class PipelineRunRequest(pydantic.BaseModel):
+    inputs: dict[str, NodeInput]
+    pipeline: str
+    env: dict[str, str] = pydantic.Field(default_factory=dict)
+    metadata: dict[str, str] = pydantic.Field(default_factory=dict)
+    stream: bool = pydantic.Field(default=False)
+    parent_span_id: Optional[uuid.UUID] = pydantic.Field(default=None)
+    trace_id: Optional[uuid.UUID] = pydantic.Field(default=None)
+
+    # uuid is not serializable by default, so we need to convert it to a string
+    def to_dict(self):
+        return {
+            "inputs": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else v
+                for k, v in self.inputs.items()
+            },
+            "pipeline": self.pipeline,
+            "env": self.env,
+            "metadata": self.metadata,
+            "stream": self.stream,
+            "parentSpanId": str(self.parent_span_id) if self.parent_span_id else None,
+            "traceId": str(self.trace_id) if self.trace_id else None,
+        }
+
+
+class PipelineRunResponse(pydantic.BaseModel):
+    outputs: dict[str, dict[str, NodeInput]]
+    run_id: str
+
+
+class PipelineRunError(Exception):
+    error_code: str
+    error_message: str
+
+    def __init__(self, response: requests.Response):
+        try:
+            resp_json = response.json()
+            self.error_code = resp_json["error_code"]
+            self.error_message = resp_json["error_message"]
+            super().__init__(self.error_message)
+        except Exception:
+            super().__init__(response.text)
+
+    def __str__(self) -> str:
+        try:
+            return str(
+                {"error_code": self.error_code, "error_message": self.error_message}
+            )
+        except Exception:
+            return super().__str__()

lmnr/sdk/utils.py

@@ -0,0 +1,102 @@
+import asyncio
+import copy
+import datetime
+import dataclasses
+import enum
+import inspect
+import pydantic
+import queue
+import typing
+import uuid
+
+from .providers import Provider, OpenAI
+
+
+def is_method(func: typing.Callable) -> bool:
+    # inspect.ismethod is True for bound methods only, but in the decorator,
+    # the method is not bound yet, so we need to check if the first parameter
+    # is either 'self' or 'cls'. This only relies on naming conventions
+
+    # `signature._parameters` is an OrderedDict,
+    # so the order of insertion is preserved
+    params = list(inspect.signature(func).parameters)
+    return len(params) > 0 and params[0] in ["self", "cls"]
+
+
+def is_async(func: typing.Callable) -> bool:
+    # `__wrapped__` is set automatically by `functools.wraps` and `functools.update_wrapper`
+    # so we can use it to get the original function
+    while hasattr(func, "__wrapped__"):
+        func = func.__wrapped__
+
+    # Check if the function is asynchronous
+    if asyncio.iscoroutinefunction(func):
+        return True
+
+    # Fallback: check if the function's code object contains 'async'. This is for
+    # cases when the decorator did not properly use `functools.wraps` or `functools.update_wrapper`
+    CO_COROUTINE = inspect.CO_COROUTINE
+    return (func.__code__.co_flags & CO_COROUTINE) != 0
+
+
+def is_async_iterator(o: typing.Any) -> bool:
+    return hasattr(o, "__aiter__") and hasattr(o, "__anext__")
+
+
+def is_iterator(o: typing.Any) -> bool:
+    return hasattr(o, "__iter__") and hasattr(o, "__next__")
+
+
+def to_dict(obj: typing.Any) -> dict[str, typing.Any]:
+    def to_dict_inner(o: typing.Any):
+        if isinstance(o, (datetime.datetime, datetime.date)):
+            return o.strftime("%Y-%m-%dT%H:%M:%S.%f%z")
+        elif o is None:
+            return None
+        elif isinstance(o, (int, float, str, bool)):
+            return o
+        elif isinstance(o, uuid.UUID):
+            return str(o)  # same as in return, but explicit
+        elif isinstance(o, enum.Enum):
+            return o.value
+        elif dataclasses.is_dataclass(o):
+            return dataclasses.asdict(o)
+        elif isinstance(o, bytes):
+            return o.decode("utf-8")
+        elif isinstance(o, pydantic.BaseModel):
+            return o.model_dump()
+        elif isinstance(o, (tuple, set, frozenset)):
+            return [to_dict_inner(item) for item in o]
+        elif isinstance(o, list):
+            return [to_dict_inner(item) for item in o]
+        elif isinstance(o, dict):
+            return {to_dict_inner(k): to_dict_inner(v) for k, v in o.items()}
+        elif isinstance(o, queue.Queue):
+            return type(o).__name__
+
+        return str(o)
+
+    for key in obj.keys():
+        obj[key] = to_dict_inner(obj[key])
+
+    return obj
+
+
+def get_input_from_func_args(
+    func: typing.Callable,
+    is_method: bool = False,
+    func_args: list[typing.Any] = [],
+    func_kwargs: dict[str, typing.Any] = {},
+) -> dict[str, typing.Any]:
+    # Remove implicitly passed "self" or "cls" argument for instance or class methods
+    res = copy.deepcopy(func_kwargs)
+    for i, k in enumerate(inspect.signature(func).parameters.keys()):
+        if is_method and k in ["self", "cls"]:
+            continue
+        res[k] = func_args[i]
+    return res
+
+
+PROVIDER_NAME_TO_OBJECT: dict[str, Provider] = {
+    "openai": OpenAI(),
+}