agenta 0.27.0a8__py3-none-any.whl → 0.27.0a12__py3-none-any.whl

agenta/sdk/__init__.py

@@ -1,6 +1,9 @@
 from typing import Optional
 
 from .utils.preinit import PreInitObject  # always the first import!
+
+import agenta.client.backend.types as client_types  # pylint: disable=wrong-import-order
+
 from .types import (
     DictInput,
     MultipleChoice,
@@ -21,12 +24,19 @@ from .tracing.conventions import Reference
 from .decorators.routing import entrypoint, app, route
 from .agenta_init import Config, AgentaSingleton, init as _init
 from .utils.costs import calculate_token_usage
-from .config_manager import ConfigManager
+from .managers.config import ConfigManager
+from .managers.variant import VariantManager
+from .managers.deployment import DeploymentManager
 
 config = PreInitObject("agenta.config", Config)
 DEFAULT_AGENTA_SINGLETON_INSTANCE = AgentaSingleton()
-tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
 
+types = client_types
+
+api = None
+async_api = None
+
+tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
 tracer = get_tracer(tracing)
 
 
@@ -36,14 +46,18 @@ def init(
     api_key: Optional[str] = None,
     config_fname: Optional[str] = None,
 ):
-    global tracing
-    global tracer
+    global api, async_api, tracing, tracer
 
     _init(
         host=host,
-        app_id=app_id,
         api_key=api_key,
         config_fname=config_fname,
+        # DEPRECATING
+        app_id=app_id,
     )
 
+    api = DEFAULT_AGENTA_SINGLETON_INSTANCE.api  # type: ignore
+    async_api = DEFAULT_AGENTA_SINGLETON_INSTANCE.async_api  # type: ignore
+
+    tracing = DEFAULT_AGENTA_SINGLETON_INSTANCE.tracing  # type: ignore
     tracer = get_tracer(tracing)

agenta/sdk/agenta_init.py

@@ -6,7 +6,7 @@ from importlib.metadata import version
 
 from agenta.sdk.utils.logging import log
 from agenta.sdk.utils.globals import set_global
-from agenta.client.backend.client import AgentaApi
+from agenta.client.backend.client import AgentaApi, AsyncAgentaApi
 from agenta.sdk.tracing import Tracing
 from agenta.client.exceptions import APIRequestError
 
@@ -22,6 +22,9 @@ class AgentaSingleton:
     config = None
     tracing = None
 
+    api = None
+    async_api = None
+
     def __new__(cls):
         if not cls._instance:
             cls._instance = super(AgentaSingleton, cls).__new__(cls)
@@ -33,14 +36,11 @@ class AgentaSingleton:
         host: Optional[str] = None,
         api_key: Optional[str] = None,
         config_fname: Optional[str] = None,
-        #
+        # DEPRECATING
         app_id: Optional[str] = None,
     ) -> None:
-        log.info(f"---------------------------")
-        log.info(f"Agenta SDK - using version: {version('agenta')}")
-        log.info(f"---------------------------")
-
-        """Main function to initialize the singleton.
+        """
+        Main function to initialize the singleton.
 
         Initializes the singleton with the given `app_id`, `host`, and `api_key`. The order of precedence for these variables is:
         1. Explicit argument provided in the function call.
@@ -62,6 +62,10 @@ class AgentaSingleton:
             ValueError: If `app_id` is not specified either as an argument, in the config file, or in the environment variables.
         """
 
+        log.info("---------------------------")
+        log.info("Agenta SDK - using version: %s", version("agenta"))
+        log.info("---------------------------")
+
         config = {}
         if config_fname:
             config = toml.load(config_fname)
@@ -97,6 +101,16 @@ class AgentaSingleton:
             app_id=self.app_id,
         )
 
+        self.api = AgentaApi(
+            base_url=self.host + "/api",
+            api_key=api_key if api_key else "",
+        )
+
+        self.async_api = AsyncAgentaApi(
+            base_url=self.host + "/api",
+            api_key=api_key if api_key else "",
+        )
+
         self.base_id = os.environ.get("AGENTA_BASE_ID")
 
         self.config = Config(

agenta/sdk/context/routing.py

@@ -7,16 +7,17 @@ routing_context = ContextVar("routing_context", default={})
 
 @contextmanager
 def routing_context_manager(
+    *,
     config: Optional[Dict[str, Any]] = None,
-    environment: Optional[str] = None,
-    version: Optional[str] = None,
-    variant: Optional[str] = None,
+    application: Optional[Dict[str, Any]] = None,
+    variant: Optional[Dict[str, Any]] = None,
+    environment: Optional[Dict[str, Any]] = None,
 ):
     context = {
         "config": config,
-        "environment": environment,
-        "version": version,
+        "application": application,
         "variant": variant,
+        "environment": environment,
     }
     token = routing_context.set(context)
     try:

agenta/sdk/decorators/routing.py

@@ -9,6 +9,7 @@ from asyncio import sleep, get_event_loop
 from traceback import format_exc, format_exception
 from pathlib import Path
 from tempfile import NamedTemporaryFile
+from os import environ
 
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import Body, FastAPI, UploadFile, HTTPException
@@ -152,7 +153,6 @@ class entrypoint:
 
             with routing_context_manager(
                 config=api_config_params,
-                environment="playground",
             ):
                 entrypoint_result = await self.execute_function(
                     func,
@@ -209,7 +209,9 @@ class entrypoint:
         @wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
             func_params = {
-                k: v for k, v in kwargs.items() if k not in ["config", "environment"]
+                k: v
+                for k, v in kwargs.items()
+                if k not in ["config", "environment", "app"]
             }
             if not config_schema:
                 if "environment" in kwargs and kwargs["environment"] is not None:
@@ -219,10 +221,19 @@ class entrypoint:
                 else:
                     ag.config.pull(config_name="default")
 
+            app_id = environ.get("AGENTA_APP_ID")
+
             with routing_context_manager(
-                config=config_params,
-                variant=kwargs["config"],
-                environment=kwargs["environment"],
+                application={
+                    "id": app_id,
+                    "slug": kwargs["app"],
+                },
+                variant={
+                    "slug": kwargs.get("config"),
+                },
+                environment={
+                    "slug": kwargs.get("environment"),
+                },
             ):
                 entrypoint_result = await self.execute_function(
                     func,

agenta/sdk/decorators/tracing.py

@@ -10,12 +10,12 @@ from agenta.sdk.tracing.conventions import parse_span_kind
 import agenta as ag
 
 
-class instrument:
+class instrument:  # pylint: disable=invalid-name
     DEFAULT_KEY = "__default__"
 
     def __init__(
         self,
-        type: str = "task",
+        type: str = "task",  # pylint: disable=redefined-builtin
         config: Optional[Dict[str, Any]] = None,
         ignore_inputs: Optional[bool] = None,
         ignore_outputs: Optional[bool] = None,
@@ -196,12 +196,16 @@ class instrument:
         ignore: Union[List[str], bool] = False,
     ) -> Dict[str, Any]:
         """
-        Redact user-defined sensitive information from inputs and outputs as defined by the ignore list or boolean flag.
+        Redact user-defined sensitive information
+        from inputs and outputs as defined by the ignore list or boolean flag.
 
         Example:
-        - ignore = ["password"] -> {"username": "admin", "password": "********"} -> {"username": "admin"}
-        - ignore = True -> {"username": "admin", "password": "********"} -> {}
-        - ignore = False -> {"username": "admin", "password": "********"} -> {"username": "admin", "password": "********"}
+        - ignore = ["password"] -> {"username": "admin", "password": "********"}
+            -> {"username": "admin"}
+        - ignore = True -> {"username": "admin", "password": "********"}
+            -> {}
+        - ignore = False -> {"username": "admin", "password": "********"}
+            -> {"username": "admin", "password": "********"}
         """
         io = {
             key: value
@@ -226,9 +230,12 @@ class instrument:
         Patch the result to ensure that it is a dictionary, with a default key when necessary.
 
         Example:
-        - result = "Hello, World!" -> {"__default__": "Hello, World!"}
-        - result = {"message": "Hello, World!", "cost": 0.0, "usage": {}} -> {"__default__": "Hello, World!"}
-        - result = {"message": "Hello, World!"} -> {"message": "Hello, World!"}
+        - result = "Hello, World!"
+            -> {"__default__": "Hello, World!"}
+        - result = {"message": "Hello, World!", "cost": 0.0, "usage": {}}
+            -> {"__default__": "Hello, World!"}
+        - result = {"message": "Hello, World!"}
+            -> {"message": "Hello, World!"}
         """
         outputs = (
             {instrument.DEFAULT_KEY: result}

agenta/sdk/litellm/litellm.py

@@ -1,16 +1,15 @@
-import agenta as ag
-
 from opentelemetry.trace import SpanKind
 
+import agenta as ag
+
 from agenta.sdk.tracing.spans import CustomSpan
-from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.exceptions import suppress  # TODO: use it !
 from agenta.sdk.utils.logging import log
 
 
 def litellm_handler():
     try:
-        from litellm.utils import ModelResponse
-        from litellm.integrations.custom_logger import (
+        from litellm.integrations.custom_logger import (  # pylint: disable=import-outside-toplevel
             CustomLogger as LitellmCustomLogger,
         )
     except ImportError as exc:
@@ -18,18 +17,23 @@ def litellm_handler():
             "The litellm SDK is not installed. Please install it using `pip install litellm`."
         ) from exc
     except Exception as exc:
-        raise Exception(
-            "Unexpected error occurred when importing litellm: {}".format(exc)
+        raise Exception(  # pylint: disable=broad-exception-raised
+            f"Unexpected error occurred when importing litellm: {exc}"
         ) from exc
 
     class LitellmHandler(LitellmCustomLogger):
-        """This handler is responsible for instrumenting certain events when using litellm to call LLMs.
+        """
+        This handler is responsible for instrumenting certain events,
+        when using litellm to call LLMs.
 
         Args:
-            LitellmCustomLogger (object): custom logger that allows us to override the events to capture.
+            LitellmCustomLogger (object): custom logger that allows us
+            to override the events to capture.
         """
 
         def __init__(self):
+            super().__init__()
+
             self.span = None
 
         def log_pre_api_call(
@@ -38,7 +42,7 @@ def litellm_handler():
             messages,
             kwargs,
         ):
-            type = (
+            type = (  # pylint: disable=redefined-builtin
                 "chat"
                 if kwargs.get("call_type") in ["completion", "acompletion"]
                 else "embedding"
@@ -59,10 +63,8 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            log.info(f"log_pre_api_call({hex(self.span.context.span_id)[2:]})")
-
             self.span.set_attributes(
-                attributes={"inputs": {"messages": kwargs["messages"]}},
+                attributes={"inputs": {"prompt": kwargs["messages"]}},
                 namespace="data",
             )
 
@@ -87,12 +89,14 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # log.info(f"log_stream({hex(self.span.context.span_id)[2:]})")
+            result = kwargs.get("complete_streaming_response")
+
+            outputs = (
+                {"__default__": result} if not isinstance(result, dict) else result
+            )
 
             self.span.set_attributes(
-                attributes={
-                    "output": {"__default__": kwargs.get("complete_streaming_response")}
-                },
+                attributes={"outputs": outputs},
                 namespace="data",
             )
 
@@ -127,14 +131,20 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # log.info(f"log_success({hex(self.span.context.span_id)[2:]})")
+            try:
+                result = []
+                for choice in response_obj.choices:
+                    message = choice.message.__dict__
+                    result.append(message)
 
-            self.span.set_attributes(
-                attributes={
-                    "output": {"__default__": response_obj.choices[0].message.content}
-                },
-                namespace="data",
-            )
+                outputs = {"completion": result}
+                self.span.set_attributes(
+                    attributes={"outputs": outputs},
+                    namespace="data",
+                )
+
+            except Exception as e:
+                pass
 
             self.span.set_attributes(
                 attributes={"total": kwargs.get("response_cost")},
@@ -167,8 +177,6 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # log.info(f"log_failure({hex(self.span.context.span_id)[2:]})")
-
             self.span.record_exception(kwargs["exception"])
 
             self.span.set_status(status="ERROR")
@@ -186,12 +194,14 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # log.info(f"async_log_stream({hex(self.span.context.span_id)[2:]})")
+            result = kwargs.get("complete_streaming_response")
+
+            outputs = (
+                {"__default__": result} if not isinstance(result, dict) else result
+            )
 
             self.span.set_attributes(
-                attributes={
-                    "output": {"__default__": kwargs.get("complete_streaming_response")}
-                },
+                attributes={"outputs": outputs},
                 namespace="data",
             )
 
@@ -226,12 +236,15 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            log.info(f"async_log_success({hex(self.span.context.span_id)[2:]})")
+            # result = kwargs.get("complete_streaming_response")
+            result = response_obj.choices[0].message.content
+
+            outputs = (
+                {"__default__": result} if not isinstance(result, dict) else result
+            )
 
             self.span.set_attributes(
-                attributes={
-                    "output": {"__default__": kwargs.get("complete_streaming_response")}
-                },
+                attributes={"outputs": outputs},
                 namespace="data",
             )
 
@@ -266,8 +279,6 @@ def litellm_handler():
                 log.error("LiteLLM callback error: span not found.")
                 return
 
-            # log.info(f"async_log_failure({hex(self.span.context.span_id)[2:]})")
-
             self.span.record_exception(kwargs["exception"])
 
             self.span.set_status(status="ERROR")

agenta/sdk/managers/__init__.py

@@ -0,0 +1,6 @@
+from agenta.sdk.managers.config import ConfigManager
+from agenta.sdk.managers.variant import VariantManager
+from agenta.sdk.managers.deployment import DeploymentManager
+
+
+__all__ = ["ConfigManager", "VariantManager", "DeploymentManager"]