agenta 0.52.6__py3-none-any.whl → 0.63.2__py3-none-any.whl

agenta/sdk/middleware/vault.py

@@ -6,23 +6,47 @@ import httpx
 from fastapi import FastAPI, Request
 from starlette.middleware.base import BaseHTTPMiddleware
 
+from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.utils.constants import TRUTHY
 from agenta.sdk.utils.cache import TTLLRUCache
 from agenta.sdk.utils.exceptions import suppress, display_exception
 from agenta.client.backend.types import SecretDto as SecretDTO
 from agenta.client.backend.types import (
-    StandardProviderKind,
     StandardProviderDto as StandardProviderDTO,
     StandardProviderSettingsDto as StandardProviderSettingsDTO,
 )
 
 import agenta as ag
 
-
-_PROVIDER_KINDS = []
-
-for provider_kind in StandardProviderKind.__args__[0].__args__:  # type: ignore
-    _PROVIDER_KINDS.append(provider_kind)
+log = get_module_logger(__name__)
+
+
+AGENTA_RUNTIME_PREFIX = getenv("AGENTA_RUNTIME_PREFIX", "")
+
+_ALWAYS_ALLOW_LIST = [
+    f"{AGENTA_RUNTIME_PREFIX}/health",
+    f"{AGENTA_RUNTIME_PREFIX}/openapi.json",
+]
+
+_PROVIDER_KINDS = [
+    "openai",
+    "cohere",
+    "anyscale",
+    "deepinfra",
+    "alephalpha",
+    "groq",
+    "mistral",
+    "mistralai",
+    "anthropic",
+    "perplexityai",
+    "togetherai",
+    "openrouter",
+    "gemini",
+]
+
+_AUTH_ENABLED = (
+    getenv("AGENTA_SERVICE_MIDDLEWARE_AUTH_ENABLED", "true").lower() in TRUTHY
+)
 
 _CACHE_ENABLED = (
     getenv("AGENTA_SERVICE_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
@@ -31,12 +55,27 @@ _CACHE_ENABLED = (
 _cache = TTLLRUCache()
 
 
+class DenyException(Exception):
+    def __init__(
+        self,
+        status_code: int = 403,
+        content: str = "Forbidden",
+    ) -> None:
+        super().__init__()
+
+        self.status_code = status_code
+        self.content = content
+
+
 class VaultMiddleware(BaseHTTPMiddleware):
     def __init__(self, app: FastAPI):
         super().__init__(app)
 
         self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
 
+        self.scope_type = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.scope_type
+        self.scope_id = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.scope_id
+
     async def dispatch(
         self,
         request: Request,
@@ -74,8 +113,12 @@ class VaultMiddleware(BaseHTTPMiddleware):
                 return secrets
 
         local_secrets: List[Dict[str, Any]] = []
+        allow_secrets = True
 
         try:
+            if not request.url.path in _ALWAYS_ALLOW_LIST:
+                await self._allow_local_secrets(credentials)
+
             for provider_kind in _PROVIDER_KINDS:
                 provider = provider_kind
                 key_name = f"{provider.upper()}_API_KEY"
@@ -85,7 +128,7 @@ class VaultMiddleware(BaseHTTPMiddleware):
                     continue
 
                 secret = SecretDTO(
-                    kind="provider_kind",  # type: ignore
+                    kind="provider_key",  # type: ignore
                     data=StandardProviderDTO(
                         kind=provider,
                         provider=StandardProviderSettingsDTO(key=key),
@@ -93,6 +136,9 @@ class VaultMiddleware(BaseHTTPMiddleware):
                 )
 
                 local_secrets.append(secret.model_dump())
+        except DenyException as e:  # pylint: disable=bare-except
+            log.warning(f"Agenta [secrets] {e.status_code}: {e.content}")
+            allow_secrets = False
         except:  # pylint: disable=bare-except
             display_exception("Vault: Local Secrets Exception")
 
@@ -101,7 +147,7 @@ class VaultMiddleware(BaseHTTPMiddleware):
         try:
             async with httpx.AsyncClient() as client:
                 response = await client.get(
-                    f"{self.host}/api/vault/v1/secrets",
+                    f"{self.host}/api/vault/v1/secrets/",
                     headers=headers,
                 )
 
@@ -133,6 +179,155 @@ class VaultMiddleware(BaseHTTPMiddleware):
 
         secrets = standard_secrets + custom_secrets
 
-        _cache.put(_hash, {"secrets": secrets})
+        if not allow_secrets:
+            _cache.put(_hash, {"secrets": secrets})
 
         return secrets
+
+    async def _allow_local_secrets(self, credentials):
+        try:
+            if not _AUTH_ENABLED:
+                return
+
+            if not credentials:
+                raise DenyException(
+                    status_code=401,
+                    content="Invalid credentials. Please check your credentials or login again.",
+                )
+
+            # HEADERS
+            headers = {"Authorization": credentials}
+            # PARAMS
+            params = {}
+            ## SCOPE
+            if self.scope_type and self.scope_id:
+                params["scope_type"] = self.scope_type
+                params["scope_id"] = self.scope_id
+            ## ACTION
+            params["action"] = "view_secret"
+            ## RESOURCE
+            params["resource_type"] = "local_secrets"
+
+            _hash = dumps(
+                {
+                    "headers": headers,
+                    "params": params,
+                },
+                sort_keys=True,
+            )
+
+            access = None
+
+            if _CACHE_ENABLED:
+                access = _cache.get(_hash)
+
+                if isinstance(access, Exception):
+                    raise access
+
+            try:
+                async with httpx.AsyncClient() as client:
+                    try:
+                        response = await client.get(
+                            f"{self.host}/api/permissions/verify",
+                            headers=headers,
+                            params=params,
+                            timeout=30.0,
+                        )
+                    except httpx.TimeoutException as exc:
+                        # log.debug(f"Timeout error while verify secrets access: {exc}")
+                        raise DenyException(
+                            status_code=504,
+                            content=f"Could not verify secrets access: connection to {self.host} timed out. Please check your network connection.",
+                        ) from exc
+                    except httpx.ConnectError as exc:
+                        # log.debug(f"Connection error while verify secrets access: {exc}")
+                        raise DenyException(
+                            status_code=503,
+                            content=f"Could not verify secrets access: connection to {self.host} failed. Please check if agenta is available.",
+                        ) from exc
+                    except httpx.NetworkError as exc:
+                        # log.debug(f"Network error while verify secrets access: {exc}")
+                        raise DenyException(
+                            status_code=503,
+                            content=f"Could not verify secrets access: connection to {self.host} failed. Please check your network connection.",
+                        ) from exc
+                    except httpx.HTTPError as exc:
+                        # log.debug(f"HTTP error while verify secrets access: {exc}")
+                        raise DenyException(
+                            status_code=502,
+                            content=f"Could not verify secrets access: connection to {self.host} failed. Please check if agenta is available.",
+                        ) from exc
+
+                    if response.status_code == 401:
+                        # log.debug("Agenta returned 401 - Invalid credentials")
+                        raise DenyException(
+                            status_code=401,
+                            content="Invalid credentials. Please check your credentials or login again.",
+                        )
+                    elif response.status_code == 403:
+                        # log.debug("Agenta returned 403 - Permission denied")
+                        raise DenyException(
+                            status_code=403,
+                            content="Out of credits. Please set your LLM provider API keys or contact support.",
+                        )
+                    elif response.status_code != 200:
+                        # log.debug(
+                        #     f"Agenta returned {response.status_code} - Unexpected status code"
+                        # )
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify secrets access: {self.host} returned unexpected status code {response.status_code}. Please try again later or contact support if the issue persists.",
+                        )
+
+                    try:
+                        auth = response.json()
+                    except ValueError as exc:
+                        # log.debug(f"Agenta returned invalid JSON response: {exc}")
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify secrets access: {self.host} returned unexpected invalid JSON response. Please try again later or contact support if the issue persists.",
+                        ) from exc
+
+                    if not isinstance(auth, dict):
+                        # log.debug(
+                        #     f"Agenta returned invalid response format: {type(auth)}"
+                        # )
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify secrets access: {self.host} returned unexpected invalid response format. Please try again later or contact support if the issue persists.",
+                        )
+
+                    effect = auth.get("effect")
+
+                    access = effect == "allow"
+
+                    if effect != "allow":
+                        # log.debug("Access denied by Agenta - effect: {effect}")
+                        raise DenyException(
+                            status_code=403,
+                            content="Out of credits. Please set your LLM provider API keys or contact support.",
+                        )
+
+                    return
+
+            except DenyException as deny:
+                _cache.put(_hash, deny)
+
+                raise deny
+            except Exception as exc:  # pylint: disable=bare-except
+                # log.debug(
+                #     f"Unexpected error while verifying credentials (remote): {exc}"
+                # )
+                raise DenyException(
+                    status_code=500,
+                    content=f"Could not verify credentials: unexpected error - {str(exc)}. Please try again later or contact support if the issue persists.",
+                ) from exc
+
+        except DenyException as deny:
+            raise deny
+        except Exception as exc:
+            # log.debug(f"Unexpected error while verifying credentials (local): {exc}")
+            raise DenyException(
+                status_code=500,
+                content=f"Could not verify credentials: unexpected error - {str(exc)}. Please try again later or contact support if the issue persists.",
+            ) from exc

agenta/sdk/middlewares/routing/auth.py

@@ -0,0 +1,263 @@
+from typing import Callable, Optional
+from os import getenv
+from json import dumps
+
+import httpx
+
+from starlette.types import ASGIApp
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request
+from fastapi.responses import JSONResponse
+
+from agenta.sdk.utils.logging import get_module_logger
+from agenta.sdk.utils.exceptions import display_exception
+from agenta.sdk.utils.cache import TTLLRUCache
+from agenta.sdk.utils.constants import TRUTHY
+
+# import agenta as ag
+
+
+log = get_module_logger(__name__)
+
+AGENTA_RUNTIME_PREFIX = getenv("AGENTA_RUNTIME_PREFIX", "")
+
+_AUTH_ENABLED = (
+    getenv("AGENTA_SERVICE_MIDDLEWARE_AUTH_ENABLED", "true").lower() in TRUTHY
+)
+
+_CACHE_ENABLED = (
+    getenv("AGENTA_SERVICE_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
+)
+
+_ALWAYS_ALLOW_LIST = [f"{AGENTA_RUNTIME_PREFIX}/health"]
+
+_cache = TTLLRUCache()
+
+
+class DenyResponse(JSONResponse):
+    def __init__(
+        self,
+        status_code: int = 401,
+        detail: str = "Unauthorized",
+    ) -> None:
+        super().__init__(
+            status_code=status_code,
+            content={"detail": detail},
+        )
+
+
+class DenyException(Exception):
+    def __init__(
+        self,
+        status_code: int = 401,
+        content: str = "Unauthorized",
+    ) -> None:
+        super().__init__()
+
+        self.status_code = status_code
+        self.content = content
+
+
+class AuthMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: ASGIApp, **options):
+        super().__init__(app)
+
+        # self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
+
+        # self.scope_type = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.scope_type
+        # self.scope_id = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.scope_id
+
+    async def dispatch(self, request: Request, call_next: Callable):
+        try:
+            if request.url.path in _ALWAYS_ALLOW_LIST:
+                request.state.auth = {}
+
+            else:
+                credentials = await self._get_credentials(request)
+
+                request.state.auth = {"credentials": credentials}
+
+            return await call_next(request)
+
+        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:
+            if not _AUTH_ENABLED:
+                return request.headers.get("authorization", None)
+
+            # HEADERS
+            authorization = request.headers.get("authorization", None)
+            headers = {"Authorization": authorization} if authorization else None
+
+            # COOKIES
+            access_token = request.cookies.get("sAccessToken", None)
+            cookies = {"sAccessToken": access_token} if access_token else None
+
+            # if not headers and not cookies:
+            #     log.debug("No auth header nor auth cookie found in the request")
+
+            # PARAMS
+            params = {}
+            ## PROJECT_ID
+            project_id = (
+                # CLEANEST
+                request.state.otel["baggage"].get("project_id")
+                # ALTERNATIVE
+                or request.query_params.get("project_id")
+            )
+            # if not project_id:
+            #     log.debug("No project ID found in request")
+
+            if project_id:
+                params["project_id"] = project_id
+            ## SCOPE
+            if self.scope_type and self.scope_id:
+                params["scope_type"] = self.scope_type
+                params["scope_id"] = self.scope_id
+            ## ACTION
+            params["action"] = "run_service"
+            ## RESOURCE
+            params["resource_type"] = "service"
+            # params["resource_id"] = None
+
+            _hash = dumps(
+                {
+                    "headers": headers,
+                    "cookies": cookies,
+                    "params": params,
+                },
+                sort_keys=True,
+            )
+
+            if _CACHE_ENABLED:
+                credentials = _cache.get(_hash)
+
+                if credentials:
+                    # log.debug("Using cached credentials")
+                    return credentials
+
+            try:
+                async with httpx.AsyncClient() as client:
+                    try:
+                        response = await client.get(
+                            f"{self.host}/api/permissions/verify",
+                            headers=headers,
+                            cookies=cookies,
+                            params=params,
+                            timeout=30.0,
+                        )
+                    except httpx.TimeoutException as exc:
+                        # log.debug(f"Timeout error while verify credentials: {exc}")
+                        raise DenyException(
+                            status_code=504,
+                            content=f"Could not verify credentials: connection to {self.host} timed out. Please check your network connection.",
+                        ) from exc
+                    except httpx.ConnectError as exc:
+                        # log.debug(f"Connection error while verify credentials: {exc}")
+                        raise DenyException(
+                            status_code=503,
+                            content=f"Could not verify credentials: connection to {self.host} failed. Please check if agenta is available.",
+                        ) from exc
+                    except httpx.NetworkError as exc:
+                        # log.debug(f"Network error while verify credentials: {exc}")
+                        raise DenyException(
+                            status_code=503,
+                            content=f"Could not verify credentials: connection to {self.host} failed. Please check your network connection.",
+                        ) from exc
+                    except httpx.HTTPError as exc:
+                        # log.debug(f"HTTP error while verify credentials: {exc}")
+                        raise DenyException(
+                            status_code=502,
+                            content=f"Could not verify credentials: connection to {self.host} failed. Please check if agenta is available.",
+                        ) from exc
+
+                    if response.status_code == 401:
+                        # log.debug("Agenta returned 401 - Invalid credentials")
+                        raise DenyException(
+                            status_code=401,
+                            content="Invalid credentials. Please check your credentials or login again.",
+                        )
+                    elif response.status_code == 403:
+                        # log.debug("Agenta returned 403 - Permission denied")
+                        raise DenyException(
+                            status_code=403,
+                            content="Permission denied. Please check your permissions or contact your administrator.",
+                        )
+                    elif response.status_code != 200:
+                        # log.debug(
+                        #     f"Agenta returned {response.status_code} - Unexpected status code"
+                        # )
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify credentials: {self.host} returned unexpected status code {response.status_code}. Please try again later or contact support if the issue persists.",
+                        )
+
+                    try:
+                        auth = response.json()
+                    except ValueError as exc:
+                        # log.debug(f"Agenta returned invalid JSON response: {exc}")
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify credentials: {self.host} returned unexpected invalid JSON response. Please try again later or contact support if the issue persists.",
+                        ) from exc
+
+                    if not isinstance(auth, dict):
+                        # log.debug(
+                        #     f"Agenta returned invalid response format: {type(auth)}"
+                        # )
+                        raise DenyException(
+                            status_code=500,
+                            content=f"Could not verify credentials: {self.host} returned unexpected invalid response format. Please try again later or contact support if the issue persists.",
+                        )
+
+                    effect = auth.get("effect")
+                    if effect != "allow":
+                        # log.debug("Access denied by Agenta - effect: {effect}")
+                        raise DenyException(
+                            status_code=403,
+                            content="Permission denied. Please check your permissions or contact your administrator.",
+                        )
+
+                    credentials = auth.get("credentials")
+
+                    # if not credentials:
+                    #     log.debug("No credentials found in the response")
+
+                    _cache.put(_hash, credentials)
+
+                    return credentials
+
+            except DenyException as deny:
+                raise deny
+            except Exception as exc:  # pylint: disable=bare-except
+                # log.debug(
+                #     f"Unexpected error while verifying credentials (remote): {exc}"
+                # )
+                raise DenyException(
+                    status_code=500,
+                    content=f"Could not verify credentials: unexpected error - {str(exc)}. Please try again later or contact support if the issue persists.",
+                ) from exc
+
+        except DenyException as deny:
+            raise deny
+        except Exception as exc:
+            # log.debug(f"Unexpected error while verifying credentials (local): {exc}")
+            raise DenyException(
+                status_code=500,
+                content=f"Could not verify credentials: unexpected error - {str(exc)}. Please try again later or contact support if the issue persists.",
+            ) from exc

agenta/sdk/middlewares/routing/cors.py

@@ -0,0 +1,30 @@
+from os import getenv
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+from fastapi.middleware.cors import CORSMiddleware as BaseCORSMiddleware
+
+from agenta.sdk.utils.constants import TRUTHY
+
+_USE_CORS = getenv("AGENTA_USE_CORS", "enable").lower() in TRUTHY
+
+
+class CORSMiddleware(BaseCORSMiddleware):
+    def __init__(self, app: ASGIApp, **options):
+        self.app = app
+
+        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/middlewares/routing/otel.py

@@ -0,0 +1,29 @@
+from typing import Callable
+
+from starlette.types import ASGIApp
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request
+
+from agenta.sdk.utils.logging import get_module_logger
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.engines.tracing.propagation import extract
+
+
+log = get_module_logger(__name__)
+
+
+class OTelMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next: Callable):
+        request.state.otel = {"baggage": {}, "traceparent": None}
+
+        headers: dict = dict(request.headers)
+
+        if "newrelic" in headers:
+            headers["traceparent"] = None
+
+        with suppress():
+            _, traceparent, baggage = extract(headers)
+
+            request.state.otel = {"baggage": baggage, "traceparent": traceparent}
+
+        return await call_next(request)