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

agenta/sdk/middleware/config.py

@@ -0,0 +1,254 @@
+from typing import Callable, Optional, Tuple, Dict
+
+from os import getenv
+from json import dumps
+
+from pydantic import BaseModel
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+import httpx
+
+from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+from agenta.sdk.utils.constants import TRUTHY
+from agenta.sdk.utils.exceptions import suppress
+
+import agenta as ag
+
+
+_CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
+
+_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
+
+
+class Reference(BaseModel):
+    id: Optional[str] = None
+    slug: Optional[str] = None
+    version: Optional[str] = None
+
+
+class ConfigMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+        self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
+        self.application_id = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.app_id
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.config = {}
+
+        with suppress():
+            parameters, references = await self._get_config(request)
+
+            request.state.config = {
+                "parameters": parameters,
+                "references": references,
+            }
+
+        return await call_next(request)
+
+    # @atimeit
+    async def _get_config(self, request: Request) -> Optional[Tuple[Dict, Dict]]:
+        credentials = request.state.auth.get("credentials")
+
+        headers = None
+        if credentials:
+            headers = {"Authorization": credentials}
+
+        application_ref = await self._parse_application_ref(request)
+        variant_ref = await self._parse_variant_ref(request)
+        environment_ref = await self._parse_environment_ref(request)
+
+        refs = {}
+        if application_ref:
+            refs["application_ref"] = application_ref.model_dump()
+        if variant_ref:
+            refs["variant_ref"] = variant_ref.model_dump()
+        if environment_ref:
+            refs["environment_ref"] = environment_ref.model_dump()
+
+        if not refs:
+            return None, None
+
+        _hash = dumps(
+            {
+                "headers": headers,
+                "refs": refs,
+            },
+            sort_keys=True,
+        )
+
+        if _CACHE_ENABLED:
+            config_cache = _cache.get(_hash)
+
+            if config_cache:
+                parameters = config_cache.get("parameters")
+                references = config_cache.get("references")
+
+                return parameters, references
+
+        config = None
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self.host}/api/variants/configs/fetch",
+                headers=headers,
+                json=refs,
+            )
+
+            if response.status_code != 200:
+                return None, None
+
+            config = response.json()
+
+        if not config:
+            _cache.put(_hash, {"parameters": None, "references": None})
+
+            return None, None
+
+        parameters = config.get("params")
+
+        references = {}
+
+        for ref_key in ["application_ref", "variant_ref", "environment_ref"]:
+            refs = config.get(ref_key)
+            ref_prefix = ref_key.split("_", maxsplit=1)[0]
+
+            for ref_part_key in ["id", "slug", "version"]:
+                ref_part = refs.get(ref_part_key)
+
+                if ref_part:
+                    references[ref_prefix + "." + ref_part_key] = ref_part
+
+        _cache.put(_hash, {"parameters": parameters, "references": references})
+
+        return parameters, references
+
+    async def _parse_application_ref(
+        self,
+        request: Request,
+    ) -> Optional[Reference]:
+        baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+        body = {}
+        try:
+            body = await request.json()
+        except:  # pylint: disable=bare-except
+            pass
+
+        application_id = (
+            # CLEANEST
+            baggage.get("application_id")
+            # ALTERNATIVE
+            or request.query_params.get("application_id")
+            # LEGACY
+            or request.query_params.get("app_id")
+            or self.application_id
+        )
+        application_slug = (
+            # CLEANEST
+            baggage.get("application_slug")
+            # ALTERNATIVE
+            or request.query_params.get("application_slug")
+            # LEGACY
+            or request.query_params.get("app_slug")
+            or body.get("app")
+        )
+
+        if not any([application_id, application_slug]):
+            return None
+
+        return Reference(
+            id=application_id,
+            slug=application_slug,
+        )
+
+    async def _parse_variant_ref(
+        self,
+        request: Request,
+    ) -> Optional[Reference]:
+        baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+        body = {}
+        try:
+            body = await request.json()
+        except:  # pylint: disable=bare-except
+            pass
+
+        variant_id = (
+            # CLEANEST
+            baggage.get("variant_id")
+            # ALTERNATIVE
+            or request.query_params.get("variant_id")
+        )
+        variant_slug = (
+            # CLEANEST
+            baggage.get("variant_slug")
+            # ALTERNATIVE
+            or request.query_params.get("variant_slug")
+            # LEGACY
+            or request.query_params.get("config")
+            or body.get("config")
+        )
+        variant_version = (
+            # CLEANEST
+            baggage.get("variant_version")
+            # ALTERNATIVE
+            or request.query_params.get("variant_version")
+        )
+
+        if not any([variant_id, variant_slug, variant_version]):
+            return None
+
+        return Reference(
+            id=variant_id,
+            slug=variant_slug,
+            version=variant_version,
+        )
+
+    async def _parse_environment_ref(
+        self,
+        request: Request,
+    ) -> Optional[Reference]:
+        baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+        body = {}
+        try:
+            body = await request.json()
+        except:  # pylint: disable=bare-except
+            pass
+
+        environment_id = (
+            # CLEANEST
+            baggage.get("environment_id")
+            # ALTERNATIVE
+            or request.query_params.get("environment_id")
+        )
+        environment_slug = (
+            # CLEANEST
+            baggage.get("environment_slug")
+            # ALTERNATIVE
+            or request.query_params.get("environment_slug")
+            # LEGACY
+            or request.query_params.get("environment")
+            or body.get("environment")
+        )
+        environment_version = (
+            # CLEANEST
+            baggage.get("environment_version")
+            # ALTERNATIVE
+            or request.query_params.get("environment_version")
+        )
+
+        if not any([environment_id, environment_slug, environment_version]):
+            return None
+
+        return Reference(
+            id=environment_id,
+            slug=environment_slug,
+            version=environment_version,
+        )

agenta/sdk/middleware/cors.py

@@ -0,0 +1,27 @@
+from os import getenv
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+from fastapi.middleware.cors import CORSMiddleware as _CORSMiddleware
+
+_TRUTHY = {"true", "1", "t", "y", "yes", "on", "enable", "enabled"}
+_USE_CORS = getenv("AGENTA_USE_CORS", "enable").lower() in _TRUTHY
+
+
+class CORSMiddleware(_CORSMiddleware):
+    def __init__(self, app: ASGIApp):
+        if _USE_CORS:
+            super().__init__(
+                app=app,
+                allow_origins=["*"],
+                allow_methods=["*"],
+                allow_headers=["*"],
+                allow_credentials=True,
+                expose_headers=None,
+                max_age=None,
+            )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if _USE_CORS:
+            return await super().__call__(scope, receive, send)
+
+        return await self.app(scope, receive, send)

agenta/sdk/middleware/otel.py

@@ -0,0 +1,40 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+
+from agenta.sdk.utils.exceptions import suppress
+
+
+class OTelMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next: Callable):
+        request.state.otel = {}
+
+        with suppress():
+            baggage = await self._get_baggage(request)
+
+            request.state.otel = {"baggage": baggage}
+
+        return await call_next(request)
+
+    async def _get_baggage(
+        self,
+        request,
+    ):
+        _baggage = {"baggage": request.headers.get("Baggage", "")}
+
+        context = W3CBaggagePropagator().extract(_baggage)
+
+        baggage = {}
+
+        if context:
+            for partial in context.values():
+                for key, value in partial.items():
+                    baggage[key] = value
+
+        return baggage

agenta/sdk/middleware/vault.py

@@ -0,0 +1,158 @@
+from os import getenv
+from json import dumps
+from typing import Callable, Dict, Optional, List, Any
+
+import httpx
+from fastapi import FastAPI, Request
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from agenta.sdk.utils.constants import TRUTHY
+from agenta.client.backend.types.provider_kind import ProviderKind
+from agenta.sdk.utils.exceptions import suppress, display_exception
+from agenta.client.backend.types.secret_dto import SecretDto as SecretDTO
+from agenta.client.backend.types.provider_key_dto import (
+    ProviderKeyDto as ProviderKeyDTO,
+)
+from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+
+import agenta as ag
+
+
+# ProviderKind (agenta.client.backend.types.provider_kind import ProviderKind) defines a type hint that allows \
+# for a fixed set of string literals representing various provider names, alongside `typing.Any`.
+PROVIDER_KINDS = []
+
+# Rationale behind the following:
+# -------------------------------
+# You cannot loop directly over the values in `typing.Literal` because:
+# - `Literal` is not iterable.
+# - `ProviderKind.__args__` includes `Literal` and `Any`, but the actual string values
+#   are nested within the `Literal`'s own `__args__` attribute.
+
+# To solve this, we programmatically extract the values from `Literal` while retaining
+# the structure of ProviderKind. This ensures:
+# 1. We don't modify the original `ProviderKind` type definition.
+# 2. We dynamically access the literal values for use at runtime when necessary.
+for arg in ProviderKind.__args__:  # type: ignore
+    if hasattr(arg, "__args__"):
+        PROVIDER_KINDS.extend(arg.__args__)
+
+
+_CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
+
+_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
+
+
+class VaultMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+        self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
+
+    def _transform_secrets_response_to_secret_dto(
+        self, secrets_list: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        secrets_dto_dict = [
+            {
+                "kind": secret.get("secret", {}).get("kind"),
+                "data": secret.get("secret", {}).get("data", {}),
+            }
+            for secret in secrets_list
+        ]
+        return secrets_dto_dict
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.vault = {}
+
+        with suppress():
+            secrets = await self._get_secrets(request)
+
+            request.state.vault = {"secrets": secrets}
+
+        return await call_next(request)
+
+    async def _get_secrets(self, request: Request) -> Optional[Dict]:
+        credentials = request.state.auth.get("credentials")
+
+        headers = None
+        if credentials:
+            headers = {"Authorization": credentials}
+
+        _hash = dumps(
+            {
+                "headers": headers,
+            },
+            sort_keys=True,
+        )
+
+        if _CACHE_ENABLED:
+            secrets_cache = _cache.get(_hash)
+
+            if secrets_cache:
+                secrets = secrets_cache.get("secrets")
+
+                return secrets
+
+        local_secrets: List[SecretDTO] = []
+
+        try:
+            for provider_kind in PROVIDER_KINDS:
+                provider = provider_kind
+                key_name = f"{provider.upper()}_API_KEY"
+                key = getenv(key_name)
+
+                if not key:
+                    continue
+
+                secret = SecretDTO(  # 'kind' attribute in SecretDTO defaults to 'provider_kind'
+                    data=ProviderKeyDTO(
+                        provider=provider,
+                        key=key,
+                    ),
+                )
+
+                local_secrets.append(secret.model_dump())
+        except:  # pylint: disable=bare-except
+            display_exception("Vault: Local Secrets Exception")
+
+        vault_secrets: List[SecretDTO] = []
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.get(
+                    f"{self.host}/api/vault/v1/secrets",
+                    headers=headers,
+                )
+
+                if response.status_code != 200:
+                    vault_secrets = []
+
+                else:
+                    secrets = response.json()
+                    vault_secrets = self._transform_secrets_response_to_secret_dto(
+                        secrets
+                    )
+        except:  # pylint: disable=bare-except
+            display_exception("Vault: Vault Secrets Exception")
+
+        merged_secrets = {}
+
+        if local_secrets:
+            for secret in local_secrets:
+                provider = secret["data"]["provider"]
+                merged_secrets[provider] = secret
+
+        if vault_secrets:
+            for secret in vault_secrets:
+                provider = secret["data"]["provider"]
+                merged_secrets[provider] = secret
+
+        secrets = list(merged_secrets.values())
+
+        _cache.put(_hash, {"secrets": secrets})
+
+        return secrets

agenta/sdk/tracing/exporters.py

@@ -9,6 +9,11 @@ from opentelemetry.sdk.trace.export import (
 )
 
 from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.context.exporting import (
+    exporting_context_manager,
+    exporting_context,
+    ExportingContext,
+)
 
 
 class InlineTraceExporter(SpanExporter):
@@ -58,8 +63,41 @@ class InlineTraceExporter(SpanExporter):
         return trace
 
 
-OTLPSpanExporter._MAX_RETRY_TIMEOUT = 2  # pylint: disable=protected-access
+class OTLPExporter(OTLPSpanExporter):
+    _MAX_RETRY_TIMEOUT = 2
+
+    def __init__(self, *args, credentials: Dict[int, str] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.credentials = credentials
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        credentials = None
+
+        if self.credentials:
+            trace_ids = set(span.get_span_context().trace_id for span in spans)
+
+            if len(trace_ids) == 1:
+                trace_id = trace_ids.pop()
+
+                if trace_id in self.credentials:
+                    credentials = self.credentials.pop(trace_id)
+
+        with exporting_context_manager(
+            context=ExportingContext(
+                credentials=credentials,
+            )
+        ):
+            return super().export(spans)
+
+    def _export(self, serialized_data: bytes):
+        credentials = exporting_context.get().credentials
+
+        if credentials:
+            self._session.headers.update({"Authorization": credentials})
+
+        return super()._export(serialized_data)
+
 
 ConsoleExporter = ConsoleSpanExporter
 InlineExporter = InlineTraceExporter
-OTLPExporter = OTLPSpanExporter

agenta/sdk/tracing/inline.py

@@ -101,8 +101,8 @@ class NodeDTO(BaseModel):
 Data = Dict[str, Any]
 Metrics = Dict[str, Any]
 Metadata = Dict[str, Any]
-Tags = Dict[str, str]
-Refs = Dict[str, str]
+Tags = Dict[str, Any]
+Refs = Dict[str, Any]
 
 
 class LinkDTO(BaseModel):

agenta/sdk/tracing/processors.py

@@ -1,5 +1,6 @@
 from typing import Optional, Dict, List
 
+from opentelemetry.baggage import get_all as get_baggage
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import Span
 from opentelemetry.sdk.trace.export import (
@@ -11,8 +12,7 @@ from opentelemetry.sdk.trace.export import (
 )
 
 from agenta.sdk.utils.logging import log
-
-# LOAD CONTEXT, HERE !
+from agenta.sdk.tracing.conventions import Reference
 
 
 class TraceProcessor(BatchSpanProcessor):
@@ -43,9 +43,17 @@ class TraceProcessor(BatchSpanProcessor):
         span: Span,
         parent_context: Optional[Context] = None,
     ) -> None:
+        baggage = get_baggage(parent_context)
+
         for key in self.references.keys():
             span.set_attribute(f"ag.refs.{key}", self.references[key])
 
+        for key in baggage.keys():
+            if key.startswith("ag.refs."):
+                _key = key.replace("ag.refs.", "")
+                if _key in [_.value for _ in Reference.__members__.values()]:
+                    span.set_attribute(key, baggage[key])
+
         if span.context.trace_id not in self._registry:
             self._registry[span.context.trace_id] = dict()
 
@@ -89,7 +97,7 @@ class TraceProcessor(BatchSpanProcessor):
         ret = super().force_flush(timeout_millis)
 
         if not ret:
-            log.warning("Agenta SDK - skipping export due to timeout.")
+            log.warning("Agenta - Skipping export due to timeout.")
 
     def is_ready(
         self,

agenta/sdk/tracing/tracing.py

@@ -41,6 +41,8 @@ class Tracing(metaclass=Singleton):
         self.headers: Dict[str, str] = dict()
         # REFERENCES
         self.references: Dict[str, str] = dict()
+        # CREDENTIALS
+        self.credentials: Dict[int, str] = dict()
 
         # TRACER PROVIDER
         self.tracer_provider: Optional[TracerProvider] = None
@@ -60,13 +62,16 @@ class Tracing(metaclass=Singleton):
     def configure(
         self,
         api_key: Optional[str] = None,
+        service_id: Optional[str] = None,
         # DEPRECATING
         app_id: Optional[str] = None,
     ):
         # HEADERS (OTLP)
         if api_key:
-            self.headers["Authorization"] = api_key
+            self.headers["Authorization"] = f"ApiKey {api_key}"
         # REFERENCES
+        if service_id:
+            self.references["service.id"] = service_id
         if app_id:
             self.references["application.id"] = app_id
 
@@ -84,31 +89,28 @@ class Tracing(metaclass=Singleton):
         self.tracer_provider.add_span_processor(self.inline)
         # TRACE PROCESSORS -- OTLP
         try:
-            log.info("--------------------------------------------")
             log.info(
-                "Agenta SDK - connecting to otlp receiver at: %s",
+                "Agenta - OLTP URL: %s",
                 self.otlp_url,
             )
-            log.info("--------------------------------------------")
-            check(
-                self.otlp_url,
-                headers=self.headers,
-                timeout=1,
-            )
+            # check(
+            #     self.otlp_url,
+            #     headers=self.headers,
+            #     timeout=1,
+            # )
 
             _otlp = TraceProcessor(
                 OTLPExporter(
                     endpoint=self.otlp_url,
                     headers=self.headers,
+                    credentials=self.credentials,
                 ),
                 references=self.references,
             )
 
             self.tracer_provider.add_span_processor(_otlp)
-            log.info("Success: traces will be exported.")
-            log.info("--------------------------------------------")
         except:  # pylint: disable=bare-except
-            log.warning("Agenta SDK - traces will not be exported.")
+            log.warning("Agenta - OLTP unreachable, skipping exports.")
 
         # GLOBAL TRACER PROVIDER -- INSTRUMENTATION LIBRARIES
         set_tracer_provider(self.tracer_provider)

agenta/sdk/utils/constants.py

@@ -0,0 +1 @@
+TRUTHY = {"true", "1", "t", "y", "yes", "on", "enable", "enabled"}