agenta 0.27.2a2__py3-none-any.whl → 0.27.4a0__py3-none-any.whl

agenta/sdk/agenta_init.py

@@ -1,6 +1,6 @@
-import os
 import logging
 import toml
+from os import getenv
 from typing import Optional
 from importlib.metadata import version
 
@@ -72,13 +72,13 @@ class AgentaSingleton:
 
         self.host = (
             host
-            or os.environ.get("AGENTA_HOST")
+            or getenv("AGENTA_HOST")
             or config.get("backend_host")
             or config.get("host")
             or "https://cloud.agenta.ai"
         )
 
-        self.app_id = app_id or config.get("app_id") or os.environ.get("AGENTA_APP_ID")
+        self.app_id = app_id or config.get("app_id") or getenv("AGENTA_APP_ID")
         # if not self.app_id:
         #     raise ValueError(
         #         "App ID must be specified. You can provide it in one of the following ways:\n"
@@ -87,9 +87,7 @@ class AgentaSingleton:
         #         "3. As an environment variable 'AGENTA_APP_ID'."
         #     )
 
-        self.api_key = (
-            api_key or os.environ.get("AGENTA_API_KEY") or config.get("api_key")
-        )
+        self.api_key = api_key or getenv("AGENTA_API_KEY") or config.get("api_key")
 
         self.tracing = Tracing(
             url=f"{self.host}/api/observability/v1/otlp/traces",  # type: ignore
@@ -103,15 +101,15 @@ class AgentaSingleton:
 
         self.api = AgentaApi(
             base_url=self.host + "/api",
-            api_key=api_key if api_key else "",
+            api_key=self.api_key if self.api_key else "",
         )
 
         self.async_api = AsyncAgentaApi(
             base_url=self.host + "/api",
-            api_key=api_key if api_key else "",
+            api_key=self.api_key if self.api_key else "",
         )
 
-        self.base_id = os.environ.get("AGENTA_BASE_ID")
+        self.base_id = getenv("AGENTA_BASE_ID")
 
         self.config = Config(
             host=self.host,

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"
@@ -226,7 +250,7 @@ class entrypoint:
             with routing_context_manager(
                 application={
                     "id": app_id,
-                    "slug": kwargs["app"],
+                    "slug": kwargs.get("app"),
                 },
                 variant={
                     "slug": kwargs.get("config"),
@@ -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, tree=trace)
+        return BaseResponse(data=data, trace=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

agenta/sdk/tracing/inline.py

@@ -903,9 +903,9 @@ def parse_to_agenta_span_dto(
     if span_dto.data:
         span_dto.data = _unmarshal_attributes(span_dto.data)
 
-        if "outputs" in span_dto.data:
-            if "__default__" in span_dto.data["outputs"]:
-                span_dto.data["outputs"] = span_dto.data["outputs"]["__default__"]
+        # if "outputs" in span_dto.data:
+        #     if "__default__" in span_dto.data["outputs"]:
+        #         span_dto.data["outputs"] = span_dto.data["outputs"]["__default__"]
 
     # METRICS
     if span_dto.metrics:
@@ -934,17 +934,6 @@ def parse_to_agenta_span_dto(
             else:
                 parse_to_agenta_span_dto(v)
 
-    # MASK LINKS FOR NOW
-    span_dto.links = None
-    # ------------------
-
-    # MASK LIFECYCLE FOR NOW
-    # span_dto.lifecycle = None
-    if span_dto.lifecycle:
-        span_dto.lifecycle.updated_at = None
-        span_dto.lifecycle.updated_by_id = None
-    # ----------------------
-
     return span_dto
 
 
@@ -956,8 +945,6 @@ def parse_to_agenta_span_dto(
 from litellm import cost_calculator
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agenta.sdk.types import AgentaNodeDto, AgentaNodesResponse
-
 
 def parse_inline_trace(
     spans: Dict[str, ReadableSpan],
@@ -1005,19 +992,51 @@ def parse_inline_trace(
     ### services.observability.service.query() ###
     ##############################################
 
-    spans = [
-        loads(
-            span_dto.model_dump_json(
-                exclude_none=True,
-                exclude_defaults=True,
-            )
-        )
-        for span_dto in agenta_span_dtos
-    ]
-    inline_trace = AgentaNodesResponse(
-        version="1.0.0",
-        nodes=[AgentaNodeDto(**span) for span in spans],
-    ).model_dump(exclude_none=True, exclude_unset=True)
+    LEGACY = True
+    inline_trace = None
+
+    if LEGACY:
+        legacy_spans = [
+            _parse_to_legacy_span(span_dto) for span_dto in span_idx.values()
+        ]
+
+        root_span = agenta_span_dtos[0]
+
+        trace_id = root_span.root.id.hex
+        latency = root_span.time.span / 1_000_000
+        cost = root_span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
+        tokens = {
+            "prompt_tokens": root_span.metrics.get("acc", {})
+            .get("tokens", {})
+            .get("prompt", 0),
+            "completion_tokens": root_span.metrics.get("acc", {})
+            .get("tokens", {})
+            .get("completion", 0),
+            "total_tokens": root_span.metrics.get("acc", {})
+            .get("tokens", {})
+            .get("total", 0),
+        }
+
+        spans = [
+            loads(span.model_dump_json(exclude_none=True)) for span in legacy_spans
+        ]
+
+        inline_trace = {
+            "trace_id": trace_id,
+            "latency": latency,
+            "cost": cost,
+            "usage": tokens,
+            "spans": spans,
+        }
+
+    else:
+        spans = [
+            loads(span_dto.model_dump_json(exclude_none=True))
+            for span_dto in agenta_span_dtos
+        ]
+
+        inline_trace = spans  # turn into Agenta Model ?
+
     return inline_trace
 
 
@@ -1101,6 +1120,98 @@ class LlmTokens(BaseModel):
     total_tokens: Optional[int] = 0
 
 
+class CreateSpan(BaseModel):
+    id: str
+    app_id: str
+    variant_id: Optional[str] = None
+    variant_name: Optional[str] = None
+    inputs: Optional[Dict[str, Optional[Any]]] = None
+    internals: Optional[Dict[str, Optional[Any]]] = None
+    outputs: Optional[Union[str, Dict[str, Optional[Any]], List[Any]]] = None
+    config: Optional[Dict[str, Optional[Any]]] = None
+    environment: Optional[str] = None
+    tags: Optional[List[str]] = None
+    token_consumption: Optional[int] = None
+    name: str
+    parent_span_id: Optional[str] = None
+    attributes: Optional[Dict[str, Optional[Any]]] = None
+    spankind: str
+    status: str
+    user: Optional[str] = None
+    start_time: datetime
+    end_time: datetime
+    tokens: Optional[LlmTokens] = None
+    cost: Optional[float] = None
+
+
+def _parse_to_legacy_span(span: SpanDTO) -> CreateSpan:
+    attributes = None
+    if span.otel:
+        attributes = span.otel.attributes or {}
+
+        if span.otel.events:
+            for event in span.otel.events:
+                if event.name == "exception":
+                    attributes.update(**event.attributes)
+
+    legacy_span = CreateSpan(
+        id=span.node.id.hex[:24],
+        spankind=span.node.type,
+        name=span.node.name,
+        #
+        status=span.status.code.name,
+        #
+        start_time=span.time.start,
+        end_time=span.time.end,
+        #
+        parent_span_id=span.parent.id.hex[:24] if span.parent else None,
+        #
+        inputs=span.data.get("inputs") if span.data else {},
+        internals=span.data.get("internals") if span.data else {},
+        outputs=span.data.get("outputs") if span.data else {},
+        #
+        environment=span.meta.get("environment") if span.meta else None,
+        config=span.meta.get("configuration") if span.meta else None,
+        #
+        tokens=(
+            LlmTokens(
+                prompt_tokens=span.metrics.get("acc", {})
+                .get("tokens", {})
+                .get("prompt", 0.0),
+                completion_tokens=span.metrics.get("acc", {})
+                .get("tokens", {})
+                .get("completion", 0.0),
+                total_tokens=span.metrics.get("acc", {})
+                .get("tokens", {})
+                .get("total", 0.0),
+            )
+            if span.metrics
+            else None
+        ),
+        cost=(
+            span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
+            if span.metrics
+            else None
+        ),
+        #
+        app_id=(
+            span.refs.get("application", {}).get("id", "missing-app-id")
+            if span.refs
+            else "missing-app-id"
+        ),
+        #
+        attributes=attributes,
+        #
+        variant_id=None,
+        variant_name=None,
+        tags=None,
+        token_consumption=None,
+        user=None,
+    )
+
+    return legacy_span
+
+
 TYPES_WITH_COSTS = [
     "embedding",
     "query",

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/sdk/types.py

@@ -4,9 +4,6 @@ from typing import Dict, List, Optional, Any, Union
 
 from pydantic import ConfigDict, BaseModel, HttpUrl
 
-from agenta.client.backend.types.agenta_node_dto import AgentaNodeDto
-from agenta.client.backend.types.agenta_nodes_response import AgentaNodesResponse
-
 
 @dataclass
 class MultipleChoice:
@@ -26,9 +23,9 @@ class LLMTokenUsage(BaseModel):
 
 
 class BaseResponse(BaseModel):
-    version: Optional[str] = "3.0"
+    version: Optional[str] = "2.0"
     data: Optional[Union[str, Dict[str, Any]]]
-    tree: AgentaNodesResponse
+    trace: Optional[Dict[str, Any]]
 
 
 class DictInput(dict):

{agenta-0.27.2a2.dist-info → agenta-0.27.4a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.27.2a2
+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