lmnr 0.4.12b2__py3-none-any.whl → 0.4.12b4__py3-none-any.whl

lmnr/sdk/decorators.py

@@ -6,6 +6,7 @@ from opentelemetry.trace import INVALID_SPAN, get_current_span
 
 from typing import Callable, Optional, cast
 
+from lmnr.traceloop_sdk.tracing.attributes import SESSION_ID, USER_ID
 from lmnr.traceloop_sdk.tracing.tracing import update_association_properties
 
 from .utils import is_async
@@ -43,11 +44,11 @@ def observe(
         if current_span != INVALID_SPAN:
             if session_id is not None:
                 current_span.set_attribute(
-                    "traceloop.association.properties.session_id", session_id
+                    SESSION_ID, session_id
                 )
             if user_id is not None:
                 current_span.set_attribute(
-                    "traceloop.association.properties.user_id", user_id
+                    USER_ID, user_id
                 )
         association_properties = {}
         if session_id is not None:

lmnr/sdk/evaluations.py

@@ -2,12 +2,26 @@ import asyncio
 import sys
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, Awaitable, Optional, Union
+from typing import Any, Awaitable, Optional, Set, Union
+import uuid
 
 from tqdm import tqdm
 
+from ..traceloop_sdk.instruments import Instruments
+from ..traceloop_sdk.tracing.attributes import SPAN_TYPE
+
 from .laminar import Laminar as L
-from .types import CreateEvaluationResponse, Datapoint, EvaluationResultDatapoint, Numeric, NumericTypes
+from .types import (
+    CreateEvaluationResponse,
+    Datapoint,
+    EvaluationResultDatapoint,
+    EvaluatorFunction,
+    ExecutorFunction,
+    Numeric,
+    NumericTypes,
+    SpanType,
+    TraceType,
+)
 from .utils import is_async
 
 DEFAULT_BATCH_SIZE = 5
@@ -39,7 +53,11 @@ class EvaluationReporter:
     def start(self, name: str, project_id: str, id: str, length: int):
         print(f"Running evaluation {name}...\n")
         print(f"Check progress and results at {get_evaluation_url(project_id, id)}\n")
-        self.cli_progress = tqdm(total=length, bar_format="{bar} {percentage:3.0f}% | ETA: {remaining}s | {n_fmt}/{total_fmt}", ncols=60)
+        self.cli_progress = tqdm(
+            total=length,
+            bar_format="{bar} {percentage:3.0f}% | ETA: {remaining}s | {n_fmt}/{total_fmt}",
+            ncols=60,
+        )
 
     def update(self, batch_length: int):
         self.cli_progress.update(batch_length)
@@ -51,7 +69,7 @@ class EvaluationReporter:
     def stop(self, average_scores: dict[str, Numeric]):
         self.cli_progress.close()
         print("\nAverage scores:")
-        for (name, score) in average_scores.items():
+        for name, score in average_scores.items():
             print(f"{name}: {score}")
         print("\n")
 
@@ -76,20 +94,21 @@ class EvaluationDataset(ABC):
 class Evaluation:
     def __init__(
         self,
-        name: str,
         data: Union[EvaluationDataset, list[Union[Datapoint, dict]]],
         executor: Any,
-        evaluators: list[Any],
+        evaluators: dict[str, EvaluatorFunction],
+        name: Optional[str] = None,
         batch_size: int = DEFAULT_BATCH_SIZE,
         project_api_key: Optional[str] = None,
         base_url: Optional[str] = None,
         http_port: Optional[int] = None,
+        grpc_port: Optional[int] = None,
+        instruments: Optional[Set[Instruments]] = None,
     ):
         """
         Initializes an instance of the Evaluations class.
 
         Parameters:
-            name (str): The name of the evaluation.
             data (Union[List[Union[EvaluationDatapoint, dict]], EvaluationDataset]): List of data points to evaluate or an evaluation dataset.
                             `data` is the input to the executor function,
                             `target` is the input to the evaluator function.
@@ -104,6 +123,8 @@ class Evaluation:
                 evaluator function. If the function is anonymous, it will be
                 named `evaluator_${index}`, where index is the index of the
                 evaluator function in the list starting from 1.
+            name (Optional[str], optional): The name of the evaluation.
+                            It will be auto-generated if not provided.
             batch_size (int, optional): The batch size for evaluation.
                             Defaults to DEFAULT_BATCH_SIZE.
             project_api_key (Optional[str], optional): The project API key.
@@ -113,33 +134,18 @@ class Evaluation:
                             Defaults to "https://api.lmnr.ai".
             http_port (Optional[int], optional): The port for the Laminar API HTTP service.
                             Defaults to 443.
+            instruments (Optional[Set[Instruments]], optional): Set of modules to auto-instrument.
+                            Defaults to None. If None, all available instruments will be used.
         """
 
         self.is_finished = False
         self.name = name
         self.reporter = EvaluationReporter()
         self.executor = executor
-        self.evaluators = dict(
-            zip(
-                [
-                    (
-                        e.__name__
-                        if e.__name__ and e.__name__ != "<lambda>"
-                        else f"evaluator_{i+1}"
-                    )
-                    for i, e in enumerate(evaluators)
-                ],
-                evaluators,
-            )
-        )
-        self.evaluator_names = list(self.evaluators.keys())
+        self.evaluators = evaluators
         if isinstance(data, list):
             self.data = [
-                (
-                    Datapoint.model_validate(point)
-                    if isinstance(point, dict)
-                    else point
-                )
+                (Datapoint.model_validate(point) if isinstance(point, dict) else point)
                 for point in data
             ]
         else:
@@ -149,7 +155,8 @@ class Evaluation:
             project_api_key=project_api_key,
             base_url=base_url,
             http_port=http_port,
-            instruments=set(),
+            grpc_port=grpc_port,
+            instruments=instruments,
         )
 
     def run(self) -> Union[None, Awaitable[None]]:
@@ -204,7 +211,7 @@ class Evaluation:
     async def evaluate_in_batches(self, evaluation: CreateEvaluationResponse):
         for i in range(0, len(self.data), self.batch_size):
             batch = (
-                self.data[i: i + self.batch_size]
+                self.data[i : i + self.batch_size]
                 if isinstance(self.data, list)
                 else self.data.slice(i, i + self.batch_size)
             )
@@ -216,71 +223,121 @@ class Evaluation:
             finally:
                 self.reporter.update(len(batch))
 
-    async def _evaluate_batch(self, batch: list[Datapoint]) -> list[EvaluationResultDatapoint]:
+    async def _evaluate_batch(
+        self, batch: list[Datapoint]
+    ) -> list[EvaluationResultDatapoint]:
         batch_promises = [self._evaluate_datapoint(datapoint) for datapoint in batch]
         results = await asyncio.gather(*batch_promises)
         return results
 
-    async def _evaluate_datapoint(self, datapoint) -> EvaluationResultDatapoint:
-        output = (
-            await self.executor(datapoint.data)
-            if is_async(self.executor)
-            else self.executor(datapoint.data)
-        )
-        target = datapoint.target
-
-        # Iterate over evaluators
-        scores: dict[str, Numeric] = {}
-        for evaluator_name in self.evaluator_names:
-            evaluator = self.evaluators[evaluator_name]
-            value = (
-                await evaluator(output, target)
-                if is_async(evaluator)
-                else evaluator(output, target)
+    async def _evaluate_datapoint(
+        self, datapoint: Datapoint
+    ) -> EvaluationResultDatapoint:
+        with L.start_as_current_span("evaluation") as evaluation_span:
+            L._set_trace_type(trace_type=TraceType.EVALUATION)
+            evaluation_span.set_attribute(SPAN_TYPE, SpanType.EVALUATION.value)
+            with L.start_as_current_span(
+                "executor", input={"data": datapoint.data}
+            ) as executor_span:
+                executor_span.set_attribute(SPAN_TYPE, SpanType.EXECUTOR.value)
+                output = (
+                    await self.executor(datapoint.data)
+                    if is_async(self.executor)
+                    else self.executor(datapoint.data)
+                )
+                L.set_span_output(output)
+            target = datapoint.target
+
+            # Iterate over evaluators
+            scores: dict[str, Numeric] = {}
+            for evaluator_name, evaluator in self.evaluators.items():
+                with L.start_as_current_span(
+                    "evaluator", input={"output": output, "target": target}
+                ) as evaluator_span:
+                    evaluator_span.set_attribute(SPAN_TYPE, SpanType.EVALUATOR.value)
+                    value = (
+                        await evaluator(output, target)
+                        if is_async(evaluator)
+                        else evaluator(output, target)
+                    )
+                    L.set_span_output(value)
+
+                # If evaluator returns a single number, use evaluator name as key
+                if isinstance(value, NumericTypes):
+                    scores[evaluator_name] = value
+                else:
+                    scores.update(value)
+
+            trace_id = uuid.UUID(int=evaluation_span.get_span_context().trace_id)
+            return EvaluationResultDatapoint(
+                data=datapoint.data,
+                target=target,
+                executor_output=output,
+                scores=scores,
+                trace_id=trace_id,
             )
 
-            # If evaluator returns a single number, use evaluator name as key
-            if isinstance(value, NumericTypes):
-                scores[evaluator_name] = value
-            else:
-                scores.update(value)
-
-        return EvaluationResultDatapoint(
-            data=datapoint.data,
-            target=target,
-            executorOutput=output,
-            scores=scores,
-        )
-
 
 def evaluate(
-    name: str,
     data: Union[EvaluationDataset, list[Union[Datapoint, dict]]],
-    executor: Any,
-    evaluators: list[Any],
+    executor: ExecutorFunction,
+    evaluators: dict[str, EvaluatorFunction],
+    name: Optional[str] = None,
     batch_size: int = DEFAULT_BATCH_SIZE,
     project_api_key: Optional[str] = None,
     base_url: Optional[str] = None,
     http_port: Optional[int] = None,
+    grpc_port: Optional[int] = None,
+    instruments: Optional[Set[Instruments]] = None,
 ) -> Optional[Awaitable[None]]:
     """
-    Run evaluation.
-
-    If `_set_global_evaluation` is `True`, sets the global evaluation to be run in another part of the program.
-
-    Otherwise, if there is no event loop, runs the evaluation in the current thread until completion.
+    If added to the file which is called through lmnr eval command, then simply registers the evaluation.
+    Otherwise, if there is no event loop, creates it and runs the evaluation until completion.
     If there is an event loop, schedules the evaluation as a task in the event loop and returns an awaitable handle.
+
+    Parameters:
+        data (Union[List[Union[EvaluationDatapoint, dict]], EvaluationDataset]): List of data points to evaluate or an evaluation dataset.
+                        `data` is the input to the executor function,
+                        `target` is the input to the evaluator function.
+        executor (Callable[..., Any]): The executor function.
+                        Takes the data point + any additional arguments
+                        and returns the output to evaluate.
+        evaluators (List[Callable[..., Any]]): List of evaluator functions.
+            Each evaluator function takes the output of the executor _and_
+            the target data, and returns a score. The score can be a
+            single number or a record of string keys and number values.
+            If the score is a single number, it will be named after the
+            evaluator function. If the function is anonymous, it will be
+            named `evaluator_${index}`, where index is the index of the
+            evaluator function in the list starting from 1.
+        name (Optional[str], optional): The name of the evaluation.
+            It will be auto-generated if not provided.
+        batch_size (int, optional): The batch size for evaluation.
+                        Defaults to DEFAULT_BATCH_SIZE.
+        project_api_key (Optional[str], optional): The project API key.
+                        Defaults to an empty string.
+        base_url (Optional[str], optional): The base URL for the Laminar API.
+                        Useful if self-hosted elsewhere.
+                        Defaults to "https://api.lmnr.ai".
+        http_port (Optional[int], optional): The port for the Laminar API HTTP service.
+                        Defaults to 443.
+        grpc_port (Optional[int], optional): The port for the Laminar API gRPC service.
+                        Defaults to 8443.
+        instruments (Optional[Set[Instruments]], optional): Set of modules to auto-instrument.
+                        Defaults to None. If None, all available instruments will be used.
     """
 
     evaluation = Evaluation(
-        name,
-        data,
-        executor,
-        evaluators,
-        batch_size,
-        project_api_key,
-        base_url,
-        http_port,
+        data=data,
+        executor=executor,
+        evaluators=evaluators,
+        name=name,
+        batch_size=batch_size,
+        project_api_key=project_api_key,
+        base_url=base_url,
+        http_port=http_port,
+        grpc_port=grpc_port,
+        instruments=instruments,
     )
 
     global _evaluation

lmnr/sdk/laminar.py

@@ -5,7 +5,6 @@ from opentelemetry.trace import (
     get_current_span,
     SpanKind,
 )
-from opentelemetry.semconv_ai import SpanAttributes
 from opentelemetry.util.types import AttributeValue
 from opentelemetry.context.context import Context
 from opentelemetry.util import types
@@ -26,7 +25,17 @@ import os
 import requests
 import uuid
 
-from lmnr.traceloop_sdk.tracing.tracing import set_association_properties, update_association_properties
+from lmnr.traceloop_sdk.tracing.attributes import (
+    SESSION_ID,
+    SPAN_INPUT,
+    SPAN_OUTPUT,
+    TRACE_TYPE,
+    USER_ID,
+)
+from lmnr.traceloop_sdk.tracing.tracing import (
+    set_association_properties,
+    update_association_properties,
+)
 
 from .log import VerboseColorfulFormatter
 
@@ -37,6 +46,7 @@ from .types import (
     PipelineRunResponse,
     NodeInput,
     PipelineRunRequest,
+    TraceType,
     UpdateEvaluationResponse,
 )
 
@@ -356,8 +366,8 @@ class Laminar:
             ) as span:
                 if input is not None:
                     span.set_attribute(
-                        SpanAttributes.TRACELOOP_ENTITY_INPUT,
-                        json.dumps({"input": input}),
+                        SPAN_INPUT,
+                        json.dumps(input),
                     )
                 yield span
 
@@ -371,9 +381,7 @@ class Laminar:
         """
         span = get_current_span()
         if output is not None and span != INVALID_SPAN:
-            span.set_attribute(
-                SpanAttributes.TRACELOOP_ENTITY_OUTPUT, json.dumps(output)
-            )
+            span.set_attribute(SPAN_OUTPUT, json.dumps(output))
 
     @classmethod
     def set_session(
@@ -396,9 +404,23 @@ class Laminar:
         """
         association_properties = {}
         if session_id is not None:
-            association_properties["session_id"] = session_id
+            association_properties[SESSION_ID] = session_id
         if user_id is not None:
-            association_properties["user_id"] = user_id
+            association_properties[USER_ID] = user_id
+        update_association_properties(association_properties)
+
+    @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
@@ -410,7 +432,7 @@ class Laminar:
         set_association_properties(props)
 
     @classmethod
-    def create_evaluation(cls, name: str) -> CreateEvaluationResponse:
+    def create_evaluation(cls, name: Optional[str]) -> CreateEvaluationResponse:
         response = requests.post(
             cls.__base_http_url + "/v1/evaluations",
             data=json.dumps({"name": name}),
@@ -430,7 +452,7 @@ class Laminar:
     ) -> requests.Response:
         body = {
             "evaluationId": str(evaluation_id),
-            "points": [datapoint.model_dump() for datapoint in data],
+            "points": [datapoint.to_dict() for datapoint in data],
         }
         response = requests.post(
             cls.__base_http_url + "/v1/evaluation-datapoints",

lmnr/sdk/types.py

@@ -1,10 +1,11 @@
 import datetime
-import requests
+from enum import Enum
 import pydantic
-import uuid
+import requests
 from typing import Any, Awaitable, Callable, Literal, Optional, Union
+import uuid
 
-from .utils import to_dict
+from .utils import serialize
 
 
 class ChatMessage(pydantic.BaseModel):
@@ -37,7 +38,7 @@ class PipelineRunRequest(pydantic.BaseModel):
     def to_dict(self):
         return {
             "inputs": {
-                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else to_dict(v)
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
                 for k, v in self.inputs.items()
             },
             "pipeline": self.pipeline,
@@ -125,5 +126,37 @@ UpdateEvaluationResponse = CreateEvaluationResponse
 class EvaluationResultDatapoint(pydantic.BaseModel):
     data: EvaluationDatapointData
     target: EvaluationDatapointTarget
-    executorOutput: ExecutorFunctionReturnType
+    executor_output: ExecutorFunctionReturnType
     scores: dict[str, Numeric]
+    trace_id: uuid.UUID
+
+    # uuid is not serializable by default, so we need to convert it to a string
+    def to_dict(self):
+        return {
+            "data": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
+                for k, v in self.data.items()
+            },
+            "target": {
+                k: v.model_dump() if isinstance(v, pydantic.BaseModel) else serialize(v)
+                for k, v in self.target.items()
+            },
+            "executorOutput": serialize(self.executor_output),
+            "scores": self.scores,
+            "traceId": str(self.trace_id),
+        }
+
+
+class SpanType(Enum):
+    DEFAULT = "DEFAULT"
+    LLM = "LLM"
+    PIPELINE = "PIPELINE"  # must not be set manually
+    EXECUTOR = "EXECUTOR"
+    EVALUATOR = "EVALUATOR"
+    EVALUATION = "EVALUATION"
+
+
+class TraceType(Enum):
+    DEFAULT = "DEFAULT"
+    EVENT = "EVENT"  # must not be set manually
+    EVALUATION = "EVALUATION"

lmnr/sdk/utils.py

@@ -1,5 +1,4 @@
 import asyncio
-import copy
 import datetime
 import dataclasses
 import enum
@@ -50,7 +49,7 @@ 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 serialize(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")
@@ -59,7 +58,7 @@ def to_dict(obj: typing.Any) -> dict[str, typing.Any]:
         elif isinstance(o, (int, float, str, bool)):
             return o
         elif isinstance(o, uuid.UUID):
-            return str(o)  # same as in return, but explicit
+            return str(o)  # same as in final return, but explicit
         elif isinstance(o, enum.Enum):
             return o.value
         elif dataclasses.is_dataclass(o):
@@ -90,11 +89,11 @@ def get_input_from_func_args(
 ) -> dict[str, typing.Any]:
     # Remove implicitly passed "self" or "cls" argument for
     # instance or class methods
-    res = copy.deepcopy(func_kwargs)
+    res = func_kwargs.copy()
     for i, k in enumerate(inspect.signature(func).parameters.keys()):
         if is_method and k in ["self", "cls"]:
             continue
         # If param has default value, then it's not present in func args
-        if len(func_args) > i:
+        if i < len(func_args):
             res[k] = func_args[i]
     return res

lmnr/traceloop_sdk/__init__.py

@@ -3,20 +3,16 @@ import sys
 from pathlib import Path
 
 from typing import Optional, Set
-from colorama import Fore
 from opentelemetry.sdk.trace import SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
-from opentelemetry.sdk.metrics.export import MetricExporter
 from opentelemetry.sdk.resources import SERVICE_NAME
 from opentelemetry.propagators.textmap import TextMapPropagator
 from opentelemetry.util.re import parse_env_headers
 
-from lmnr.traceloop_sdk.metrics.metrics import MetricsWrapper
 from lmnr.traceloop_sdk.instruments import Instruments
 from lmnr.traceloop_sdk.config import (
     is_content_tracing_enabled,
     is_tracing_enabled,
-    is_metrics_enabled,
 )
 from lmnr.traceloop_sdk.tracing.tracing import TracerWrapper
 from typing import Dict
@@ -38,8 +34,6 @@ class Traceloop:
         headers: Dict[str, str] = {},
         disable_batch=False,
         exporter: Optional[SpanExporter] = None,
-        metrics_exporter: Optional[MetricExporter] = None,
-        metrics_headers: Optional[Dict[str, str]] = None,
         processor: Optional[SpanProcessor] = None,
         propagator: Optional[TextMapPropagator] = None,
         should_enrich_metrics: bool = True,
@@ -50,7 +44,7 @@ class Traceloop:
         api_key = os.getenv("TRACELOOP_API_KEY") or api_key
 
         if not is_tracing_enabled():
-            print(Fore.YELLOW + "Tracing is disabled" + Fore.RESET)
+            # print(Fore.YELLOW + "Tracing is disabled" + Fore.RESET)
             return
 
         enable_content_tracing = is_content_tracing_enabled()
@@ -67,12 +61,10 @@ class Traceloop:
             and not api_key
         ):
             print(
-                Fore.RED
-                + "Error: Missing API key,"
+                "Error: Missing API key,"
                 + " go to project settings to create one"
             )
             print("Set the LMNR_PROJECT_API_KEY environment variable to the key")
-            print(Fore.RESET)
             return
 
         if api_key and not exporter and not processor and not headers:
@@ -80,7 +72,7 @@ class Traceloop:
                 "Authorization": f"Bearer {api_key}",
             }
 
-        print(Fore.RESET)
+        # print(Fore.RESET)
 
         # Tracer init
         resource_attributes.update({SERVICE_NAME: app_name})
@@ -95,21 +87,3 @@ class Traceloop:
             should_enrich_metrics=should_enrich_metrics,
             instruments=instruments,
         )
-
-        if not metrics_exporter and exporter:
-            return
-
-        metrics_endpoint = os.getenv("TRACELOOP_METRICS_ENDPOINT") or api_endpoint
-        metrics_headers = (
-            os.getenv("TRACELOOP_METRICS_HEADERS") or metrics_headers or headers
-        )
-
-        if not is_metrics_enabled() or not metrics_exporter and exporter:
-            print(Fore.YELLOW + "Metrics are disabled" + Fore.RESET)
-            return
-
-        MetricsWrapper.set_static_params(
-            resource_attributes, metrics_endpoint, metrics_headers
-        )
-
-        Traceloop.__metrics_wrapper = MetricsWrapper(exporter=metrics_exporter)

lmnr/traceloop_sdk/config/__init__.py

@@ -7,7 +7,3 @@ def is_tracing_enabled() -> bool:
 
 def is_content_tracing_enabled() -> bool:
     return (os.getenv("TRACELOOP_TRACE_CONTENT") or "true").lower() == "true"
-
-
-def is_metrics_enabled() -> bool:
-    return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"

lmnr/traceloop_sdk/decorators/base.py

@@ -7,9 +7,10 @@ import warnings
 
 from opentelemetry import trace
 from opentelemetry import context as context_api
-from opentelemetry.semconv_ai import SpanAttributes
 
+from lmnr.sdk.utils import get_input_from_func_args, is_method
 from lmnr.traceloop_sdk.tracing import get_tracer
+from lmnr.traceloop_sdk.tracing.attributes import SPAN_INPUT, SPAN_OUTPUT
 from lmnr.traceloop_sdk.tracing.tracing import TracerWrapper
 from lmnr.traceloop_sdk.utils.json_encoder import JSONEncoder
 
@@ -52,8 +53,12 @@ def entity_method(
                 try:
                     if _should_send_prompts():
                         span.set_attribute(
-                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
-                            _json_dumps({"args": args, "kwargs": kwargs}),
+                            SPAN_INPUT,
+                            _json_dumps(
+                                get_input_from_func_args(
+                                    fn, is_method(fn), args, kwargs
+                                )
+                            ),
                         )
                 except TypeError:
                     pass
@@ -67,7 +72,7 @@ def entity_method(
                 try:
                     if _should_send_prompts():
                         span.set_attribute(
-                            SpanAttributes.TRACELOOP_ENTITY_OUTPUT,
+                            SPAN_OUTPUT,
                             _json_dumps(res),
                         )
                 except TypeError:
@@ -105,8 +110,12 @@ def aentity_method(
                 try:
                     if _should_send_prompts():
                         span.set_attribute(
-                            SpanAttributes.TRACELOOP_ENTITY_INPUT,
-                            _json_dumps({"args": args, "kwargs": kwargs}),
+                            SPAN_INPUT,
+                            _json_dumps(
+                                get_input_from_func_args(
+                                    fn, is_method(fn), args, kwargs
+                                )
+                            ),
                         )
                 except TypeError:
                     pass
@@ -119,9 +128,7 @@ def aentity_method(
 
                 try:
                     if _should_send_prompts():
-                        span.set_attribute(
-                            SpanAttributes.TRACELOOP_ENTITY_OUTPUT, json.dumps(res)
-                        )
+                        span.set_attribute(SPAN_OUTPUT, json.dumps(res))
                 except TypeError:
                     pass