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

agenta/sdk/tracing/processors.py

@@ -0,0 +1,117 @@
+from typing import Optional, Dict, List
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import Span
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    ReadableSpan,
+    BatchSpanProcessor,
+    _DEFAULT_MAX_QUEUE_SIZE,
+    _DEFAULT_MAX_EXPORT_BATCH_SIZE,
+)
+
+from agenta.sdk.utils.logging import log
+
+# LOAD CONTEXT, HERE
+
+
+class TraceProcessor(BatchSpanProcessor):
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        references: Dict[str, str] = None,
+        max_queue_size: int = None,
+        schedule_delay_millis: float = None,
+        max_export_batch_size: int = None,
+        export_timeout_millis: float = None,
+    ):
+        super().__init__(
+            span_exporter,
+            _DEFAULT_MAX_QUEUE_SIZE,
+            12 * 60 * 60 * 1000,  # 12 hours
+            _DEFAULT_MAX_EXPORT_BATCH_SIZE,
+            500,  # < 1 second (0.5 seconds)
+        )
+
+        self._registry = dict()
+        self._exporter = span_exporter
+        self.references = references or dict()
+        self.spans: Dict[int, List[ReadableSpan]] = dict()
+
+    def on_start(
+        self,
+        span: Span,
+        parent_context: Optional[Context] = None,
+    ) -> None:
+        # ADD LINKS FROM CONTEXT, HERE
+
+        for key in self.references.keys():
+            span.set_attribute(f"ag.refs.{key}", self.references[key])
+
+        if span.context.trace_id not in self._registry:
+            self._registry[span.context.trace_id] = dict()
+
+        self._registry[span.context.trace_id][span.context.span_id] = True
+
+    def on_end(
+        self,
+        span: ReadableSpan,
+    ):
+        if self.done:
+            return
+
+        if span.context.trace_id not in self.spans:
+            self.spans[span.context.trace_id] = list()
+
+        self.spans[span.context.trace_id].append(span)
+
+        del self._registry[span.context.trace_id][span.context.span_id]
+
+        if len(self._registry[span.context.trace_id]) == 0:
+            self.export(span.context.trace_id)
+
+    def export(
+        self,
+        trace_id: int,
+    ):
+        spans = self.spans[trace_id]
+
+        for span in spans:
+            self.queue.appendleft(span)
+
+        with self.condition:
+            self.condition.notify()
+
+        del self.spans[trace_id]
+
+    def force_flush(
+        self,
+        timeout_millis: int = None,
+    ) -> bool:
+        ret = super().force_flush(timeout_millis)
+
+        if not ret:
+            log.error("--------------------------------------------")
+            log.error("Agenta SDK - skipping export due to timeout.")
+            log.error("--------------------------------------------")
+
+    def is_ready(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> bool:
+        is_ready = True
+
+        try:
+            is_ready = self._exporter.is_ready(trace_id)
+        except:  # pylint: disable=bare-except
+            pass
+
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> Dict[str, ReadableSpan]:
+        trace = self._exporter.fetch(trace_id)  # type: ignore
+
+        return trace

agenta/sdk/tracing/spans.py

@@ -0,0 +1,136 @@
+from typing import Optional, Union, Any, Dict
+
+from opentelemetry.trace import SpanContext
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.sdk.trace import Span
+
+from agenta.sdk.tracing.attributes import serialize
+
+
+class CustomSpan(Span):  # INHERITANCE FOR TYPING ONLY
+    def __init__(
+        self,
+        span: Span,
+    ) -> None:
+        super().__init__(  # INHERITANCE FOR TYPING ONLY
+            name=span.name,
+            context=span.context,
+            parent=span.parent,
+            sampler=span._sampler,
+            trace_config=span._trace_config,
+            resource=span.resource,
+            attributes=span.attributes,
+            events=span.events,
+            links=span.links,
+            kind=span.kind,
+            span_processor=span._span_processor,
+            instrumentation_info=span.instrumentation_info,
+            record_exception=span._record_exception,
+            set_status_on_exception=span._set_status_on_exception,
+            limits=span._limits,
+            instrumentation_scope=span.instrumentation_scope,
+        )
+
+        self._span = span
+
+    ## --- PROXY METHODS --- ##
+
+    def get_span_context(self):
+        return self._span.get_span_context()
+
+    def is_recording(self) -> bool:
+        return self._span.is_recording()
+
+    def update_name(
+        self,
+        name: str,
+    ) -> None:
+        self._span.update_name(name)
+
+    def set_status(
+        self,
+        status: Union[Status, StatusCode],
+        description: Optional[str] = None,
+    ) -> None:
+        self._span.set_status(
+            status=status,
+            description=description,
+        )
+
+    def end(self) -> None:
+        self._span.end()
+
+    ## --- CUSTOM METHODS W/ ATTRIBUTES SERALIZATION --- ##
+
+    def set_attributes(
+        self,
+        attributes: Dict[str, Any],
+        namespace: Optional[str] = None,
+        max_depth: Optional[int] = None,
+    ) -> None:
+        self._span.set_attributes(
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+                max_depth=max_depth,
+            )
+        )
+
+    def set_attribute(
+        self,
+        key: str,
+        value: Any,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self.set_attributes(
+            attributes={key: value},
+            namespace=namespace,
+        )
+
+    def add_event(
+        self,
+        name: str,
+        attributes: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[int] = None,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.add_event(
+            name=name,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+            timestamp=timestamp,
+        )
+
+    def add_link(
+        self,
+        context: SpanContext,
+        attributes: Optional[Dict[str, Any]] = None,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.add_link(
+            context=context,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+        )
+
+    def record_exception(
+        self,
+        exception: BaseException,
+        attributes: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[int] = None,
+        escaped: bool = False,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.record_exception(
+            exception=exception,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+            timestamp=timestamp,
+            escaped=escaped,
+        )

agenta/sdk/tracing/tracing.py

@@ -0,0 +1,233 @@
+from typing import Optional, Any, Dict
+from enum import Enum
+
+from httpx import get as check
+
+from opentelemetry.trace import (
+    get_current_span,
+    set_tracer_provider,
+    get_tracer_provider,
+    Status,
+    StatusCode,
+)
+from opentelemetry.sdk.trace import Span, Tracer, TracerProvider
+from opentelemetry.sdk.resources import Resource
+
+from agenta.sdk.utils.singleton import Singleton
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.logging import log
+from agenta.sdk.tracing.processors import TraceProcessor
+from agenta.sdk.tracing.exporters import InlineExporter, OTLPExporter
+from agenta.sdk.tracing.spans import CustomSpan
+from agenta.sdk.tracing.inline import parse_inline_trace
+from agenta.sdk.tracing.conventions import Reference, is_valid_attribute_key
+
+
+class Tracing(metaclass=Singleton):
+    VERSION = "0.1.0"
+
+    Status = Status
+    StatusCode = StatusCode
+
+    def __init__(
+        self,
+        url: str,
+    ) -> None:
+        # ENDPOINT (OTLP)
+        self.otlp_url = url
+        # HEADERS (OTLP)
+        self.headers: Dict[str, str] = dict()
+        # REFERENCES
+        self.references: Dict[str, str] = dict()
+
+        # TRACER PROVIDER
+        self.tracer_provider: Optional[TracerProvider] = None
+        # TRACE PROCESSORS -- INLINE
+        self.inline: Optional[TraceProcessor] = None
+        # TRACER
+        self.tracer: Optional[Tracer] = None
+        # INLINE SPANS for INLINE TRACES (INLINE PROCESSOR)
+        self.inline_spans: Dict[str, Any] = dict()
+
+    # PUBLIC
+
+    def configure(
+        self,
+        api_key: Optional[str] = None,
+        # DEPRECATING
+        app_id: Optional[str] = None,
+    ):
+        # HEADERS (OTLP)
+        if api_key:
+            self.headers["Authorization"] = api_key
+        # REFERENCES
+        if app_id:
+            self.references["application.id"] = app_id
+
+        # TRACER PROVIDER
+        self.tracer_provider = TracerProvider(
+            resource=Resource(attributes={"service.name": "agenta-sdk"})
+        )
+        # TRACE PROCESSORS -- INLINE
+        self.inline = TraceProcessor(
+            InlineExporter(
+                registry=self.inline_spans,
+            ),
+            references=self.references,
+        )
+        self.tracer_provider.add_span_processor(self.inline)
+        # TRACE PROCESSORS -- OTLP
+        try:
+            log.info("--------------------------------------------")
+            log.info(
+                "Agenta SDK - connecting to otlp receiver at: %s",
+                self.otlp_url,
+            )
+            log.info("--------------------------------------------")
+
+            check(
+                self.otlp_url,
+                headers=self.headers,
+                timeout=1,
+            )
+
+            _otlp = TraceProcessor(
+                OTLPExporter(
+                    endpoint=self.otlp_url,
+                    headers=self.headers,
+                ),
+                references=self.references,
+            )
+
+            self.tracer_provider.add_span_processor(_otlp)
+
+            log.info("Success: traces will be exported.")
+            log.info("--------------------------------------------")
+
+        except:  # pylint: disable=bare-except
+            log.warning("Failure: traces will not be exported.")
+            log.warning("--------------------------------------------")
+
+        # GLOBAL TRACER PROVIDER -- INSTRUMENTATION LIBRARIES
+        set_tracer_provider(self.tracer_provider)
+        # TRACER
+        self.tracer: Tracer = self.tracer_provider.get_tracer("agenta.tracer")
+
+    def get_current_span(self):
+        _span = None
+
+        with suppress():
+            _span = get_current_span()
+
+            if _span.is_recording():
+                return CustomSpan(_span)
+
+        return _span
+
+    def store_internals(
+        self,
+        attributes: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            span.set_attributes(
+                attributes={"internals": attributes},
+                namespace="data",
+            )
+
+    def store_refs(
+        self,
+        refs: Dict[str, str],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in refs.keys():
+                if key in [_.value for _ in Reference.__members__.values()]:
+                    # ADD REFERENCE TO THIS SPAN
+                    span.set_attribute(
+                        key.value if isinstance(key, Enum) else key,
+                        refs[key],
+                        namespace="refs",
+                    )
+
+                    # AND TO ALL SPANS CREATED AFTER THIS ONE
+                    self.references[key] = refs[key]
+                    # TODO: THIS SHOULD BE REPLACED BY A TRACE CONTEXT !!!
+
+    def store_meta(
+        self,
+        meta: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in meta.keys():
+                if is_valid_attribute_key(key):
+                    span.set_attribute(
+                        key,
+                        meta[key],
+                        namespace="meta",
+                    )
+
+    def store_metrics(
+        self,
+        metrics: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in metrics.keys():
+                if is_valid_attribute_key(key):
+                    span.set_attribute(
+                        key,
+                        metrics[key],
+                        namespace="metrics",
+                    )
+
+    def is_inline_trace_ready(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> bool:
+        is_ready = True
+
+        with suppress():
+            if trace_id is not None:
+                is_ready = self.inline.is_ready(trace_id)
+
+        return is_ready
+
+    def get_inline_trace(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        _inline_trace = {}
+
+        with suppress():
+            is_ready = self.inline.is_ready(trace_id)
+
+            if is_ready is True:
+                otel_spans = self.inline.fetch(trace_id)
+
+                if otel_spans:
+                    _inline_trace = parse_inline_trace(otel_spans)
+
+        return _inline_trace
+
+
+def get_tracer(
+    tracing: Tracing,
+) -> Tracer:
+    if tracing is None or tracing.tracer is None or tracing.tracer_provider is None:
+        return get_tracer_provider().get_tracer("default.tracer")
+
+    return tracing.tracer

agenta/sdk/types.py

@@ -1,9 +1,8 @@
 import json
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Any, Union
 
 from pydantic import ConfigDict, BaseModel, HttpUrl
-from dataclasses import dataclass
-from typing import Union
 
 
 @dataclass
@@ -198,3 +197,51 @@ class Context(BaseModel):
     def from_json(cls, json_str: str):
         data = json.loads(json_str)
         return cls(**data)
+
+
+class ReferencesResponse(BaseModel):
+    app_id: Optional[str] = None
+    app_slug: Optional[str] = None
+    variant_id: Optional[str] = None
+    variant_slug: Optional[str] = None
+    variant_version: Optional[int] = None
+    environment_id: Optional[str] = None
+    environment_slug: Optional[str] = None
+    environment_version: Optional[int] = None
+
+    def __str__(self):
+        return str(self.model_dump(exclude_none=True))
+
+
+class LifecyclesResponse(ReferencesResponse):
+    committed_at: Optional[str] = None
+    committed_by: Optional[str] = None
+    committed_by_id: Optional[str] = None
+    deployed_at: Optional[str] = None
+    deployed_by: Optional[str] = None
+    deployed_by_id: Optional[str] = None
+
+    def __str__(self):
+        return self.model_dump_json(indent=4)
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class ConfigurationResponse(LifecyclesResponse):
+    params: Dict[str, Any]
+
+
+class DeploymentResponse(LifecyclesResponse):
+    pass
+
+
+class Prompt(BaseModel):
+    temperature: float
+    model: str
+    max_tokens: int
+    prompt_system: str
+    prompt_user: str
+    top_p: float
+    frequency_penalty: float
+    presence_penalty: float

agenta/sdk/utils/{helper/openai_cost.py → costs.py}

@@ -49,6 +49,9 @@ MODEL_COST_PER_1K_TOKENS = {
     "gpt-35-turbo-16k-completion": 0.004,
     "gpt-35-turbo-16k-0613-completion": 0.004,
     # Others
+    "text-embedding-ada-002": 0.1,
+    "text-ada-002": 0.1,
+    "adav2": 0.1,
     "text-ada-001": 0.0004,
     "ada": 0.0004,
     "text-babbage-001": 0.0005,

agenta/sdk/utils/debug.py

@@ -1,7 +1,7 @@
 from inspect import iscoroutinefunction
 from functools import wraps
 
-from agenta.sdk.tracing.logger import llm_logger as logging
+from agenta.sdk.utils.logging import log
 
 DEBUG = False
 SHIFT = 7
@@ -14,7 +14,7 @@ def debug(shift=1, req=False, res=False, chars=[">", "<"]):
         @wraps(f)
         async def async_log_wrapper(*args, **kwargs):
             if DEBUG:
-                logging.debug(
+                log.debug(
                     " ".join(
                         [
                             chars[0] * shift + " " * (SHIFT - shift),
@@ -26,7 +26,7 @@ def debug(shift=1, req=False, res=False, chars=[">", "<"]):
                 )
             result = await f(*args, **kwargs)
             if DEBUG:
-                logging.debug(
+                log.debug(
                     " ".join(
                         [
                             chars[1] * shift + " " * (SHIFT - shift),
@@ -40,7 +40,7 @@ def debug(shift=1, req=False, res=False, chars=[">", "<"]):
         @wraps(f)
         def log_wrapper(*args, **kwargs):
             if DEBUG:
-                logging.debug(
+                log.debug(
                     " ".join(
                         [
                             chars[0] * shift + " " * (SHIFT - shift),
@@ -52,7 +52,7 @@ def debug(shift=1, req=False, res=False, chars=[">", "<"]):
                 )
             result = f(*args, **kwargs)
             if DEBUG:
-                logging.debug(
+                log.debug(
                     " ".join(
                         [
                             chars[1] * shift + " " * (SHIFT - shift),

agenta/sdk/utils/exceptions.py

@@ -0,0 +1,52 @@
+from contextlib import AbstractContextManager
+from traceback import format_exc
+from functools import wraps
+from inspect import iscoroutinefunction
+
+from agenta.sdk.utils.logging import log
+
+
+class suppress(AbstractContextManager):  # pylint: disable=invalid-name
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        if exc_type is None:
+            return True
+        else:
+            log.error("-------------------------------------------------")
+            log.error("Agenta SDK - suppressing tracing exception below:")
+            log.error("-------------------------------------------------")
+            log.error(format_exc().strip("\n"))
+            log.error("-------------------------------------------------")
+            return True
+
+
+def handle_exceptions():
+    def decorator(func):
+        is_coroutine_function = iscoroutinefunction(func)
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except Exception as e:
+                log.error("--- HANDLING EXCEPTION ---")
+                log.error("--------------------------")
+                raise e
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                log.error("--- HANDLING EXCEPTION ---")
+                log.error("--------------------------")
+                raise e
+
+        return async_wrapper if is_coroutine_function else sync_wrapper
+
+    return decorator

agenta/sdk/utils/globals.py

@@ -1,15 +1,13 @@
 import agenta
 
 
-def set_global(setup=None, config=None, tracing=None):
-    """Allows usage of agenta.config and agenta.setup in the user's code.
+def set_global(config=None, tracing=None):
+    """Allows usage of agenta.config and agenta.tracing in the user's code.
 
     Args:
-        setup: _description_. Defaults to None.
         config: _description_. Defaults to None.
+        tracing: _description_. Defaults to None.
     """
-    if setup is not None:
-        agenta.setup = setup
     if config is not None:
         agenta.config = config
     if tracing is not None:

agenta/sdk/{tracing/logger.py → utils/logging.py}

@@ -1,12 +1,11 @@
 import logging
 
 
-class LLMLogger:
-    def __init__(self, name="LLMLogger", level=logging.DEBUG):
+class Logger:
+    def __init__(self, name="agenta.logger", level=logging.DEBUG):
         self.logger = logging.getLogger(name)
         self.logger.setLevel(level)
 
-        # Add a stream logger to view the logs in the console
         console_handler = logging.StreamHandler()
         self.logger.addHandler(console_handler)
 
@@ -15,5 +14,4 @@ class LLMLogger:
         return self.logger
 
 
-# Initialize llm logger
-llm_logger = LLMLogger().log
+log = Logger().log