agenta 0.27.3__py3-none-any.whl → 0.27.4a1__py3-none-any.whl

agenta/client/backend/types/parent_dto.py

@@ -0,0 +1,21 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+import pydantic
+
+
+class ParentDto(UniversalBaseModel):
+    id: str
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

agenta/client/backend/types/root_dto.py

@@ -0,0 +1,21 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+import pydantic
+
+
+class RootDto(UniversalBaseModel):
+    id: str
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

agenta/client/backend/types/span_dto.py

@@ -0,0 +1,54 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from __future__ import annotations
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .lifecycle_dto import LifecycleDto
+from .root_dto import RootDto
+from .tree_dto import TreeDto
+from .node_dto import NodeDto
+from .parent_dto import ParentDto
+from .time_dto import TimeDto
+from .status_dto import StatusDto
+from .exception_dto import ExceptionDto
+from .link_dto import LinkDto
+from .o_tel_extra_dto import OTelExtraDto
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+from ..core.pydantic_utilities import update_forward_refs
+
+
+class SpanDto(UniversalBaseModel):
+    lifecycle: typing.Optional[LifecycleDto] = None
+    root: RootDto
+    tree: TreeDto
+    node: NodeDto
+    parent: typing.Optional[ParentDto] = None
+    time: TimeDto
+    status: StatusDto
+    exception: typing.Optional[ExceptionDto] = None
+    data: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+    metrics: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+    meta: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+    refs: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+    links: typing.Optional[typing.List[LinkDto]] = None
+    otel: typing.Optional[OTelExtraDto] = None
+    nodes: typing.Optional[
+        typing.Dict[str, typing.Optional["SpanDtoNodesValue"]]
+    ] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow
+
+
+from .span_dto_nodes_value import SpanDtoNodesValue  # noqa: E402
+
+update_forward_refs(SpanDto)

agenta/client/backend/types/span_dto_nodes_value.py

@@ -0,0 +1,9 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from __future__ import annotations
+import typing
+import typing
+
+if typing.TYPE_CHECKING:
+    from .span_dto import SpanDto
+SpanDtoNodesValue = typing.Union["SpanDto", typing.List["SpanDto"]]

agenta/client/backend/types/status_code.py

@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+StatusCode = typing.Union[typing.Literal["UNSET", "OK", "ERROR"], typing.Any]

agenta/client/backend/types/status_dto.py

@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+from .status_code import StatusCode
+import typing
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class StatusDto(UniversalBaseModel):
+    code: StatusCode
+    message: typing.Optional[str] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

agenta/client/backend/types/time_dto.py

@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import datetime as dt
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import typing
+import pydantic
+
+
+class TimeDto(UniversalBaseModel):
+    start: dt.datetime
+    end: dt.datetime
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

agenta/client/backend/types/tree_dto.py

@@ -0,0 +1,23 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .tree_type import TreeType
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class TreeDto(UniversalBaseModel):
+    id: str
+    type: typing.Optional[TreeType] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            extra="allow", frozen=True
+        )  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

agenta/client/backend/types/tree_type.py

@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+TreeType = typing.Literal["invocation"]

agenta/client/backend/variants/client.py

@@ -1104,14 +1104,14 @@ class VariantsClient:
     def configs_delete(
         self,
         *,
-        variant_ref: typing.Optional[ReferenceRequestModel] = OMIT,
+        variant_ref: ReferenceRequestModel,
         application_ref: typing.Optional[ReferenceRequestModel] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> int:
         """
         Parameters
         ----------
-        variant_ref : typing.Optional[ReferenceRequestModel]
+        variant_ref : ReferenceRequestModel
 
         application_ref : typing.Optional[ReferenceRequestModel]
 
@@ -1125,13 +1125,15 @@ class VariantsClient:
 
         Examples
         --------
-        from agenta import AgentaApi
+        from agenta import AgentaApi, ReferenceRequestModel
 
         client = AgentaApi(
             api_key="YOUR_API_KEY",
             base_url="https://yourhost.com/path/to/api",
         )
-        client.variants.configs_delete()
+        client.variants.configs_delete(
+            variant_ref=ReferenceRequestModel(),
+        )
         """
         _response = self._client_wrapper.httpx_client.request(
             "variants/configs/delete",
@@ -1244,14 +1246,14 @@ class VariantsClient:
     def configs_history(
         self,
         *,
-        variant_ref: typing.Optional[ReferenceRequestModel] = OMIT,
+        variant_ref: ReferenceRequestModel,
         application_ref: typing.Optional[ReferenceRequestModel] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> typing.List[ConfigResponseModel]:
         """
         Parameters
         ----------
-        variant_ref : typing.Optional[ReferenceRequestModel]
+        variant_ref : ReferenceRequestModel
 
         application_ref : typing.Optional[ReferenceRequestModel]
 
@@ -1265,13 +1267,15 @@ class VariantsClient:
 
         Examples
         --------
-        from agenta import AgentaApi
+        from agenta import AgentaApi, ReferenceRequestModel
 
         client = AgentaApi(
             api_key="YOUR_API_KEY",
             base_url="https://yourhost.com/path/to/api",
         )
-        client.variants.configs_history()
+        client.variants.configs_history(
+            variant_ref=ReferenceRequestModel(),
+        )
         """
         _response = self._client_wrapper.httpx_client.request(
             "variants/configs/history",
@@ -2504,14 +2508,14 @@ class AsyncVariantsClient:
     async def configs_delete(
         self,
         *,
-        variant_ref: typing.Optional[ReferenceRequestModel] = OMIT,
+        variant_ref: ReferenceRequestModel,
         application_ref: typing.Optional[ReferenceRequestModel] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> int:
         """
         Parameters
         ----------
-        variant_ref : typing.Optional[ReferenceRequestModel]
+        variant_ref : ReferenceRequestModel
 
         application_ref : typing.Optional[ReferenceRequestModel]
 
@@ -2527,7 +2531,7 @@ class AsyncVariantsClient:
         --------
         import asyncio
 
-        from agenta import AsyncAgentaApi
+        from agenta import AsyncAgentaApi, ReferenceRequestModel
 
         client = AsyncAgentaApi(
             api_key="YOUR_API_KEY",
@@ -2536,7 +2540,9 @@ class AsyncVariantsClient:
 
 
         async def main() -> None:
-            await client.variants.configs_delete()
+            await client.variants.configs_delete(
+                variant_ref=ReferenceRequestModel(),
+            )
 
 
         asyncio.run(main())
@@ -2660,14 +2666,14 @@ class AsyncVariantsClient:
     async def configs_history(
         self,
         *,
-        variant_ref: typing.Optional[ReferenceRequestModel] = OMIT,
+        variant_ref: ReferenceRequestModel,
         application_ref: typing.Optional[ReferenceRequestModel] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> typing.List[ConfigResponseModel]:
         """
         Parameters
         ----------
-        variant_ref : typing.Optional[ReferenceRequestModel]
+        variant_ref : ReferenceRequestModel
 
         application_ref : typing.Optional[ReferenceRequestModel]
 
@@ -2683,7 +2689,7 @@ class AsyncVariantsClient:
         --------
         import asyncio
 
-        from agenta import AsyncAgentaApi
+        from agenta import AsyncAgentaApi, ReferenceRequestModel
 
         client = AsyncAgentaApi(
             api_key="YOUR_API_KEY",
@@ -2692,7 +2698,9 @@ class AsyncVariantsClient:
 
 
         async def main() -> None:
-            await client.variants.configs_history()
+            await client.variants.configs_history(
+                variant_ref=ReferenceRequestModel(),
+            )
 
 
         asyncio.run(main())

agenta/sdk/__init__.py

@@ -17,6 +17,8 @@ from .types import (
     FileInputURL,
     BinaryParam,
     Prompt,
+    AgentaNodeDto,
+    AgentaNodesResponse,
 )
 
 from .tracing import Tracing, get_tracer

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())
 
@@ -362,7 +386,7 @@ class entrypoint:
         log.info(f"Agenta SDK - exiting with success: 200")
         log.info(f"----------------------------------")
 
-        return BaseResponse(data=data, trace=trace)
+        return BaseResponse(data=data, tree=trace)
 
     def handle_failure(self, error: Exception):
         log.error("--------------------------------------------------")

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/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/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