agenta 0.70.1__py3-none-any.whl → 0.75.0__py3-none-any.whl

agenta/sdk/decorators/tracing.py

@@ -1,36 +1,27 @@
 # /agenta/sdk/decorators/tracing.py
 
-from typing import Callable, Optional, Any, Dict, List, Union
-
-from opentelemetry import context as otel_context
-from opentelemetry.context import attach, detach
-
-
 from functools import wraps
-from itertools import chain
 from inspect import (
     getfullargspec,
+    isasyncgenfunction,
     iscoroutinefunction,
     isgeneratorfunction,
-    isasyncgenfunction,
 )
+from itertools import chain
+from typing import Any, Callable, Dict, List, Optional, Union
 
-from pydantic import BaseModel
-
-from opentelemetry import baggage
-from opentelemetry.context import attach, detach, get_current
-from opentelemetry.baggage import set_baggage, get_all
-
-from agenta.sdk.utils.logging import get_module_logger
-from agenta.sdk.utils.exceptions import suppress
+import agenta as ag
 from agenta.sdk.contexts.tracing import (
     TracingContext,
     tracing_context_manager,
 )
 from agenta.sdk.tracing.conventions import parse_span_kind
-
-import agenta as ag
-
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.logging import get_module_logger
+from opentelemetry import context as otel_context
+from opentelemetry.baggage import get_all, set_baggage
+from opentelemetry.context import attach, detach, get_current
+from pydantic import BaseModel
 
 log = get_module_logger(__name__)
 
@@ -88,11 +79,12 @@ class instrument:  # pylint: disable=invalid-name
                 with tracing_context_manager(context=TracingContext.get()):
                     # debug_otel_context("[BEFORE STREAM] [BEFORE SETUP]")
 
-                    captured_ctx = otel_context.get_current()
-
                     self._parse_type_and_kind()
 
-                    self._attach_baggage()
+                    baggage_token = self._attach_baggage()
+
+                    # Capture AFTER baggage attach so we do not wipe it later.
+                    captured_ctx = otel_context.get_current()
 
                     ctx = self._get_traceparent()
 
@@ -141,6 +133,7 @@ class instrument:  # pylint: disable=invalid-name
                             otel_context.detach(otel_token)
 
                             # debug_otel_context("[WITHIN STREAM] [AFTER DETACH]")
+                            self._detach_baggage(baggage_token)
 
                 return wrapped_generator()
 
@@ -311,15 +304,43 @@ class instrument:  # pylint: disable=invalid-name
 
     def _attach_baggage(self):
         context = TracingContext.get()
+        otel_ctx = get_current()
+
+        # 1. Propagate any incoming `ag.*` baggage as-is (for example
+        # `ag.meta.session_id`) so all nested spans inherit it.
+        if context.baggage:
+            for k, v in context.baggage.items():
+                if not isinstance(k, str) or not k.startswith("ag."):
+                    continue
+                if v is None:
+                    continue
+                otel_ctx = set_baggage(name=k, value=str(v), context=otel_ctx)
+
+        # 2. Propagate Agenta references in baggage (used for linking traces to
+        # application/variant/environment).
+        if context.references:
+            for k, v in context.references.items():
+                if v is None:
+                    continue
+                if isinstance(v, BaseModel):
+                    try:
+                        v = v.model_dump(mode="json", exclude_none=True)
+                    except Exception:  # pylint: disable=bare-except
+                        pass
+                if isinstance(v, dict):
+                    for field, value in v.items():
+                        otel_ctx = set_baggage(
+                            name=f"ag.refs.{k}.{field}",
+                            value=str(value),
+                            context=otel_ctx,
+                        )
+                    continue
+                otel_ctx = set_baggage(
+                    name=f"ag.refs.{k}", value=str(v), context=otel_ctx
+                )
 
-        references = context.references
-
-        token = None
-        if references:
-            for k, v in references.items():
-                token = attach(baggage.set_baggage(f"ag.refs.{k}", v))
-
-        return token
+        # Attach once so we can reliably detach later.
+        return attach(otel_ctx)
 
     def _detach_baggage(
         self,

agenta/sdk/engines/tracing/inline.py

@@ -957,9 +957,10 @@ def parse_to_agenta_span_dto(
 ########################################
 
 
-from litellm import cost_calculator
 from opentelemetry.sdk.trace import ReadableSpan
 
+from agenta.sdk.utils.lazy import _load_litellm
+
 from agenta.sdk.types import AgentaNodeDto, AgentaNodesResponse
 
 
@@ -1120,6 +1121,12 @@ TYPES_WITH_COSTS = [
 
 
 def calculate_costs(span_idx: Dict[str, SpanDTO]):
+    litellm = _load_litellm()
+    if not litellm:
+        return
+
+    cost_calculator = litellm.cost_calculator
+
     for span in span_idx.values():
         if (
             span.node.type

agenta/sdk/engines/tracing/processors.py

@@ -1,18 +1,17 @@
-from typing import Optional, Dict, List
 from threading import Lock
+from typing import Dict, List, Optional
 
+from agenta.sdk.models.tracing import BaseModel
+from agenta.sdk.utils.logging import get_module_logger
 from opentelemetry.baggage import get_all as get_baggage
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import Span, SpanProcessor
 from opentelemetry.sdk.trace.export import (
-    SpanExporter,
-    ReadableSpan,
     BatchSpanProcessor,
+    ReadableSpan,
+    SpanExporter,
 )
 
-from agenta.sdk.utils.logging import get_module_logger
-from agenta.sdk.engines.tracing.conventions import Reference
-
 log = get_module_logger(__name__)
 
 
@@ -51,15 +50,27 @@ class TraceProcessor(SpanProcessor):
         parent_context: Optional[Context] = None,
     ) -> None:
         for key in self.references.keys():
-            span.set_attribute(f"ag.refs.{key}", self.references[key])
+            ref = self.references[key]
+            if ref is None:
+                continue
+            if isinstance(ref, BaseModel):
+                try:
+                    ref = ref.model_dump(mode="json", exclude_none=True)
+                except Exception:  # pylint: disable=bare-except
+                    pass
+            if isinstance(ref, dict):
+                for field, value in ref.items():
+                    span.set_attribute(f"ag.refs.{key}.{field}", str(value))
+            else:
+                span.set_attribute(f"ag.refs.{key}", str(ref))
 
         baggage = get_baggage(parent_context)
 
-        for key in baggage.keys():
-            if key.startswith("ag.refs."):
-                _key = key.replace("ag.refs.", "")
-                if _key in [_.value for _ in Reference.__members__.values()]:
-                    span.set_attribute(key, baggage[key])
+        # Copy any `ag.*` baggage entries onto the span attributes so they can be
+        # used for filtering and grouping (for example `ag.meta.session_id`).
+        for key, value in baggage.items():
+            if key.startswith("ag."):
+                span.set_attribute(key, value)
 
         trace_id = span.context.trace_id
         span_id = span.context.span_id

agenta/sdk/evaluations/preview/evaluate.py

@@ -126,10 +126,11 @@ async def _upsert_entities(
                 for testcases_data in simple_evaluation_data.testset_steps:
                     if isinstance(testcases_data, List):
                         if all(isinstance(step, Dict) for step in testcases_data):
-                            testset_revision_id = await acreate_testset(
+                            created_revision = await acreate_testset(
                                 data=testcases_data,
                             )
-                            testset_steps[str(testset_revision_id)] = "custom"
+                            if created_revision and created_revision.id:
+                                testset_steps[str(created_revision.id)] = "custom"
 
             simple_evaluation_data.testset_steps = testset_steps
 
@@ -215,15 +216,16 @@ async def _retrieve_entities(
     Dict[UUID, EvaluatorRevision],
 ]:
     testset_revisions: Dict[UUID, TestsetRevision] = {}
-    # for testset_revision_id, origin in simple_evaluation_data.testset_steps.items():
-    #     testset_revision = await retrieve_testset(
-    #         testset_revision_id=testset_revision_id,
-    #     )
-    for testset_id, origin in simple_evaluation_data.testset_steps.items():
+    for testset_ref, origin in simple_evaluation_data.testset_steps.items():
         testset_revision = await aretrieve_testset(
-            testset_id=testset_id,
+            testset_revision_id=testset_ref,
         )
 
+        if not testset_revision or not testset_revision.id:
+            testset_revision = await aretrieve_testset(
+                testset_id=testset_ref,
+            )
+
         if not testset_revision or not testset_revision.id:
             continue
 
@@ -308,6 +310,32 @@ async def aevaluate(
         "────────────────────────────────────────────────────────────────────────────"
     )
 
+    # Normalize testset_steps to revision ids (no JIT transfers in backend)
+    if simple_evaluation_data.testset_steps and isinstance(
+        simple_evaluation_data.testset_steps, dict
+    ):
+        normalized_testset_steps: Dict[str, Origin] = {}
+        for testset_id_str, origin in simple_evaluation_data.testset_steps.items():
+            try:
+                testset_uuid = UUID(str(testset_id_str))
+            except Exception:
+                continue
+
+            testset_revision = await aretrieve_testset(
+                testset_revision_id=testset_uuid,
+            )
+
+            if not testset_revision or not testset_revision.id:
+                # Fallback: treat as testset_id (latest revision)
+                testset_revision = await aretrieve_testset(
+                    testset_id=testset_uuid,
+                )
+
+            if testset_revision and testset_revision.id:
+                normalized_testset_steps[str(testset_revision.id)] = origin
+
+        simple_evaluation_data.testset_steps = normalized_testset_steps
+
     suffix = _timestamp_suffix()
     name = f"{name}{suffix}"
 

agenta/sdk/evaluations/runs.py

@@ -68,7 +68,8 @@ async def acreate(
                 repeats=repeats,
             ),
             #
-            jit={"testsets": True, "evaluators": False},
+            # Default: expect callers to pass testset revision ids; no JIT migration
+            jit={"testsets": False, "evaluators": False},
         )
     )
 

agenta/sdk/litellm/mockllm.py

@@ -2,9 +2,8 @@ from typing import Optional, Protocol, Any
 from os import environ
 from contextlib import contextmanager
 
-import litellm
-
 from agenta.sdk.utils.logging import get_module_logger
+from agenta.sdk.utils.lazy import _load_litellm
 
 from agenta.sdk.litellm.mocks import MOCKS
 from agenta.sdk.contexts.routing import RoutingContext
@@ -81,6 +80,7 @@ async def acompletion(*args, **kwargs):
 
         return MOCKS[mock](*args, **kwargs)
 
+    litellm = _load_litellm(injected=globals().get("litellm"))
     if not litellm:
         raise ValueError("litellm not found")
 

agenta/sdk/managers/config.py

@@ -2,10 +2,10 @@ import json
 from pathlib import Path
 from typing import Optional, Type, TypeVar, Dict, Any, Union
 
-import yaml
 from pydantic import BaseModel
 
 from agenta.sdk.utils.logging import get_module_logger
+from agenta.sdk.utils.lazy import _load_yaml
 from agenta.sdk.managers.shared import SharedManager
 from agenta.sdk.contexts.routing import RoutingContext
 
@@ -174,6 +174,8 @@ class ConfigManager:
         """
         file_path = Path(filename)
 
+        yaml = _load_yaml()
+
         with open(file_path, "r", encoding="utf-8") as file:
             parameters = yaml.safe_load(file)
 

agenta/sdk/managers/secrets.py

@@ -15,10 +15,15 @@ log = get_module_logger(__name__)
 
 class SecretsManager:
     @staticmethod
-    def get_from_route() -> Optional[List[Dict[str, Any]]]:
+    def get_from_route(scope: str = "all") -> Optional[List[Dict[str, Any]]]:
         context = RoutingContext.get()
 
-        secrets = context.secrets
+        if scope == "local":
+            secrets = context.local_secrets
+        elif scope == "vault":
+            secrets = context.vault_secrets
+        else:
+            secrets = context.secrets
 
         if not secrets:
             return []
@@ -140,7 +145,7 @@ class SecretsManager:
         return modified_model
 
     @staticmethod
-    def get_provider_settings(model: str) -> Optional[Dict]:
+    def get_provider_settings(model: str, scope: str = "all") -> Optional[Dict]:
         """
         Builds the LLM request with appropriate kwargs based on the custom provider/model
 
@@ -154,7 +159,7 @@ class SecretsManager:
         request_provider_model = model
 
         # STEP 1: get vault secrets from route context and transform it
-        secrets = SecretsManager.get_from_route()
+        secrets = SecretsManager.get_from_route(scope=scope)
         if not secrets:
             return None
 
@@ -231,7 +236,7 @@ class SecretsManager:
         return provider_settings
 
     @staticmethod
-    async def retrieve_secrets():
+    async def retrieve_secrets() -> tuple[list, list, list]:
         return await get_secrets(
             f"{ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host}/api",
             RunningContext.get().credentials,
@@ -241,14 +246,20 @@ class SecretsManager:
     async def ensure_secrets_in_workflow():
         ctx = RunningContext.get()
 
-        ctx.secrets = await SecretsManager.retrieve_secrets()
+        secrets, vault_secrets, local_secrets = await SecretsManager.retrieve_secrets()
+
+        ctx.secrets = secrets
+        ctx.vault_secrets = vault_secrets
+        ctx.local_secrets = local_secrets
 
         RunningContext.set(ctx)
 
         return ctx.secrets
 
     @staticmethod
-    def get_provider_settings_from_workflow(model: str) -> Optional[Dict]:
+    def get_provider_settings_from_workflow(
+        model: str, scope: str = "all"
+    ) -> Optional[Dict]:
         """
         Builds the LLM request with appropriate kwargs based on the custom provider/model
 
@@ -262,7 +273,13 @@ class SecretsManager:
         request_provider_model = model
 
         # STEP 1: get vault secrets from route context and transform it
-        secrets = RunningContext.get().secrets
+        ctx = RunningContext.get()
+        if scope == "local":
+            secrets = ctx.local_secrets
+        elif scope == "vault":
+            secrets = ctx.vault_secrets
+        else:
+            secrets = ctx.secrets
         if not secrets:
             return None