agenta 0.26.0a0__py3-none-any.whl → 0.27.0__py3-none-any.whl

agenta/sdk/managers/variant.py

@@ -0,0 +1,182 @@
+from typing import Optional
+
+from agenta.sdk.managers.shared import SharedManager
+
+
+class VariantManager(SharedManager):
+    @classmethod
+    def create(
+        cls,
+        *,
+        parameters: dict,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variant = SharedManager.add(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+
+        if variant:
+            variant = SharedManager.commit(
+                parameters=parameters,
+                app_id=app_id,
+                app_slug=app_slug,
+                variant_slug=variant_slug,
+            )
+
+        return variant
+
+    @classmethod
+    async def acreate(
+        cls,
+        *,
+        parameters: dict,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variant = await SharedManager.aadd(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        if variant:
+            variant = await SharedManager.acommit(
+                parameters=parameters,
+                app_id=app_id,
+                app_slug=app_slug,
+                variant_slug=variant_slug,
+            )
+
+        return variant
+
+    @classmethod
+    def commit(
+        cls,
+        *,
+        parameters: dict,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variant = SharedManager.commit(
+            parameters=parameters,
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return variant
+
+    @classmethod
+    async def acommit(
+        cls,
+        *,
+        parameters: dict,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variant = await SharedManager.acommit(
+            parameters=parameters,
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return variant
+
+    @classmethod
+    def delete(
+        cls,
+        *,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        message = SharedManager.delete(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return message
+
+    @classmethod
+    async def adelete(
+        cls,
+        *,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        message = await SharedManager.adelete(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return message
+
+    @classmethod
+    def list(
+        cls,
+        *,
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variants = SharedManager.list(
+            app_id=app_id,
+            app_slug=app_slug,
+        )
+        return variants
+
+    @classmethod
+    async def alist(
+        cls,
+        *,
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variants = await SharedManager.alist(
+            app_id=app_id,
+            app_slug=app_slug,
+        )
+        return variants
+
+    @classmethod
+    def history(
+        cls,
+        *,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variants = SharedManager.history(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return variants
+
+    @classmethod
+    async def ahistory(
+        cls,
+        *,
+        variant_slug: str,
+        #
+        app_id: Optional[str] = None,
+        app_slug: Optional[str] = None,
+    ):
+        variants = await SharedManager.ahistory(
+            app_id=app_id,
+            app_slug=app_slug,
+            variant_slug=variant_slug,
+        )
+        return variants

agenta/sdk/router.py

@@ -1,15 +1,8 @@
 from fastapi import APIRouter
-from .context import get_contexts
 
 router = APIRouter()
 
 
-@router.get("/contexts/")
-def get_all_contexts():
-    contexts = get_contexts()
-    return {"contexts": contexts}
-
-
 @router.get("/health")
 def health():
     return {"status": "ok"}

agenta/sdk/tracing/__init__.py

@@ -0,0 +1 @@
+from .tracing import Tracing, get_tracer

agenta/sdk/tracing/attributes.py

@@ -0,0 +1,141 @@
+from json import loads, dumps
+from typing import Optional, Union, Sequence, Any, Dict
+
+Primitive = Union[str, int, float, bool, bytes]
+PrimitivesSequence = Sequence[Primitive]
+Attribute = Union[Primitive, PrimitivesSequence]
+
+
+def _marshal(
+    unmarshalled: Dict[str, Any],
+    *,
+    parent_key: Optional[str] = "",
+    depth: Optional[int] = 0,
+    max_depth: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    Marshals a dictionary of unmarshalled attributes into a flat dictionary
+
+    Example:
+    unmarshalled = {
+        "ag": {
+            "type": "tree",
+            "node": {
+                "name": "root",
+                "children": [
+                    {
+                        "name": "child1",
+                    },
+                    {
+                        "name": "child2",
+                    }
+                ]
+            }
+        }
+    }
+    marshalled = {
+        "ag.type": "tree",
+        "ag.node.name": "root",
+        "ag.node.children.0.name": "child1",
+        "ag.node.children.1.name": "child2"
+    }
+    """
+    marshalled = {}
+
+    # If max_depth is set and we've reached it,
+    # just return the unmarshalled attributes
+    if max_depth is not None and depth >= max_depth:
+        marshalled[parent_key] = unmarshalled
+        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+        return marshalled
+
+    # Otherwise,
+    # iterate over the unmarshalled attributes and marshall them
+    for key, value in unmarshalled.items():
+        child_key = f"{parent_key}.{key}" if parent_key else key
+
+        if isinstance(value, dict):
+            dict_key = child_key
+
+            marshalled.update(
+                _marshal(
+                    value,
+                    parent_key=dict_key,
+                    depth=depth + 1,
+                    max_depth=max_depth,
+                )
+            )
+        elif isinstance(value, list):
+            if max_depth is not None and depth + 1 >= max_depth:
+                marshalled[child_key] = value
+                # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+            else:
+                for i, item in enumerate(value):
+                    list_key = f"{child_key}.{i}"
+
+                    if isinstance(item, (dict, list)):
+                        marshalled.update(
+                            _marshal(
+                                item,
+                                parent_key=list_key,
+                                depth=depth + 1,
+                                max_depth=max_depth,
+                            )
+                        )
+                    else:
+                        marshalled[list_key] = item
+                        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+        else:
+            marshalled[child_key] = value
+            # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+    return marshalled
+
+
+def _encode_key(
+    namespace: Optional[str] = None,
+    key: str = "",
+) -> str:
+    if namespace is None:
+        return key
+
+    return f"ag.{namespace}.{key}"
+
+
+def _encode_value(
+    value: Any,
+) -> Optional[Attribute]:
+    if value is None:
+        return None
+
+    if isinstance(value, (str, int, float, bool, bytes)):
+        return value
+
+    if isinstance(value, dict) or isinstance(value, list):
+        encoded = dumps(value)
+        value = "@ag.type=json:" + encoded
+        return value
+
+    return repr(value)
+
+
+def serialize(
+    *,
+    namespace: str,
+    attributes: Dict[str, Any],
+    max_depth: Optional[int] = None,
+) -> Dict[str, str]:
+    if not isinstance(attributes, dict):
+        return {}
+
+    _attributes = {
+        k: v
+        for k, v in {
+            _encode_key(namespace, key): _encode_value(value)
+            for key, value in _marshal(attributes, max_depth=max_depth).items()
+        }.items()
+        if v is not None
+    }
+
+    return _attributes

agenta/sdk/tracing/context.py

@@ -0,0 +1,24 @@
+from contextvars import ContextVar
+from contextlib import contextmanager
+from traceback import format_exc
+
+from agenta.sdk.utils.logging import log
+
+tracing_context = ContextVar("tracing_context", default={})
+
+
+@contextmanager
+def tracing_context_manager():
+    _tracing_context = {"health": {"status": "ok"}}
+
+    token = tracing_context.set(_tracing_context)
+    try:
+        yield
+    except Exception as e:
+        log.error("----------------------------------------------")
+        log.error("Agenta SDK - handling tracing exception below:")
+        log.error("----------------------------------------------")
+        log.error(format_exc().strip("\n"))
+        log.error("----------------------------------------------")
+    finally:
+        tracing_context.reset(token)

agenta/sdk/tracing/conventions.py

@@ -0,0 +1,49 @@
+from enum import Enum
+from re import fullmatch
+
+from opentelemetry.trace import SpanKind
+
+
+class Reference(str, Enum):
+    #
+    VARIANT_ID = "variant.id"
+    VARIANT_SLUG = "variant.slug"
+    VARIANT_VERSION = "variant.version"
+    #
+    ENVIRONMENT_ID = "environment.id"
+    ENVIRONMENT_SLUG = "environment.slug"
+    ENVIRONMENT_VERSION = "environment.version"
+    #
+    APPLICATION_ID = "application.id"
+    APPLICATION_SLUG = "application.slug"
+    #
+
+
+_PATTERN = r"[A-Za-z0-9._-]+"
+
+
+def is_valid_attribute_key(
+    string: str,
+):
+    return bool(fullmatch(_PATTERN, string))
+
+
+def parse_span_kind(type: str) -> SpanKind:
+    kind = SpanKind.INTERNAL
+    if type in [
+        "agent",
+        "chain",
+        "workflow",
+    ]:
+        kind = SpanKind.SERVER
+    elif type in [
+        "tool",
+        "embedding",
+        "query",
+        "completion",
+        "chat",
+        "rerank",
+    ]:
+        kind = SpanKind.CLIENT
+
+    return kind

agenta/sdk/tracing/exporters.py

@@ -0,0 +1,65 @@
+from typing import Sequence, Dict, List
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import (
+    ConsoleSpanExporter,
+    SpanExporter,
+    SpanExportResult,
+    ReadableSpan,
+)
+
+from agenta.sdk.utils.exceptions import suppress
+
+
+class InlineTraceExporter(SpanExporter):
+    def __init__(self, registry: Dict[str, List[ReadableSpan]]):
+        self._shutdown = False
+        self._registry = registry
+
+    def export(
+        self,
+        spans: Sequence[ReadableSpan],
+    ) -> SpanExportResult:
+        if self._shutdown:
+            return
+
+        with suppress():
+            for span in spans:
+                trace_id = span.get_span_context().trace_id
+
+                if trace_id not in self._registry:
+                    self._registry[trace_id] = []
+
+                self._registry[trace_id].append(span)
+
+    def shutdown(self) -> None:
+        self._shutdown = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def is_ready(
+        self,
+        trace_id: int,
+    ) -> bool:
+        is_ready = trace_id in self._registry
+
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: int,
+    ) -> List[ReadableSpan]:
+        trace = self._registry.get(trace_id, [])
+
+        if trace_id in self._registry:
+            del self._registry[trace_id]
+
+        return trace
+
+
+OTLPSpanExporter._MAX_RETRY_TIMEOUT = 2
+
+ConsoleExporter = ConsoleSpanExporter
+InlineExporter = InlineTraceExporter
+OTLPExporter = OTLPSpanExporter