agenta 0.14.0__py3-none-any.whl → 0.14.1a0__py3-none-any.whl

agenta/__init__.py

@@ -16,6 +16,7 @@ from .sdk.types import (
 )
 from .sdk.tracing.decorators import span
 from .sdk.agenta_init import Config, init, llm_tracing
+from .sdk.tracing.callbacks import agenta_litellm_handler
 from .sdk.utils.helper.openai_cost import calculate_token_usage
 from .sdk.client import Agenta
 

agenta/sdk/__init__.py

@@ -17,6 +17,7 @@ from .types import (
 )
 from .tracing.decorators import span
 from .agenta_init import Config, init, llm_tracing
+from .tracing.callbacks import agenta_litellm_handler
 from .utils.helper.openai_cost import calculate_token_usage
 
 

agenta/sdk/agenta_decorator.py

@@ -17,6 +17,7 @@ from fastapi import Body, FastAPI, UploadFile, HTTPException
 
 import agenta
 from .context import save_context
+from .tracing.llm_tracing import Tracing
 from .router import router as router
 from .types import (
     Context,
@@ -155,10 +156,7 @@ def entrypoint(func: Callable[..., Any]) -> Callable[..., Any]:
 
     if is_main_script(func):
         handle_terminal_run(
-            func,
-            func_signature.parameters,
-            config_params,
-            ingestible_files,
+            func, func_signature.parameters, config_params, ingestible_files, tracing
         )
     return None
 
@@ -350,6 +348,7 @@ def handle_terminal_run(
     func_params: Dict[str, Any],
     config_params: Dict[str, Any],
     ingestible_files: Dict,
+    tracing: Tracing,
 ) -> None:
     """
     Parses command line arguments and sets configuration when script is run from the terminal.
@@ -357,9 +356,8 @@ def handle_terminal_run(
     Args:
         func_params (dict): A dictionary containing the function parameters and their annotations.
         config_params (dict): A dictionary containing the configuration parameters.
-
-    Example:
-        handle_terminal_run(func_params=inspect.signature(my_function).parameters, config_params=config.all())
+        ingestible_files (dict): A dictionary containing the files that should be ingested.
+        tracing (Tracing): The tracing object
     """
 
     # For required parameters, we add them as arguments
@@ -398,13 +396,29 @@ def handle_terminal_run(
         )
     agenta.config.set(**args_config_params)
 
+    # Start tracing
+    tracing.start_parent_span(
+        name=func.__name__,
+        inputs=args_func_params,
+        config=args_config_params,
+        environment="shell",  # type: ignore
+    )
+
     loop = asyncio.get_event_loop()
     result = loop.run_until_complete(
         execute_function(
             func, **{"params": args_func_params, "config_params": args_config_params}
         )
     )
-    print(result)
+
+    # End trace recording
+    tracing.end_recording(
+        outputs=result.dict(),
+        span=tracing.active_trace,  # type: ignore
+    )
+    print(
+        f"\n========== Result ==========\n\nMessage: {result.message}\nCost: {result.cost}\nToken Usage: {result.usage}"
+    )
 
 
 def override_schema(openapi_schema: dict, func_name: str, endpoint: str, params: dict):

agenta/sdk/tracing/callbacks.py

@@ -0,0 +1,125 @@
+# Own Imports
+from agenta.sdk import llm_tracing
+
+# Third Party Imports
+from litellm.utils import ModelResponse
+from litellm.integrations.custom_logger import CustomLogger as LitellmCustomLogger
+
+
+class AgentaLiteLLMHandler(LitellmCustomLogger):
+    """This handler is responsible for logging certain events when using litellm to call LLMs.
+
+    Args:
+        LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
+    """
+
+    @property
+    def _trace(self):
+        return llm_tracing()
+
+    def log_pre_api_call(self, model, messages, kwargs):
+        self._trace.start_span(
+            name="pre_api_call",
+            input=(
+                {"messages": messages}
+                if isinstance(messages, list)
+                else {"inputs": messages}
+            ),
+            spankind=(
+                "llm"
+                if kwargs["call_type"] in ["completion", "acompletion"]
+                else "unset"
+            ),
+        )
+        self._trace.set_span_attribute(
+            "model_config",
+            {
+                "model": kwargs.get("model"),
+                "temperature": kwargs["optional_params"]["temperature"],
+            },
+        )
+
+    def log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        self._trace.update_span_status(span=self._trace.active_span, value="OK")
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["complete_streaming_response"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+    def log_success_event(
+        self, kwargs, response_obj: ModelResponse, start_time, end_time
+    ):
+        self._trace.update_span_status(span=self._trace.active_span, value="OK")
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["message"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+    def log_failure_event(
+        self, kwargs, response_obj: ModelResponse, start_time, end_time
+    ):
+        self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
+        self._trace.set_span_attribute(
+            attributes={
+                "traceback_exception": kwargs["traceback_exception"],
+                "call_end_time": kwargs["end_time"],
+            },
+        )
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["exception"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+    async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        self._trace.update_span_status(span=self._trace.active_span, value="OK")
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["complete_streaming_response"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self._trace.update_span_status(span=self._trace.active_span, value="OK")
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["message"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        self._trace.update_span_status(span=self._trace.active_span, value="ERROR")
+        self._trace.set_span_attribute(
+            attributes={
+                "traceback_exception": kwargs["traceback_exception"],
+                "call_end_time": kwargs["end_time"],
+            },
+        )
+        self._trace.end_span(
+            outputs={
+                "message": kwargs["exception"],
+                "usage": kwargs["usage"],
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+
+agenta_litellm_handler = AgentaLiteLLMHandler()

agenta/sdk/tracing/llm_tracing.py

@@ -1,4 +1,5 @@
 # Stdlib Imports
+from threading import Lock
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any, List, Union
 
@@ -13,10 +14,43 @@ from agenta.client.backend.types.create_span import CreateSpan, SpanKind, SpanSt
 from bson.objectid import ObjectId
 
 
-class Tracing(object):
-    """Agenta llm tracing object.
+class SingletonMeta(type):
+    """
+    Thread-safe implementation of Singleton.
+    """
+
+    _instances = {}
+
+    # We need the lock mechanism to synchronize threads \
+    # during the initial access to the Singleton object.
+    _lock: Lock = Lock()
 
-    Args:
+    def __call__(cls, *args, **kwargs):
+        """
+        Possible changes to the value of the `__init__` argument do not affect
+        the returned instance.
+        """
+        # Now, imagine that the program has just been launched. Since there's no
+        # Singleton instance yet, multiple threads can simultaneously pass the
+        # previous conditional and reach this point almost at the same time. The
+        # first of them will acquire lock and will proceed further, while the
+        # rest will wait here.
+        with cls._lock:
+            # The first thread to acquire the lock, reaches this conditional,
+            # goes inside and creates the Singleton instance. Once it leaves the
+            # lock block, a thread that might have been waiting for the lock
+            # release may then enter this section. But since the Singleton field
+            # is already initialized, the thread won't create a new object.
+            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:
         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
@@ -24,13 +58,6 @@ class Tracing(object):
         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,

{agenta-0.14.0.dist-info → agenta-0.14.1a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.14.0
+Version: 0.14.1a0
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering
@@ -22,6 +22,7 @@ Requires-Dist: fastapi (>=0.96.1)
 Requires-Dist: httpx (>=0.27.0,<0.28.0)
 Requires-Dist: importlib-metadata (>=6.7.0,<7.0.0)
 Requires-Dist: ipdb (>=0.13)
+Requires-Dist: litellm (>=1.35.33,<2.0.0)
 Requires-Dist: posthog (>=3.1.0,<4.0.0)
 Requires-Dist: pydantic (==1.10.13)
 Requires-Dist: pymongo (>=4.6.3,<5.0.0)

{agenta-0.14.0.dist-info → agenta-0.14.1a0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-agenta/__init__.py,sha256=rCZ-mUOGnyKQpfWQFzXIEFX7KBlGC3qFMighdIVJgDc,610
+agenta/__init__.py,sha256=ZJZlhn56j2rjl-QKTuNU6HHlJLVlmd_NcU67VIBvGGI,668
 agenta/cli/evaluation_commands.py,sha256=fs6492tprPId9p8eGO02Xy-NCBm2RZNJLZWcUxugwd8,474
 agenta/cli/helper.py,sha256=vRxHyeNaltzNIGrfU2vO0H28_rXDzx9QqIZ_S-W6zL4,6212
 agenta/cli/main.py,sha256=GgYu6UsrnHbqPV7zPlO14b61IyaDiTIjGMYQS9DlqC4,9551
@@ -126,15 +126,16 @@ agenta/docker/docker-assets/entrypoint.sh,sha256=29XK8VQjQsx4hN2j-4JDy-6kQb5y4LC
 agenta/docker/docker-assets/lambda_function.py,sha256=h4UZSSfqwpfsCgERv6frqwm_4JrYu9rLz3I-LxCfeEg,83
 agenta/docker/docker-assets/main.py,sha256=7MI-21n81U7N7A0GxebNi0cmGWtJKcR2sPB6FcH2QfA,251
 agenta/docker/docker_utils.py,sha256=5uHMCzXkCvIsDdEiwbnnn97KkzsFbBvyMwogCsv_Z5U,3509
-agenta/sdk/__init__.py,sha256=jmeLRuXrew02ZruODZYIu4kpw0S8vV6JhMPQWFGtj30,648
-agenta/sdk/agenta_decorator.py,sha256=6vz0G3YCRKRzK8JrQFyy8c2RIXy2kVMwyxTS093_8vQ,17296
+agenta/sdk/__init__.py,sha256=3At8FvbuYFoP5IIYs4FZWYQK8l4XlVD1zv4E6QsgQv0,702
+agenta/sdk/agenta_decorator.py,sha256=ucoGffoIefsYaxniMfMgsepKDo8jrlVMA7gGoTYyQUc,17805
 agenta/sdk/agenta_init.py,sha256=wDfStpe8_3ZXRLtikarwDKI_VpA1YW4eIz_3fXq39is,9044
 agenta/sdk/client.py,sha256=trKyBOYFZRk0v5Eptxvh87yPf50Y9CqY6Qgv4Fy-VH4,2142
 agenta/sdk/context.py,sha256=q-PxL05-I84puunUAs9LGsffEXcYhDxhQxjuOz2vK90,901
 agenta/sdk/router.py,sha256=0sbajvn5C7t18anH6yNo7-oYxldHnYfwcbmQnIXBePw,269
+agenta/sdk/tracing/callbacks.py,sha256=hwOHgu1pGp5C7KTs0Tm07y7Uv_0gCo2YoyckKfnjmqw,4342
 agenta/sdk/tracing/context_manager.py,sha256=HskDaiORoOhjeN375gm05wYnieQzh5UnoIsnSAHkAyc,252
 agenta/sdk/tracing/decorators.py,sha256=ujtU8gf3GDoHYuLTfEYK_2eIYZ-1oX5dpv02Mf4l_II,1191
-agenta/sdk/tracing/llm_tracing.py,sha256=UiotJ56EFA3VPt7LREkcK2w51D9-0T1QNvBy4zNWEdY,7348
+agenta/sdk/tracing/llm_tracing.py,sha256=f69FtUIA187gNob8N6IZ-eiF6QQyBVl2JUMJ4Nnnp7M,8671
 agenta/sdk/tracing/logger.py,sha256=4zG9c51p8xPdKA5SL8MOgBfkpCnBSuV6JfWiXO0A7oc,473
 agenta/sdk/tracing/tasks_manager.py,sha256=XVGBEOwmHa6KcCC0PApk0_bZ0Ilk2ESuduNObB1rw2s,3792
 agenta/sdk/types.py,sha256=Mn0yBlHh_Yr_5oQXUfsYI3V7sJAVWkJgkxEOBDOOMS0,5852
@@ -156,7 +157,7 @@ agenta/templates/simple_prompt/app.py,sha256=kODgF6lhzsaJPdgL5b21bUki6jkvqjWZzWR
 agenta/templates/simple_prompt/env.example,sha256=g9AE5bYcGPpxawXMJ96gh8oenEPCHTabsiOnfQo3c5k,70
 agenta/templates/simple_prompt/requirements.txt,sha256=ywRglRy7pPkw8bljmMEJJ4aOOQKrt9FGKULZ-DGkoBU,23
 agenta/templates/simple_prompt/template.toml,sha256=DQBtRrF4GU8LBEXOZ-GGuINXMQDKGTEG5y37tnvIUIE,60
-agenta-0.14.0.dist-info/METADATA,sha256=pY6HvnMLhCHryK46gnKyCLI2WTg7wwv1Q4nCa6558oY,26474
-agenta-0.14.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-agenta-0.14.0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.14.0.dist-info/RECORD,,
+agenta-0.14.1a0.dist-info/METADATA,sha256=Z4cOQ0WJuec014xZcsQcU5CtcaKIKsAOMlLVux4rnJ4,26518
+agenta-0.14.1a0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+agenta-0.14.1a0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.14.1a0.dist-info/RECORD,,