agenta 0.14.7__tar.gz → 0.14.7a1__tar.gz

{agenta-0.14.7 → agenta-0.14.7a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.14.7
+Version: 0.14.7a1
 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.24,<0.28)
 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.7 → agenta-0.14.7a1}/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-0.14.7 → agenta-0.14.7a1}/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-0.14.7 → agenta-0.14.7a1}/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-0.14.7a1/agenta/sdk/tracing/callbacks.py

@@ -0,0 +1,129 @@
+# 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.get("call_type") in ["completion", "acompletion"]
+                else (
+                    "embedding"
+                    if kwargs.get("call_type") in ["embedding", "aembedding"]
+                    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.get("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.get("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.get("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.get("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.get("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.get("usage"),
+                "cost": kwargs.get("response_cost"),
+            },
+            span=self._trace.active_span,
+        )
+
+
+agenta_litellm_handler = AgentaLiteLLMHandler()

{agenta-0.14.7 → agenta-0.14.7a1}/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.7 → agenta-0.14.7a1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.14.7"
+version = "0.14.7a1"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]
@@ -30,6 +30,7 @@ posthog = "^3.1.0"
 pydantic = "1.10.13"
 httpx = ">=0.24, <0.28"
 pymongo = "^4.6.3"
+litellm = "^1.35.33"
 cachetools = "^5.3.3"
 
 [tool.poetry.dev-dependencies]