agenta 0.30.0a1__py3-none-any.whl → 0.30.0a3__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,8 +1,12 @@
 from typing import Callable, Optional, Any, Dict, List, Union
+
 from functools import wraps
 from itertools import chain
 from inspect import iscoroutinefunction, getfullargspec
 
+from opentelemetry import baggage as baggage
+from opentelemetry.context import attach, detach
+
 from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.context.tracing import tracing_context
 from agenta.sdk.tracing.conventions import parse_span_kind
@@ -39,10 +43,12 @@ class instrument:  # pylint: disable=invalid-name
         is_coroutine_function = iscoroutinefunction(func)
 
         @wraps(func)
-        async def async_wrapper(*args, **kwargs):
-            async def _async_auto_instrumented(*args, **kwargs):
+        async def awrapper(*args, **kwargs):
+            async def aauto_instrumented(*args, **kwargs):
                 self._parse_type_and_kind()
 
+                token = self._attach_baggage()
+
                 with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
                     self._pre_instrument(func, *args, **kwargs)
 
@@ -52,13 +58,17 @@ class instrument:  # pylint: disable=invalid-name
 
                     return result
 
-            return await _async_auto_instrumented(*args, **kwargs)
+                self._detach_baggage(token)
+
+            return await aauto_instrumented(*args, **kwargs)
 
         @wraps(func)
-        def sync_wrapper(*args, **kwargs):
-            def _sync_auto_instrumented(*args, **kwargs):
+        def wrapper(*args, **kwargs):
+            def auto_instrumented(*args, **kwargs):
                 self._parse_type_and_kind()
 
+                token = self._attach_baggage()
+
                 with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
                     self._pre_instrument(func, *args, **kwargs)
 
@@ -68,9 +78,11 @@ class instrument:  # pylint: disable=invalid-name
 
                     return result
 
-            return _sync_auto_instrumented(*args, **kwargs)
+                self._detach_baggage(token)
+
+            return auto_instrumented(*args, **kwargs)
 
-        return async_wrapper if is_coroutine_function else sync_wrapper
+        return awrapper if is_coroutine_function else wrapper
 
     def _parse_type_and_kind(self):
         if not ag.tracing.get_current_span().is_recording():
@@ -78,6 +90,25 @@ class instrument:  # pylint: disable=invalid-name
 
         self.kind = parse_span_kind(self.type)
 
+    def _attach_baggage(self):
+        context = tracing_context.get()
+
+        references = context.references
+
+        token = None
+        if references:
+            for k, v in references.items():
+                token = attach(baggage.set_baggage(f"ag.refs.{k}", v))
+
+        return token
+
+    def _detach_baggage(
+        self,
+        token,
+    ):
+        if token:
+            detach(token)
+
     def _pre_instrument(
         self,
         func,
@@ -86,29 +117,21 @@ class instrument:  # pylint: disable=invalid-name
     ):
         span = ag.tracing.get_current_span()
 
+        context = tracing_context.get()
+
         with suppress():
+            trace_id = span.context.trace_id
+
+            ag.tracing.credentials[trace_id] = context.credentials
+
             span.set_attributes(
                 attributes={"node": self.type},
                 namespace="type",
             )
 
             if span.parent is None:
-                rctx = tracing_context.get()
-
-                span.set_attributes(
-                    attributes={"configuration": rctx.get("config", {})},
-                    namespace="meta",
-                )
-                span.set_attributes(
-                    attributes={"environment": rctx.get("environment", {})},
-                    namespace="meta",
-                )
                 span.set_attributes(
-                    attributes={"version": rctx.get("version", {})},
-                    namespace="meta",
-                )
-                span.set_attributes(
-                    attributes={"variant": rctx.get("variant", {})},
+                    attributes={"configuration": context.parameters or {}},
                     namespace="meta",
                 )
 
@@ -118,6 +141,7 @@ class instrument:  # pylint: disable=invalid-name
                 io=self._parse(func, *args, **kwargs),
                 ignore=self.ignore_inputs,
             )
+
             span.set_attributes(
                 attributes={"inputs": _inputs},
                 namespace="data",
@@ -161,6 +185,7 @@ class instrument:  # pylint: disable=invalid-name
                 io=self._patch(result),
                 ignore=self.ignore_outputs,
             )
+
             span.set_attributes(
                 attributes={"outputs": _outputs},
                 namespace="data",
@@ -171,15 +196,12 @@ class instrument:  # pylint: disable=invalid-name
 
         with suppress():
             if hasattr(span, "parent") and span.parent is None:
-                tracing_context.set(
-                    tracing_context.get()
-                    | {
-                        "root": {
-                            "trace_id": span.get_span_context().trace_id,
-                            "span_id": span.get_span_context().span_id,
-                        }
-                    }
-                )
+                context = tracing_context.get()
+                context.link = {
+                    "tree_id": span.get_span_context().trace_id,
+                    "node_id": span.get_span_context().span_id,
+                }
+                tracing_context.set(context)
 
     def _parse(
         self,

agenta/sdk/managers/config.py

@@ -7,7 +7,7 @@ import yaml
 from pydantic import BaseModel
 
 from agenta.sdk.managers.shared import SharedManager
-from agenta.sdk.decorators.routing import routing_context
+from agenta.sdk.context.routing import routing_context
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -20,7 +20,7 @@ class ConfigManager:
     @staticmethod
     def get_from_route(
         schema: Optional[Type[T]] = None,
-    ) -> Union[Dict[str, Any], T]:
+    ) -> Optional[Union[Dict[str, Any], T]]:
         """
         Retrieves the configuration from the route context and returns a config object.
 
@@ -47,125 +47,15 @@ class ConfigManager:
 
         context = routing_context.get()
 
-        parameters = None
-
-        if "config" in context and context["config"]:
-            parameters = context["config"]
-
-        else:
-            app_id: Optional[str] = None
-            app_slug: Optional[str] = None
-            variant_id: Optional[str] = None
-            variant_slug: Optional[str] = None
-            variant_version: Optional[int] = None
-            environment_id: Optional[str] = None
-            environment_slug: Optional[str] = None
-            environment_version: Optional[int] = None
-
-            if "application" in context:
-                app_id = context["application"].get("id")
-                app_slug = context["application"].get("slug")
-
-            if "variant" in context:
-                variant_id = context["variant"].get("id")
-                variant_slug = context["variant"].get("slug")
-                variant_version = context["variant"].get("version")
-
-            if "environment" in context:
-                environment_id = context["environment"].get("id")
-                environment_slug = context["environment"].get("slug")
-                environment_version = context["environment"].get("version")
-
-            parameters = ConfigManager.get_from_registry(
-                app_id=app_id,
-                app_slug=app_slug,
-                variant_id=variant_id,
-                variant_slug=variant_slug,
-                variant_version=variant_version,
-                environment_id=environment_id,
-                environment_slug=environment_slug,
-                environment_version=environment_version,
-            )
+        parameters = context.parameters
 
-        if schema:
-            return schema(**parameters)
-
-        return parameters
+        if not parameters:
+            return None
 
-    @staticmethod
-    async def aget_from_route(
-        schema: Optional[Type[T]] = None,
-    ) -> Union[Dict[str, Any], T]:
-        """
-        Asynchronously retrieves the configuration from the route context and returns a config object.
+        if not schema:
+            return parameters
 
-        This method checks the route context for configuration information and returns
-        an instance of the specified schema based on the available context data.
-
-        Args:
-            schema (Type[T]): A Pydantic model class that defines the structure of the configuration.
-
-        Returns:
-            T: An instance of the specified schema populated with the configuration data.
-
-        Raises:
-            ValueError: If conflicting configuration sources are provided or if no valid
-                        configuration source is found in the context.
-
-        Note:
-            The method prioritizes the inputs in the following way:
-            1. 'config' (i.e. when called explicitly from the playground)
-            2. 'environment'
-            3. 'variant'
-            Only one of these should be provided.
-        """
-
-        context = routing_context.get()
-
-        parameters = None
-
-        if "config" in context and context["config"]:
-            parameters = context["config"]
-
-        else:
-            app_id: Optional[str] = None
-            app_slug: Optional[str] = None
-            variant_id: Optional[str] = None
-            variant_slug: Optional[str] = None
-            variant_version: Optional[int] = None
-            environment_id: Optional[str] = None
-            environment_slug: Optional[str] = None
-            environment_version: Optional[int] = None
-
-            if "application" in context:
-                app_id = context["application"].get("id")
-                app_slug = context["application"].get("slug")
-
-            if "variant" in context:
-                variant_id = context["variant"].get("id")
-                variant_slug = context["variant"].get("slug")
-                variant_version = context["variant"].get("version")
-
-            if "environment" in context:
-                environment_id = context["environment"].get("id")
-                environment_slug = context["environment"].get("slug")
-                environment_version = context["environment"].get("version")
-
-            parameters = await ConfigManager.async_get_from_registry(
-                app_id=app_id,
-                app_slug=app_slug,
-                variant_id=variant_id,
-                variant_slug=variant_slug,
-                variant_version=variant_version,
-                environment_id=environment_id,
-                environment_slug=environment_slug,
-                environment_version=environment_version,
-            )
-
-        if schema:
-            return schema(**parameters)
-
-        return parameters
+        return schema(**parameters)
 
     @staticmethod
     def get_from_registry(

agenta/sdk/managers/secrets.py

@@ -0,0 +1,38 @@
+from typing import Optional, Dict, Any
+
+from agenta.sdk.context.routing import routing_context
+
+from agenta.sdk.assets import model_to_provider_mapping
+
+
+class SecretsManager:
+    @staticmethod
+    def get_from_route() -> Optional[Dict[str, Any]]:
+        context = routing_context.get()
+
+        secrets = context.secrets
+
+        if not secrets:
+            return None
+
+        return secrets
+
+    @staticmethod
+    def get_api_key_for_model(model: str) -> str:
+        secrets = SecretsManager.get_from_route()
+
+        if not secrets:
+            return None
+
+        provider = model_to_provider_mapping.get(model)
+
+        if not provider:
+            return None
+
+        provider = provider.lower().replace(" ", "")
+
+        for secret in secrets:
+            if secret["data"]["provider"] == provider:
+                return secret["data"]["key"]
+
+        return None

agenta/sdk/middleware/auth.py

@@ -1,90 +1,116 @@
 from typing import Callable, Optional
-from os import environ
-from uuid import UUID
+
+from os import getenv
 from json import dumps
-from traceback import format_exc
 
 import httpx
 from starlette.middleware.base import BaseHTTPMiddleware
-from fastapi import FastAPI, Request, Response
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
 
-from agenta.sdk.utils.logging import log
-from agenta.sdk.middleware.cache import TTLLRUCache
+from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+from agenta.sdk.utils.constants import TRUTHY
+from agenta.sdk.utils.exceptions import display_exception
 
-AGENTA_SDK_AUTH_CACHE_CAPACITY = environ.get(
-    "AGENTA_SDK_AUTH_CACHE_CAPACITY",
-    512,
-)
+import agenta as ag
 
-AGENTA_SDK_AUTH_CACHE_TTL = environ.get(
-    "AGENTA_SDK_AUTH_CACHE_TTL",
-    15 * 60,  # 15 minutes
-)
 
-AGENTA_SDK_AUTH_CACHE = str(environ.get("AGENTA_SDK_AUTH_CACHE", True)).lower() in (
-    "true",
-    "1",
-    "t",
+_SHARED_SERVICE = getenv("AGENTA_SHARED_SERVICE", "false").lower() in TRUTHY
+_CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
+_UNAUTHORIZED_ALLOWED = (
+    getenv("AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED", "false").lower() in TRUTHY
 )
+_ALWAYS_ALLOW_LIST = ["/health"]
 
-AGENTA_SDK_AUTH_CACHE = False
+_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
 
-AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED = str(
-    environ.get("AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED", False)
-).lower() in ("true", "1", "t")
 
+class DenyResponse(JSONResponse):
+    def __init__(
+        self,
+        status_code: int = 401,
+        detail: str = "Unauthorized",
+    ) -> None:
+        super().__init__(
+            status_code=status_code,
+            content={"detail": detail},
+        )
 
-class Deny(Response):
-    def __init__(self) -> None:
-        super().__init__(status_code=401, content="Unauthorized")
 
+class DenyException(Exception):
+    def __init__(
+        self,
+        status_code: int = 401,
+        content: str = "Unauthorized",
+    ) -> None:
+        super().__init__()
 
-cache = TTLLRUCache(
-    capacity=AGENTA_SDK_AUTH_CACHE_CAPACITY,
-    ttl=AGENTA_SDK_AUTH_CACHE_TTL,
-)
+        self.status_code = status_code
+        self.content = content
 
 
-class AuthorizationMiddleware(BaseHTTPMiddleware):
-    def __init__(
-        self,
-        app: FastAPI,
-        host: str,
-        resource_id: UUID,
-        resource_type: str,
-    ):
+class AuthMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
         super().__init__(app)
 
-        self.host = host
-        self.resource_id = resource_id
-        self.resource_type = resource_type
+        self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
+        self.resource_id = (
+            ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.service_id
+            if not _SHARED_SERVICE
+            else None
+        )
+
+    async def dispatch(self, request: Request, call_next: Callable):
+        try:
+            if _UNAUTHORIZED_ALLOWED or request.url.path in _ALWAYS_ALLOW_LIST:
+                request.state.auth = {}
+
+            else:
+                credentials = await self._get_credentials(request)
+
+                request.state.auth = {"credentials": credentials}
 
-    async def dispatch(
-        self,
-        request: Request,
-        call_next: Callable,
-    ):
-        if AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED:
             return await call_next(request)
 
-        try:
-            authorization = (
-                request.headers.get("Authorization")
-                or request.headers.get("authorization")
-                or None
+        except DenyException as deny:
+            display_exception("Auth Middleware Exception")
+
+            return DenyResponse(
+                status_code=deny.status_code,
+                detail=deny.content,
             )
 
+        except:  # pylint: disable=bare-except
+            display_exception("Auth Middleware Exception")
+
+            return DenyResponse(
+                status_code=500,
+                detail="Auth: Unexpected Error.",
+            )
+
+    async def _get_credentials(self, request: Request) -> Optional[str]:
+        try:
+            authorization = request.headers.get("authorization", None)
+
             headers = {"Authorization": authorization} if authorization else None
 
-            cookies = {"sAccessToken": request.cookies.get("sAccessToken")}
+            access_token = request.cookies.get("sAccessToken", None)
+
+            cookies = {"sAccessToken": access_token} if access_token else None
+
+            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+            project_id = (
+                # CLEANEST
+                baggage.get("project_id")
+                # ALTERNATIVE
+                or request.query_params.get("project_id")
+            )
 
-            params = {
-                "action": "run_service",
-                "resource_type": self.resource_type,
-                "resource_id": self.resource_id,
-            }
+            params = {"action": "run_service", "resource_type": "service"}
 
-            project_id = request.query_params.get("project_id")
+            if self.resource_id:
+                params["resource_id"] = self.resource_id
 
             if project_id:
                 params["project_id"] = project_id
@@ -98,48 +124,57 @@ class AuthorizationMiddleware(BaseHTTPMiddleware):
                 sort_keys=True,
             )
 
-            policy = None
-            if AGENTA_SDK_AUTH_CACHE:
-                policy = cache.get(_hash)
-
-            if not policy:
-                async with httpx.AsyncClient() as client:
-                    response = await client.get(
-                        f"{self.host}/api/permissions/verify",
-                        headers=headers,
-                        cookies=cookies,
-                        params=params,
-                    )
+            if _CACHE_ENABLED:
+                credentials = _cache.get(_hash)
+
+                if credentials:
+                    return credentials
 
-                    if response.status_code != 200:
-                        cache.put(_hash, {"effect": "deny"})
-                        return Deny()
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    f"{self.host}/api/permissions/verify",
+                    headers=headers,
+                    cookies=cookies,
+                    params=params,
+                )
 
-                    auth = response.json()
+                if response.status_code == 401:
+                    raise DenyException(
+                        status_code=401,
+                        content="Invalid credentials",
+                    )
+                elif response.status_code == 403:
+                    raise DenyException(
+                        status_code=403,
+                        content="Service execution not allowed.",
+                    )
+                elif response.status_code != 200:
+                    raise DenyException(
+                        status_code=400,
+                        content="Auth: Unexpected Error.",
+                    )
 
-                    if auth.get("effect") != "allow":
-                        cache.put(_hash, {"effect": "deny"})
-                        return Deny()
+                auth = response.json()
 
-                    policy = {
-                        "effect": "allow",
-                        "credentials": auth.get("credentials"),
-                    }
+                if auth.get("effect") != "allow":
+                    raise DenyException(
+                        status_code=403,
+                        content="Service execution not allowed.",
+                    )
 
-                    cache.put(_hash, policy)
+                credentials = auth.get("credentials")
 
-            if not policy or policy.get("effect") == "deny":
-                return Deny()
+                _cache.put(_hash, credentials)
 
-            request.state.credentials = policy.get("credentials")
+                return credentials
 
-            return await call_next(request)
+        except DenyException as deny:
+            raise deny
 
-        except:  # pylint: disable=bare-except
-            log.warning("------------------------------------------------------")
-            log.warning("Agenta SDK - handling auth middleware exception below:")
-            log.warning("------------------------------------------------------")
-            log.warning(format_exc().strip("\n"))
-            log.warning("------------------------------------------------------")
+        except Exception as exc:  # pylint: disable=bare-except
+            display_exception("Auth Middleware Exception (suppressed)")
 
-            return Deny()
+            raise DenyException(
+                status_code=500,
+                content="Auth: Unexpected Error.",
+            ) from exc

agenta/sdk/middleware/cache.py

@@ -1,6 +1,10 @@
+from os import getenv
 from time import time
 from collections import OrderedDict
 
+CACHE_CAPACITY = int(getenv("AGENTA_MIDDLEWARE_CACHE_CAPACITY", "512"))
+CACHE_TTL = int(getenv("AGENTA_MIDDLEWARE_CACHE_TTL", str(5 * 60)))  # 5 minutes
+
 
 class TTLLRUCache:
     def __init__(self, capacity: int, ttl: int):