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

agenta/sdk/tracing/llm_tracing.py

@@ -1,59 +1,79 @@
-# Stdlib Imports
+import os
+from threading import Lock
 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.
 
-    Args:
-        base_url (str): The URL of the backend host
+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
         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,
-        base_url: str,
+        host: str,
         app_id: str,
-        variant_id: str,
-        variant_name: Optional[str] = None,
         api_key: Optional[str] = None,
         max_workers: Optional[int] = None,
     ):
-        self.base_url = base_url + "/api"
+        self.host = host + "/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 = CreateSpan
-        self.active_trace = CreateSpan
-        self.recording_trace_id: Union[str, None] = None
-        self.recorded_spans: List[CreateSpan] = []
+        self.active_span: Optional[CreateSpan] = None
+        self.active_trace_id: Optional[str] = None
+        self.pending_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
@@ -65,131 +85,130 @@ class Tracing(object):
         """
 
         return AsyncAgentaApi(
-            base_url=self.base_url, api_key=self.api_key, timeout=120  # type: ignore
+            base_url=self.host, 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,
-        span_attributes: Dict[str, Any],
-        key: str,
-        value: Any,
-        parent_key: Optional[str] = None,
+        attributes: Dict[str, Any] = {},
     ):
-        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
+        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
         else:
-            span_attributes[key] = value
+            for key, value in attributes.items():
+                self.active_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: Dict[str, Any] = {},
+        config: Optional[Dict[str, Any]] = None,
+        **kwargs,
     ) -> CreateSpan:
         span_id = self._create_span_id()
-        self.llm_logger.info(f"Recording {spankind} span...")
+        self.llm_logger.info(
+            f"Recording {'parent' if spankind == 'workflow' else 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,
         )
 
-        self.active_span = span
+        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.parent_span_id = span.id
-        self.llm_logger.info(
-            f"Recorded active_span and setting parent_span_id: {span.id}"
-        )
+        self.active_span = span
+
+        self.llm_logger.info(f"Recorded span and setting parent_span_id: {span.id}")
         return span
 
     def update_span_status(self, span: CreateSpan, value: str):
-        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"),
-        )
+        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)
 
         # Push span to list of recorded spans
-        self.recorded_spans.append(updated_span)
+        self.pending_spans.append(self.active_span)
         self.llm_logger.info(
-            f"Pushed {updated_span.spankind} span {updated_span.id} to recorded spans."
+            f"Pushed {self.active_span.spankind} span {self.active_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_recording(self, outputs: Dict[str, Any], span: CreateSpan, **kwargs):
-        self.end_span(outputs=outputs, span=span, **kwargs)
+    def end_trace(self, parent_span: CreateSpan):
         if self.api_key == "":
             return
 
-        self.llm_logger.info(f"Preparing to send recorded spans for processing.")
-        self.llm_logger.info(f"Recorded spans => {len(self.recorded_spans)}")
+        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.tasks_manager.add_task(
-            self.active_trace.id,
+            self.active_trace_id,
             "trace",
             self.client.create_traces(
-                trace=self.recording_trace_id, spans=self.recorded_spans  # type: ignore
+                trace=self.active_trace_id, spans=self.pending_spans  # type: ignore
             ),
             self.client,
         )
         self.llm_logger.info(
-            f"Tracing for {span.id} recorded successfully and sent for processing."
+            f"Tracing for {parent_span.id} recorded successfully and sent for processing."
         )
-        self._clear_recorded_spans()
+        self._clear_pending_spans()
+        self.active_trace_id = None
+        self.active_span = None
+        self.trace_config_cache.clear()
 
     def _create_trace_id(self) -> str:
         """Creates a unique mongo id for the trace object.
@@ -209,12 +228,12 @@ class Tracing(object):
 
         return str(ObjectId())
 
-    def _clear_recorded_spans(self) -> None:
+    def _clear_pending_spans(self) -> None:
         """
         Clear the list of recorded spans to prepare for next batch processing.
         """
 
-        self.recorded_spans = []
+        self.pending_spans = []
         self.llm_logger.info(
-            f"Cleared all recorded spans from batch: {self.recorded_spans}"
+            f"Cleared all recorded spans from batch: {self.pending_spans}"
         )

agenta/sdk/tracing/logger.py

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

agenta/sdk/tracing/tasks_manager.py

@@ -106,9 +106,7 @@ class TaskQueue(object):
                 future.result()
             except Exception as exc:
                 self._logger.error(f"Error running task: {str(exc)}")
-                self._logger.error(
-                    f"Recording trace {task.coroutine_type} status to ERROR."
-                )
+                self._logger.error(f"Recording {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 Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
-from pydantic import ConfigDict, BaseModel, HttpUrl
+from pydantic import BaseModel, Extra, HttpUrl, Field
 
 
 class InFile:
@@ -24,75 +24,87 @@ class FuncResponse(BaseModel):
 
 
 class DictInput(dict):
-    def __new__(cls, default_keys: Optional[List[str]] = None):
+    def __new__(cls, default_keys=None):
         instance = super().__new__(cls, default_keys)
         if default_keys is None:
             default_keys = []
-        instance.data = [key for key in default_keys]  # type: ignore
+        instance.data = [key for key in default_keys]
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "dict"}
+    def __modify_schema__(cls, field_schema):
+        field_schema.update({"x-parameter": "dict"})
 
 
 class TextParam(str):
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "text", "type": "string"}
+    def __modify_schema__(cls, field_schema):
+        field_schema.update({"x-parameter": "text"})
 
 
 class BinaryParam(int):
     def __new__(cls, value: bool = False):
         instance = super().__new__(cls, int(value))
-        instance.default = value  # type: ignore
+        instance.default = value
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {
-            "x-parameter": "bool",
-            "type": "boolean",
-        }
+    def __modify_schema__(cls, field_schema):
+        field_schema.update(
+            {
+                "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  # type: ignore
-        instance.maxval = maxval  # type: ignore
+        instance.minval = minval
+        instance.maxval = maxval
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "int", "type": "integer"}
+    def __modify_schema__(cls, field_schema):
+        field_schema.update(
+            {
+                "x-parameter": "int",
+                "type": "integer",
+                "minimum": 1,
+                "maximum": 10,
+            }
+        )
 
 
 class FloatParam(float):
     def __new__(cls, default: float = 0.5, minval: float = 0.0, maxval: float = 1.0):
         instance = super().__new__(cls, default)
-        instance.default = default  # type: ignore
-        instance.minval = minval  # type: ignore
-        instance.maxval = maxval  # type: ignore
+        instance.minval = minval
+        instance.maxval = maxval
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "float", "type": "number"}
+    def __modify_schema__(cls, field_schema):
+        field_schema.update(
+            {
+                "x-parameter": "float",
+                "type": "number",
+                "minimum": 0.0,
+                "maximum": 1.0,
+            }
+        )
 
 
 class MultipleChoiceParam(str):
-    def __new__(
-        cls, default: Optional[str] = None, choices: Optional[List[str]] = None
-    ):
-        if default is not None and type(default) is list:
+    def __new__(cls, default: str = None, choices: List[str] = None):
+        if 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 not default and choices is not None:
+        if default is None and choices:
             # if a default value is not provided,
-            # set the first value in the choices list
+            # uset the first value in the choices list
             default = choices[0]
 
         if default is None and not choices:
@@ -100,25 +112,28 @@ class MultipleChoiceParam(str):
             raise ValueError("You must provide either a default value or choices")
 
         instance = super().__new__(cls, default)
-        instance.choices = choices  # type: ignore
-        instance.default = default  # type: ignore
+        instance.choices = choices
+        instance.default = default
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "choice", "type": "string", "enum": []}
+    def __modify_schema__(cls, field_schema: dict[str, Any]):
+        field_schema.update(
+            {
+                "x-parameter": "choice",
+                "type": "string",
+                "enum": [],
+            }
+        )
 
 
 class GroupedMultipleChoiceParam(str):
-    def __new__(
-        cls,
-        default: Optional[str] = None,
-        choices: Optional[Dict[str, List[str]]] = None,
-    ):
+    def __new__(cls, default: str = None, choices: Dict[str, List[str]] = None):
         if choices is None:
             choices = {}
-
-        if default and not any(default in choices for choices in choices.values()):
+        if default and not any(
+            default in choice_list for choice_list in choices.values()
+        ):
             if not choices:
                 print(
                     f"Warning: Default value {default} provided but choices are empty."
@@ -129,23 +144,31 @@ class GroupedMultipleChoiceParam(str):
                 )
 
         if not default:
-            default_selected_choice = next(
-                (choices for choices in choices.values()), None
-            )
-            if default_selected_choice:
-                default = default_selected_choice[0]
+            for choices in choices.values():
+                if choices:
+                    default = choices[0]
+                    break
 
         instance = super().__new__(cls, default)
-        instance.choices = choices  # type: ignore
-        instance.default = default  # type: ignore
+        instance.choices = choices
+        instance.default = default
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {
-            "x-parameter": "grouped_choice",
-            "type": "string",
-        }
+    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
 
 
 class MessagesInput(list):
@@ -160,32 +183,28 @@ class MessagesInput(list):
 
     """
 
-    def __new__(cls, messages: List[Dict[str, str]] = []):
-        instance = super().__new__(cls)
-        instance.default = messages  # type: ignore
+    def __new__(cls, messages: List[Dict[str, str]] = None):
+        instance = super().__new__(cls, messages)
+        instance.default = messages
         return instance
 
     @classmethod
-    def __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "messages", "type": "array"}
+    def __modify_schema__(cls, field_schema: dict[str, Any]):
+        field_schema.update({"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 __schema_type_properties__(cls) -> dict:
-        return {"x-parameter": "file_url", "type": "string"}
+    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
+        field_schema.update({"x-parameter": "file_url", "type": "string"})
 
 
 class Context(BaseModel):
-    model_config = ConfigDict(extra="allow")
+    class Config:
+        extra = Extra.allow
 
     def to_json(self):
-        return self.model_dump()
+        return self.json()
 
     @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):
+def set_global(setup=None, config=None, tracing=None):
     """Allows usage of agenta.config and agenta.setup in the user's code.
 
     Args:
@@ -12,3 +12,5 @@ def set_global(setup=None, config=None):
         agenta.setup = setup
     if config is not None:
         agenta.config = config
+    if tracing is not None:
+        agenta.tracing = tracing

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.14.14a0
+Version: 0.15.0
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering
@@ -23,7 +23,7 @@ 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 (>=2.7.1,<3.0.0)
+Requires-Dist: pydantic (==1.10.13)
 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)