agenta 0.27.3__tar.gz → 0.27.4a0__tar.gz

{agenta-0.27.3 → agenta-0.27.4a0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.27.3
+Version: 0.27.4a0
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering

{agenta-0.27.3 → agenta-0.27.4a0}/agenta/sdk/decorators/routing.py

@@ -14,9 +14,10 @@ from os import environ
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import Body, FastAPI, UploadFile, HTTPException
 
+from agenta.sdk.middleware.auth import AuthorizationMiddleware
 from agenta.sdk.context.routing import routing_context_manager, routing_context
 from agenta.sdk.context.tracing import tracing_context
-from agenta.sdk.router import router as router
+from agenta.sdk.router import router
 from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.utils.logging import log
 from agenta.sdk.types import (
@@ -50,6 +51,9 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+_MIDDLEWARES = True
+
+
 app.include_router(router, prefix="")
 
 
@@ -121,6 +125,26 @@ class entrypoint:
         route_path="",
         config_schema: Optional[BaseModel] = None,
     ):
+        ### --- Update Middleware --- #
+        try:
+            global _MIDDLEWARES  # pylint: disable=global-statement
+
+            if _MIDDLEWARES:
+                app.add_middleware(
+                    AuthorizationMiddleware,
+                    host=ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host,
+                    resource_id=ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.app_id,
+                    resource_type="application",
+                )
+
+                _MIDDLEWARES = False
+
+        except:  # pylint: disable=bare-except
+            log.error("------------------------------------")
+            log.error("Agenta SDK - failed to secure route: %s", route_path)
+            log.error("------------------------------------")
+        ### --- Update Middleware --- #
+
         DEFAULT_PATH = "generate"
         PLAYGROUND_PATH = "/playground"
         RUN_PATH = "/run"
@@ -330,9 +354,9 @@ class entrypoint:
         *args,
         **func_params,
     ):
-        log.info(f"---------------------------")
+        log.info("---------------------------")
         log.info(f"Agenta SDK - running route: {repr(self.route_path or '/')}")
-        log.info(f"---------------------------")
+        log.info("---------------------------")
 
         tracing_context.set(routing_context.get())
 

agenta-0.27.4a0/agenta/sdk/middleware/auth.py

@@ -0,0 +1,136 @@
+from typing import Callable, Optional
+from os import environ
+from uuid import UUID
+from json import dumps
+from traceback import format_exc
+
+import httpx
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import FastAPI, Request, Response
+
+from agenta.sdk.utils.logging import log
+from agenta.sdk.middleware.cache import TTLLRUCache
+
+AGENTA_SDK_AUTH_CACHE_CAPACITY = environ.get(
+    "AGENTA_SDK_AUTH_CACHE_CAPACITY",
+    512,
+)
+
+AGENTA_SDK_AUTH_CACHE_TTL = environ.get(
+    "AGENTA_SDK_AUTH_CACHE_TTL",
+    15 * 60,  # 15 minutes
+)
+
+AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED = str(
+    environ.get("AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED", False)
+).lower() in ("true", "1", "t")
+
+
+class Deny(Response):
+    def __init__(self) -> None:
+        super().__init__(status_code=401, content="Unauthorized")
+
+
+cache = TTLLRUCache(
+    capacity=AGENTA_SDK_AUTH_CACHE_CAPACITY,
+    ttl=AGENTA_SDK_AUTH_CACHE_TTL,
+)
+
+
+class AuthorizationMiddleware(BaseHTTPMiddleware):
+    def __init__(
+        self,
+        app: FastAPI,
+        host: str,
+        resource_id: UUID,
+        resource_type: str,
+    ):
+        super().__init__(app)
+
+        self.host = host
+        self.resource_id = resource_id
+        self.resource_type = resource_type
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+        project_id: Optional[UUID] = None,
+    ):
+        if AGENTA_UNAUTHORIZED_EXECUTION_ALLOWED:
+            return await call_next(request)
+
+        try:
+            authorization = (
+                request.headers.get("Authorization")
+                or request.headers.get("authorization")
+                or None
+            )
+
+            headers = {"Authorization": authorization} if authorization else None
+
+            cookies = {"sAccessToken": request.cookies.get("sAccessToken")}
+
+            params = {
+                "action": "run_service",
+                "resource_type": self.resource_type,
+                "resource_id": self.resource_id,
+            }
+
+            if project_id:
+                params["project_id"] = project_id
+
+            _hash = dumps(
+                {
+                    "headers": headers,
+                    "cookies": cookies,
+                    "params": params,
+                },
+                sort_keys=True,
+            )
+
+            cached_policy = cache.get(_hash)
+
+            if not cached_policy:
+                async with httpx.AsyncClient() as client:
+                    response = await client.get(
+                        f"{self.host}/api/permissions/verify",
+                        headers=headers,
+                        cookies=cookies,
+                        params=params,
+                    )
+
+                    if response.status_code != 200:
+                        cache.put(_hash, {"effect": "deny"})
+                        return Deny()
+
+                    auth = response.json()
+
+                    if auth.get("effect") != "allow":
+                        cache.put(_hash, {"effect": "deny"})
+                        return Deny()
+
+                    cached_policy = {
+                        "effect": "allow",
+                        "credentials": auth.get("credentials"),
+                    }
+
+                    cache.put(_hash, cached_policy)
+
+            if cached_policy.get("effect") == "deny":
+                return Deny()
+
+            request.state.credentials = cached_policy.get("credentials")
+
+            print(f"credentials: {request.state.credentials}")
+
+            return await call_next(request)
+
+        except:  # pylint: disable=bare-except
+            log.error("------------------------------------------------------")
+            log.error("Agenta SDK - handling auth middleware exception below:")
+            log.error("------------------------------------------------------")
+            log.error(format_exc().strip("\n"))
+            log.error("------------------------------------------------------")
+
+            return Deny()

agenta-0.27.4a0/agenta/sdk/middleware/cache.py

@@ -0,0 +1,43 @@
+from time import time
+from collections import OrderedDict
+
+
+class TTLLRUCache:
+    def __init__(self, capacity: int, ttl: int):
+        self.cache = OrderedDict()
+        self.capacity = capacity
+        self.ttl = ttl
+
+    def get(self, key):
+        # CACHE
+        if key not in self.cache:
+            return None
+
+        value, expiry = self.cache[key]
+        # -----
+
+        # TTL
+        if time() > expiry:
+            del self.cache[key]
+
+            return None
+        # ---
+
+        # LRU
+        self.cache.move_to_end(key)
+        # ---
+
+        return value
+
+    def put(self, key, value):
+        # CACHE
+        if key in self.cache:
+            del self.cache[key]
+        # CACHE & LRU
+        elif len(self.cache) >= self.capacity:
+            self.cache.popitem(last=False)
+        # -----------
+
+        # TTL
+        self.cache[key] = (value, time() + self.ttl)
+        # ---

{agenta-0.27.3 → agenta-0.27.4a0}/agenta/sdk/tracing/exporters.py

@@ -58,7 +58,7 @@ class InlineTraceExporter(SpanExporter):
         return trace
 
 
-OTLPSpanExporter._MAX_RETRY_TIMEOUT = 2
+OTLPSpanExporter._MAX_RETRY_TIMEOUT = 2  # pylint: disable=protected-access
 
 ConsoleExporter = ConsoleSpanExporter
 InlineExporter = InlineTraceExporter

{agenta-0.27.3 → agenta-0.27.4a0}/agenta/sdk/tracing/processors.py

@@ -12,7 +12,7 @@ from opentelemetry.sdk.trace.export import (
 
 from agenta.sdk.utils.logging import log
 
-# LOAD CONTEXT, HERE
+# LOAD CONTEXT, HERE !
 
 
 class TraceProcessor(BatchSpanProcessor):

{agenta-0.27.3 → agenta-0.27.4a0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.27.3"
+version = "0.27.4a0"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]