agenta 0.15.0__py3-none-any.whl → 0.15.0a0__py3-none-any.whl

agenta/sdk/tracing/llm_tracing.py

@@ -1,79 +1,59 @@
-import os
-from threading import Lock
+# Stdlib Imports
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any, List, Union
 
+# Own Imports
 from agenta.sdk.tracing.logger import llm_logger
 from agenta.sdk.tracing.tasks_manager import TaskQueue
 from agenta.client.backend.client import AsyncAgentaApi
 from agenta.client.backend.client import AsyncObservabilityClient
 from agenta.client.backend.types.create_span import CreateSpan, SpanKind, SpanStatusCode
 
+# Third Party Imports
 from bson.objectid import ObjectId
 
-VARIANT_TRACKING_FEATURE_FLAG = False
 
+class Tracing(object):
+    """Agenta llm tracing object.
 
-class SingletonMeta(type):
-    """
-    Thread-safe implementation of Singleton.
-    """
-
-    _instances = {}  # type: ignore
-
-    # We need the lock mechanism to synchronize threads \
-    # during the initial access to the Singleton object.
-    _lock: Lock = Lock()
-
-    def __call__(cls, *args, **kwargs):
-        """
-        Ensures that changes to the `__init__` arguments do not affect the
-        returned instance.
-
-        Uses a lock to make this method thread-safe. If an instance of the class
-        does not already exist, it creates one. Otherwise, it returns the
-        existing instance.
-        """
-
-        with cls._lock:
-            if cls not in cls._instances:
-                instance = super().__call__(*args, **kwargs)
-                cls._instances[cls] = instance
-        return cls._instances[cls]
-
-
-class Tracing(metaclass=SingletonMeta):
-    """The `Tracing` class is an agent for LLM tracing with specific initialization arguments.
-
-    __init__ args:
-        host (str): The URL of the backend host
+    Args:
+        base_url (str): The URL of the backend host
         api_key (str): The API Key of the backend host
         tasks_manager (TaskQueue): The tasks manager dedicated to handling asynchronous tasks
         llm_logger (Logger): The logger associated with the LLM tracing
         max_workers (int): The maximum number of workers to run tracing
     """
 
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
     def __init__(
         self,
-        host: str,
+        base_url: str,
         app_id: str,
+        variant_id: str,
+        variant_name: Optional[str] = None,
         api_key: Optional[str] = None,
         max_workers: Optional[int] = None,
     ):
-        self.host = host + "/api"
+        self.base_url = base_url + "/api"
         self.api_key = api_key if api_key is not None else ""
         self.llm_logger = llm_logger
         self.app_id = app_id
+        self.variant_id = variant_id
+        self.variant_name = variant_name
         self.tasks_manager = TaskQueue(
             max_workers if max_workers else 4, logger=llm_logger
         )
-        self.active_span: Optional[CreateSpan] = None
-        self.active_trace_id: Optional[str] = None
-        self.pending_spans: List[CreateSpan] = []
+        self.active_span = CreateSpan
+        self.active_trace = CreateSpan
+        self.recording_trace_id: Union[str, None] = None
+        self.recorded_spans: List[CreateSpan] = []
         self.tags: List[str] = []
-        self.trace_config_cache: Dict[
-            str, Any
-        ] = {}  # used to save the trace configuration before starting the first span
         self.span_dict: Dict[str, CreateSpan] = {}  # type: ignore
 
     @property
@@ -85,130 +65,131 @@ class Tracing(metaclass=SingletonMeta):
         """
 
         return AsyncAgentaApi(
-            base_url=self.host, api_key=self.api_key, timeout=120  # type: ignore
+            base_url=self.base_url, api_key=self.api_key, timeout=120  # type: ignore
         ).observability
 
     def set_span_attribute(
+        self, parent_key: Optional[str] = None, attributes: Dict[str, Any] = {}
+    ):
+        span = self.span_dict[self.active_span.id]  # type: ignore
+        for key, value in attributes.items():
+            self.set_attribute(span.attributes, key, value, parent_key)  # type: ignore
+
+    def set_attribute(
         self,
-        attributes: Dict[str, Any] = {},
+        span_attributes: Dict[str, Any],
+        key: str,
+        value: Any,
+        parent_key: Optional[str] = None,
     ):
-        if (
-            self.active_span is None
-        ):  # This is the case where entrypoint wants to save the trace information but the parent span has not been initialized yet
-            for key, value in attributes.items():
-                self.trace_config_cache[key] = value
+        if parent_key is not None:
+            model_config = span_attributes.get(parent_key, None)
+            if not model_config:
+                span_attributes[parent_key] = {}
+            span_attributes[parent_key][key] = value
         else:
-            for key, value in attributes.items():
-                self.active_span.attributes[key] = value
+            span_attributes[key] = value
 
     def set_trace_tags(self, tags: List[str]):
         self.tags.extend(tags)
 
+    def start_parent_span(
+        self, name: str, inputs: Dict[str, Any], config: Dict[str, Any], **kwargs
+    ):
+        trace_id = self._create_trace_id()
+        span_id = self._create_span_id()
+        self.llm_logger.info("Recording parent span...")
+        span = CreateSpan(
+            id=span_id,
+            app_id=self.app_id,
+            variant_id=self.variant_id,
+            variant_name=self.variant_name,
+            inputs=inputs,
+            name=name,
+            config=config,
+            environment=kwargs.get("environment"),
+            spankind=SpanKind.WORKFLOW.value,
+            status=SpanStatusCode.UNSET.value,
+            start_time=datetime.now(timezone.utc),
+        )
+        self.active_trace = span
+        self.recording_trace_id = trace_id
+        self.parent_span_id = span.id
+        self.llm_logger.info(
+            f"Recorded active_trace and setting parent_span_id: {span.id}"
+        )
+
     def start_span(
         self,
         name: str,
         spankind: str,
         input: Dict[str, Any],
-        config: Optional[Dict[str, Any]] = None,
-        **kwargs,
+        config: Dict[str, Any] = {},
     ) -> CreateSpan:
         span_id = self._create_span_id()
-        self.llm_logger.info(
-            f"Recording {'parent' if spankind == 'workflow' else spankind} span..."
-        )
+        self.llm_logger.info(f"Recording {spankind} span...")
         span = CreateSpan(
             id=span_id,
             inputs=input,
             name=name,
             app_id=self.app_id,
+            variant_id=self.variant_id,
+            variant_name=self.variant_name,
             config=config,
+            environment=self.active_trace.environment,
+            parent_span_id=self.parent_span_id,
             spankind=spankind.upper(),
             attributes={},
             status=SpanStatusCode.UNSET.value,
             start_time=datetime.now(timezone.utc),
-            outputs=None,
-            tags=None,
-            user=None,
-            end_time=None,
-            tokens=None,
-            cost=None,
-            token_consumption=None,
-            parent_span_id=None,
         )
 
-        if self.active_trace_id is None:  # This is a parent span
-            self.active_trace_id = self._create_trace_id()
-            span.environment = (
-                self.trace_config_cache.get("environment")
-                if self.trace_config_cache is not None
-                else os.environ.get("environment", "unset")
-            )
-            span.config = (
-                self.trace_config_cache.get("config")
-                if not config and self.trace_config_cache is not None
-                else None
-            )
-            if VARIANT_TRACKING_FEATURE_FLAG:
-                # TODO: we should get the variant_id and variant_name (and environment) from the config object
-                span.variant_id = config.variant_id
-                span.variant_name = (config.variant_name,)
-
-        else:
-            span.parent_span_id = self.active_span.id
-        self.span_dict[span.id] = span
         self.active_span = span
-
-        self.llm_logger.info(f"Recorded span and setting parent_span_id: {span.id}")
+        self.span_dict[span.id] = span
+        self.parent_span_id = span.id
+        self.llm_logger.info(
+            f"Recorded active_span and setting parent_span_id: {span.id}"
+        )
         return span
 
     def update_span_status(self, span: CreateSpan, value: str):
-        span.status = value
-
-    def end_span(self, outputs: Dict[str, Any]):
-        """
-        Ends the active span, if it is a parent span, ends the trace too.
-        """
-        if self.active_span is None:
-            raise ValueError("There is no active span to end.")
-        self.active_span.end_time = datetime.now(timezone.utc)
-        self.active_span.outputs = [outputs.get("message", "")]
-        self.active_span.cost = outputs.get("cost", None)
-        self.active_span.tokens = outputs.get("usage", None)
+        updated_span = CreateSpan(**{**span.dict(), "status": value})
+        self.active_span = updated_span
+
+    def end_span(self, outputs: Dict[str, Any], span: CreateSpan, **kwargs):
+        updated_span = CreateSpan(
+            **span.dict(),
+            end_time=datetime.now(timezone.utc),
+            outputs=[outputs["message"]],
+            cost=outputs.get("cost", None),
+            tokens=outputs.get("usage"),
+        )
 
         # Push span to list of recorded spans
-        self.pending_spans.append(self.active_span)
+        self.recorded_spans.append(updated_span)
         self.llm_logger.info(
-            f"Pushed {self.active_span.spankind} span {self.active_span.id} to recorded spans."
+            f"Pushed {updated_span.spankind} span {updated_span.id} to recorded spans."
         )
-        if self.active_span.parent_span_id is None:
-            self.end_trace(parent_span=self.active_span)
-        else:
-            self.active_span = self.span_dict[self.active_span.parent_span_id]
 
-    def end_trace(self, parent_span: CreateSpan):
+    def end_recording(self, outputs: Dict[str, Any], span: CreateSpan, **kwargs):
+        self.end_span(outputs=outputs, span=span, **kwargs)
         if self.api_key == "":
             return
 
-        if not self.active_trace_id:
-            raise RuntimeError("No active trace to end.")
-
-        self.llm_logger.info("Preparing to send recorded spans for processing.")
-        self.llm_logger.info(f"Recorded spans => {len(self.pending_spans)}")
+        self.llm_logger.info(f"Preparing to send recorded spans for processing.")
+        self.llm_logger.info(f"Recorded spans => {len(self.recorded_spans)}")
         self.tasks_manager.add_task(
-            self.active_trace_id,
+            self.active_trace.id,
             "trace",
             self.client.create_traces(
-                trace=self.active_trace_id, spans=self.pending_spans  # type: ignore
+                trace=self.recording_trace_id, spans=self.recorded_spans  # type: ignore
             ),
             self.client,
         )
         self.llm_logger.info(
-            f"Tracing for {parent_span.id} recorded successfully and sent for processing."
+            f"Tracing for {span.id} recorded successfully and sent for processing."
         )
-        self._clear_pending_spans()
-        self.active_trace_id = None
-        self.active_span = None
-        self.trace_config_cache.clear()
+        self._clear_recorded_spans()
 
     def _create_trace_id(self) -> str:
         """Creates a unique mongo id for the trace object.
@@ -228,12 +209,12 @@ class Tracing(metaclass=SingletonMeta):
 
         return str(ObjectId())
 
-    def _clear_pending_spans(self) -> None:
+    def _clear_recorded_spans(self) -> None:
         """
         Clear the list of recorded spans to prepare for next batch processing.
         """
 
-        self.pending_spans = []
+        self.recorded_spans = []
         self.llm_logger.info(
-            f"Cleared all recorded spans from batch: {self.pending_spans}"
+            f"Cleared all recorded spans from batch: {self.recorded_spans}"
         )

agenta/sdk/tracing/logger.py

@@ -2,7 +2,7 @@ import logging
 
 
 class LLMLogger:
-    def __init__(self, name="LLMLogger", level=logging.DEBUG):
+    def __init__(self, name="LLMLogger", level=logging.INFO):
         self.logger = logging.getLogger(name)
         self.logger.setLevel(level)
 

agenta/sdk/tracing/tasks_manager.py

@@ -106,7 +106,9 @@ class TaskQueue(object):
                 future.result()
             except Exception as exc:
                 self._logger.error(f"Error running task: {str(exc)}")
-                self._logger.error(f"Recording {task.coroutine_type} status to ERROR.")
+                self._logger.error(
+                    f"Recording trace {task.coroutine_type} status to ERROR."
+                )
                 break
             finally:
                 self.tasks.task_done()

agenta/sdk/types.py

@@ -1,7 +1,7 @@
 import json
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
-from pydantic import BaseModel, Extra, HttpUrl, Field
+from pydantic import ConfigDict, BaseModel, HttpUrl
 
 
 class InFile:
@@ -24,87 +24,75 @@ class FuncResponse(BaseModel):
 
 
 class DictInput(dict):
-    def __new__(cls, default_keys=None):
+    def __new__(cls, default_keys: Optional[List[str]] = None):
         instance = super().__new__(cls, default_keys)
         if default_keys is None:
             default_keys = []
-        instance.data = [key for key in default_keys]
+        instance.data = [key for key in default_keys]  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update({"x-parameter": "dict"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "dict"}
 
 
 class TextParam(str):
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update({"x-parameter": "text"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "text", "type": "string"}
 
 
 class BinaryParam(int):
     def __new__(cls, value: bool = False):
         instance = super().__new__(cls, int(value))
-        instance.default = value
+        instance.default = value  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "bool",
-                "type": "boolean",
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {
+            "x-parameter": "bool",
+            "type": "boolean",
+        }
 
 
 class IntParam(int):
     def __new__(cls, default: int = 6, minval: float = 1, maxval: float = 10):
         instance = super().__new__(cls, default)
-        instance.minval = minval
-        instance.maxval = maxval
+        instance.minval = minval  # type: ignore
+        instance.maxval = maxval  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "int",
-                "type": "integer",
-                "minimum": 1,
-                "maximum": 10,
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "int", "type": "integer"}
 
 
 class FloatParam(float):
     def __new__(cls, default: float = 0.5, minval: float = 0.0, maxval: float = 1.0):
         instance = super().__new__(cls, default)
-        instance.minval = minval
-        instance.maxval = maxval
+        instance.default = default  # type: ignore
+        instance.minval = minval  # type: ignore
+        instance.maxval = maxval  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "float",
-                "type": "number",
-                "minimum": 0.0,
-                "maximum": 1.0,
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "float", "type": "number"}
 
 
 class MultipleChoiceParam(str):
-    def __new__(cls, default: str = None, choices: List[str] = None):
-        if type(default) is list:
+    def __new__(
+        cls, default: Optional[str] = None, choices: Optional[List[str]] = None
+    ):
+        if default is not None and type(default) is list:
             raise ValueError(
                 "The order of the parameters for MultipleChoiceParam is wrong! It's MultipleChoiceParam(default, choices) and not the opposite"
             )
-        if default is None and choices:
+
+        if not default and choices is not None:
             # if a default value is not provided,
-            # uset the first value in the choices list
+            # set the first value in the choices list
             default = choices[0]
 
         if default is None and not choices:
@@ -112,28 +100,25 @@ class MultipleChoiceParam(str):
             raise ValueError("You must provide either a default value or choices")
 
         instance = super().__new__(cls, default)
-        instance.choices = choices
-        instance.default = default
+        instance.choices = choices  # type: ignore
+        instance.default = default  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema: dict[str, Any]):
-        field_schema.update(
-            {
-                "x-parameter": "choice",
-                "type": "string",
-                "enum": [],
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "choice", "type": "string", "enum": []}
 
 
 class GroupedMultipleChoiceParam(str):
-    def __new__(cls, default: str = None, choices: Dict[str, List[str]] = None):
+    def __new__(
+        cls,
+        default: Optional[str] = None,
+        choices: Optional[Dict[str, List[str]]] = None,
+    ):
         if choices is None:
             choices = {}
-        if default and not any(
-            default in choice_list for choice_list in choices.values()
-        ):
+
+        if default and not any(default in choices for choices in choices.values()):
             if not choices:
                 print(
                     f"Warning: Default value {default} provided but choices are empty."
@@ -144,31 +129,23 @@ class GroupedMultipleChoiceParam(str):
                 )
 
         if not default:
-            for choices in choices.values():
-                if choices:
-                    default = choices[0]
-                    break
+            default_selected_choice = next(
+                (choices for choices in choices.values()), None
+            )
+            if default_selected_choice:
+                default = default_selected_choice[0]
 
         instance = super().__new__(cls, default)
-        instance.choices = choices
-        instance.default = default
+        instance.choices = choices  # type: ignore
+        instance.default = default  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema: dict[str, Any], **kwargs):
-        choices = kwargs.get("choices", {})
-        field_schema.update(
-            {
-                "x-parameter": "grouped_choice",
-                "type": "string",
-                "choices": choices,
-            }
-        )
-
-
-class Message(BaseModel):
-    role: str
-    content: str
+    def __schema_type_properties__(cls) -> dict:
+        return {
+            "x-parameter": "grouped_choice",
+            "type": "string",
+        }
 
 
 class MessagesInput(list):
@@ -183,28 +160,32 @@ class MessagesInput(list):
 
     """
 
-    def __new__(cls, messages: List[Dict[str, str]] = None):
-        instance = super().__new__(cls, messages)
-        instance.default = messages
+    def __new__(cls, messages: List[Dict[str, str]] = []):
+        instance = super().__new__(cls)
+        instance.default = messages  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema: dict[str, Any]):
-        field_schema.update({"x-parameter": "messages", "type": "array"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "messages", "type": "array"}
 
 
 class FileInputURL(HttpUrl):
+    def __new__(cls, url: str):
+        instance = super().__new__(cls, url)
+        instance.default = url  # type: ignore
+        return instance
+
     @classmethod
-    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
-        field_schema.update({"x-parameter": "file_url", "type": "string"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "file_url", "type": "string"}
 
 
 class Context(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     def to_json(self):
-        return self.json()
+        return self.model_dump()
 
     @classmethod
     def from_json(cls, json_str: str):

agenta/sdk/utils/globals.py

@@ -1,7 +1,7 @@
 import agenta
 
 
-def set_global(setup=None, config=None, tracing=None):
+def set_global(setup=None, config=None):
     """Allows usage of agenta.config and agenta.setup in the user's code.
 
     Args:
@@ -12,5 +12,3 @@ def set_global(setup=None, config=None, tracing=None):
         agenta.setup = setup
     if config is not None:
         agenta.config = config
-    if tracing is not None:
-        agenta.tracing = tracing

{agenta-0.15.0.dist-info → agenta-0.15.0a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.15.0
+Version: 0.15.0a0
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering
@@ -18,12 +18,12 @@ Classifier: Topic :: Software Development :: Libraries
 Requires-Dist: cachetools (>=5.3.3,<6.0.0)
 Requires-Dist: click (>=8.1.3,<9.0.0)
 Requires-Dist: docker (>=6.1.1,<8.0.0)
-Requires-Dist: fastapi (>=0.96.1)
+Requires-Dist: fastapi (>=0.111.0,<0.112.0)
 Requires-Dist: httpx (>=0.24,<0.28)
 Requires-Dist: importlib-metadata (>=6.7,<8.0)
 Requires-Dist: ipdb (>=0.13)
 Requires-Dist: posthog (>=3.1.0,<4.0.0)
-Requires-Dist: pydantic (==1.10.13)
+Requires-Dist: pydantic (>=2.7.1,<3.0.0)
 Requires-Dist: pymongo (>=4.6.3,<5.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
 Requires-Dist: python-multipart (>=0.0.6,<0.0.10)