agenta 0.12.3__py3-none-any.whl → 0.32.0a1__py3-none-any.whl

agenta/sdk/middleware/inline.py

@@ -0,0 +1,38 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+
+from agenta.sdk.utils.exceptions import suppress
+
+from agenta.sdk.utils.constants import TRUTHY
+
+
+class InlineMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.inline = False
+
+        with suppress():
+            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+            inline = (
+                str(
+                    # CLEANEST
+                    baggage.get("inline")
+                    # ALTERNATIVE
+                    or request.query_params.get("inline")
+                )
+                in TRUTHY
+            )
+
+            request.state.inline = inline
+
+        return await call_next(request)

agenta/sdk/middleware/mock.py

@@ -0,0 +1,33 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+
+from agenta.sdk.utils.exceptions import suppress
+
+
+class MockMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.mock = None
+
+        with suppress():
+            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+            mock = (
+                # CLEANEST
+                baggage.get("mock")
+                # ALTERNATIVE
+                or request.query_params.get("mock")
+            )
+
+            request.state.mock = mock
+
+        return await call_next(request)

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,145 @@
+from os import getenv
+from json import dumps
+from typing import Callable, Dict, Optional, List, Any, get_args
+
+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
+
+
+_PROVIDER_KINDS = []
+
+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=...  # 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/router.py

@@ -1,15 +1,8 @@
 from fastapi import APIRouter
-from .context import get_contexts
 
 router = APIRouter()
 
 
-@router.get("/contexts/")
-def get_all_contexts():
-    contexts = get_contexts()
-    return {"contexts": contexts}
-
-
 @router.get("/health")
 def health():
     return {"status": "ok"}

agenta/sdk/tracing/__init__.py

@@ -0,0 +1 @@
+from .tracing import Tracing, get_tracer

agenta/sdk/tracing/attributes.py

@@ -0,0 +1,141 @@
+from json import loads, dumps
+from typing import Optional, Union, Sequence, Any, Dict
+
+Primitive = Union[str, int, float, bool, bytes]
+PrimitivesSequence = Sequence[Primitive]
+Attribute = Union[Primitive, PrimitivesSequence]
+
+
+def _marshal(
+    unmarshalled: Dict[str, Any],
+    *,
+    parent_key: Optional[str] = "",
+    depth: Optional[int] = 0,
+    max_depth: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    Marshals a dictionary of unmarshalled attributes into a flat dictionary
+
+    Example:
+    unmarshalled = {
+        "ag": {
+            "type": "tree",
+            "node": {
+                "name": "root",
+                "children": [
+                    {
+                        "name": "child1",
+                    },
+                    {
+                        "name": "child2",
+                    }
+                ]
+            }
+        }
+    }
+    marshalled = {
+        "ag.type": "tree",
+        "ag.node.name": "root",
+        "ag.node.children.0.name": "child1",
+        "ag.node.children.1.name": "child2"
+    }
+    """
+    marshalled = {}
+
+    # If max_depth is set and we've reached it,
+    # just return the unmarshalled attributes
+    if max_depth is not None and depth >= max_depth:
+        marshalled[parent_key] = unmarshalled
+        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+        return marshalled
+
+    # Otherwise,
+    # iterate over the unmarshalled attributes and marshall them
+    for key, value in unmarshalled.items():
+        child_key = f"{parent_key}.{key}" if parent_key else key
+
+        if isinstance(value, dict):
+            dict_key = child_key
+
+            marshalled.update(
+                _marshal(
+                    value,
+                    parent_key=dict_key,
+                    depth=depth + 1,
+                    max_depth=max_depth,
+                )
+            )
+        elif isinstance(value, list):
+            if max_depth is not None and depth + 1 >= max_depth:
+                marshalled[child_key] = value
+                # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+            else:
+                for i, item in enumerate(value):
+                    list_key = f"{child_key}.{i}"
+
+                    if isinstance(item, (dict, list)):
+                        marshalled.update(
+                            _marshal(
+                                item,
+                                parent_key=list_key,
+                                depth=depth + 1,
+                                max_depth=max_depth,
+                            )
+                        )
+                    else:
+                        marshalled[list_key] = item
+                        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+        else:
+            marshalled[child_key] = value
+            # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+    return marshalled
+
+
+def _encode_key(
+    namespace: Optional[str] = None,
+    key: str = "",
+) -> str:
+    if namespace is None:
+        return key
+
+    return f"ag.{namespace}.{key}"
+
+
+def _encode_value(
+    value: Any,
+) -> Optional[Attribute]:
+    if value is None:
+        return None
+
+    if isinstance(value, (str, int, float, bool, bytes)):
+        return value
+
+    if isinstance(value, dict) or isinstance(value, list):
+        encoded = dumps(value)
+        value = "@ag.type=json:" + encoded
+        return value
+
+    return repr(value)
+
+
+def serialize(
+    *,
+    namespace: str,
+    attributes: Dict[str, Any],
+    max_depth: Optional[int] = None,
+) -> Dict[str, str]:
+    if not isinstance(attributes, dict):
+        return {}
+
+    _attributes = {
+        k: v
+        for k, v in {
+            _encode_key(namespace, key): _encode_value(value)
+            for key, value in _marshal(attributes, max_depth=max_depth).items()
+        }.items()
+        if v is not None
+    }
+
+    return _attributes

agenta/sdk/tracing/conventions.py

@@ -0,0 +1,49 @@
+from enum import Enum
+from re import fullmatch
+
+from opentelemetry.trace import SpanKind
+
+
+class Reference(str, Enum):
+    #
+    VARIANT_ID = "variant.id"
+    VARIANT_SLUG = "variant.slug"
+    VARIANT_VERSION = "variant.version"
+    #
+    ENVIRONMENT_ID = "environment.id"
+    ENVIRONMENT_SLUG = "environment.slug"
+    ENVIRONMENT_VERSION = "environment.version"
+    #
+    APPLICATION_ID = "application.id"
+    APPLICATION_SLUG = "application.slug"
+    #
+
+
+_PATTERN = r"[A-Za-z0-9._-]+"
+
+
+def is_valid_attribute_key(
+    string: str,
+):
+    return bool(fullmatch(_PATTERN, string))
+
+
+def parse_span_kind(type: str) -> SpanKind:
+    kind = SpanKind.INTERNAL
+    if type in [
+        "agent",
+        "chain",
+        "workflow",
+    ]:
+        kind = SpanKind.SERVER
+    elif type in [
+        "tool",
+        "embedding",
+        "query",
+        "completion",
+        "chat",
+        "rerank",
+    ]:
+        kind = SpanKind.CLIENT
+
+    return kind

agenta/sdk/tracing/exporters.py

@@ -0,0 +1,103 @@
+from typing import Sequence, Dict, List
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import (
+    ConsoleSpanExporter,
+    SpanExporter,
+    SpanExportResult,
+    ReadableSpan,
+)
+
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.context.exporting import (
+    exporting_context_manager,
+    exporting_context,
+    ExportingContext,
+)
+
+
+class InlineTraceExporter(SpanExporter):
+    def __init__(self, registry: Dict[str, List[ReadableSpan]]):
+        self._shutdown = False
+        self._registry = registry
+
+    def export(
+        self,
+        spans: Sequence[ReadableSpan],
+    ) -> SpanExportResult:
+        if self._shutdown:
+            return
+
+        with suppress():
+            for span in spans:
+                trace_id = span.get_span_context().trace_id
+
+                if trace_id not in self._registry:
+                    self._registry[trace_id] = []
+
+                self._registry[trace_id].append(span)
+
+    def shutdown(self) -> None:
+        self._shutdown = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def is_ready(
+        self,
+        trace_id: int,
+    ) -> bool:
+        is_ready = trace_id in self._registry
+
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: int,
+    ) -> List[ReadableSpan]:
+        trace = self._registry.get(trace_id, [])
+
+        if trace_id in self._registry:
+            del self._registry[trace_id]
+
+        return trace
+
+
+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