agenta 0.19.4__py3-none-any.whl → 0.19.6__py3-none-any.whl

agenta/__init__.py

@@ -14,6 +14,7 @@ from .sdk.types import (
     BinaryParam,
 )
 
+from .sdk.tracing.logger import llm_logger as logging
 from .sdk.tracing.llm_tracing import Tracing
 from .sdk.decorators.tracing import instrument
 from .sdk.decorators.llm_entrypoint import entrypoint, app

agenta/sdk/agenta_init.py

@@ -239,6 +239,7 @@ def init(
     singleton = AgentaSingleton()
 
     singleton.init(app_id=app_id, host=host, api_key=api_key, config_fname=config_fname)
+
     tracing = Tracing(
         host=singleton.host,  # type: ignore
         app_id=singleton.app_id,  # type: ignore

agenta/sdk/decorators/llm_entrypoint.py

@@ -18,6 +18,7 @@ from fastapi import Body, FastAPI, UploadFile, HTTPException
 import agenta
 from agenta.sdk.context import save_context
 from agenta.sdk.router import router as router
+from agenta.sdk.tracing.logger import llm_logger as logging
 from agenta.sdk.tracing.llm_tracing import Tracing
 from agenta.sdk.decorators.base import BaseDecorator
 from agenta.sdk.types import (
@@ -52,6 +53,12 @@ app.add_middleware(
 app.include_router(router, prefix="")
 
 
+from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
+
+
+logging.setLevel("DEBUG")
+
+
 class entrypoint(BaseDecorator):
     """Decorator class to wrap a function for HTTP POST, terminal exposure and enable tracing.
 
@@ -72,6 +79,7 @@ class entrypoint(BaseDecorator):
         config_params = agenta.config.all()
         ingestible_files = self.extract_ingestible_files(func_signature)
 
+        @debug()
         @functools.wraps(func)
         async def wrapper(*args, **kwargs) -> Any:
             func_params, api_config_params = self.split_kwargs(kwargs, config_params)
@@ -79,7 +87,7 @@ class entrypoint(BaseDecorator):
             agenta.config.set(**api_config_params)
 
             # Set the configuration and environment of the LLM app parent span at run-time
-            agenta.tracing.set_span_attribute(
+            agenta.tracing.update_baggage(
                 {"config": config_params, "environment": "playground"}
             )
 
@@ -90,6 +98,7 @@ class entrypoint(BaseDecorator):
 
             return llm_result
 
+        @debug()
         @functools.wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
             func_params = {
@@ -104,13 +113,14 @@ class entrypoint(BaseDecorator):
                 agenta.config.pull(config_name="default")
 
             # Set the configuration and environment of the LLM app parent span at run-time
-            agenta.tracing.set_span_attribute(
+            agenta.tracing.update_baggage(
                 {"config": config_params, "environment": kwargs["environment"]}
             )
 
             llm_result = await self.execute_function(
                 func, *args, params=func_params, config_params=config_params
             )
+
             return llm_result
 
         self.update_function_signature(
@@ -190,6 +200,7 @@ class entrypoint(BaseDecorator):
             """
             is_coroutine_function = inspect.iscoroutinefunction(func)
             start_time = time.perf_counter()
+
             if is_coroutine_function:
                 result = await func(*args, **func_params["params"])
             else:
@@ -410,7 +421,7 @@ class entrypoint(BaseDecorator):
         agenta.config.set(**args_config_params)
 
         # Set the configuration and environment of the LLM app parent span at run-time
-        agenta.tracing.set_span_attribute(
+        agenta.tracing.update_baggage(
             {"config": agenta.config.all(), "environment": "bash"}
         )
 

agenta/sdk/decorators/tracing.py

@@ -7,6 +7,12 @@ from typing import Any, Callable, Optional
 # Own Imports
 import agenta as ag
 from agenta.sdk.decorators.base import BaseDecorator
+from agenta.sdk.tracing.logger import llm_logger as logging
+from agenta.sdk.tracing.tracing_context import tracing_context, TracingContext
+from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
+
+
+logging.setLevel("DEBUG")
 
 
 class instrument(BaseDecorator):
@@ -41,6 +47,7 @@ class instrument(BaseDecorator):
     def __call__(self, func: Callable[..., Any]):
         is_coroutine_function = inspect.iscoroutinefunction(func)
 
+        @debug()
         @wraps(func)
         async def async_wrapper(*args, **kwargs):
             result = None
@@ -48,34 +55,65 @@ class instrument(BaseDecorator):
             input_dict = {name: value for name, value in zip(func_args, args)}
             input_dict.update(kwargs)
 
-            span = self.tracing.start_span(
-                name=func.__name__,
-                input=input_dict,
-                spankind=self.spankind,
-                config=self.config,
-            )
-
-            try:
-                result = await func(*args, **kwargs)
-                self.tracing.update_span_status(span=span, value="OK")
-                self.tracing.end_span(
-                    outputs=(
-                        {"message": result} if not isinstance(result, dict) else result
-                    )
-                )
-                return result
-
-            except Exception as e:
-                result = {
-                    "message": str(e),
-                    "stacktrace": traceback.format_exc(),
-                }
-                self.tracing.set_span_attribute(
-                    {"traceback_exception": traceback.format_exc()}
+            async def wrapped_func(*args, **kwargs):
+                # logging.debug(" ".join([">..", str(tracing_context.get())]))
+
+                token = None
+                if tracing_context.get() is None:
+                    token = tracing_context.set(TracingContext())
+
+                # logging.debug(" ".join([">>.", str(tracing_context.get())]))
+
+                self.tracing.start_span(
+                    name=func.__name__,
+                    input=input_dict,
+                    spankind=self.spankind,
+                    config=self.config,
                 )
-                self.tracing.update_span_status(span=span, value="ERROR")
-                self.tracing.end_span(outputs=result)
-                raise e
+
+                try:
+                    result = await func(*args, **kwargs)
+
+                    self.tracing.set_status(status="OK")
+                    self.tracing.end_span(
+                        outputs=(
+                            {"message": result}
+                            if not isinstance(result, dict)
+                            else result
+                        )
+                    )
+
+                    # logging.debug(" ".join(["<<.", str(tracing_context.get())]))
+
+                    if token is not None:
+                        tracing_context.reset(token)
+
+                    # logging.debug(" ".join(["<..", str(tracing_context.get())]))
+
+                    return result
+
+                except Exception as e:
+                    result = {
+                        "message": str(e),
+                        "stacktrace": traceback.format_exc(),
+                    }
+
+                    self.tracing.set_attributes(
+                        {"traceback_exception": traceback.format_exc()}
+                    )
+                    self.tracing.set_status(status="ERROR")
+                    self.tracing.end_span(outputs=result)
+
+                    # logging.debug(" ".join(["<<.", str(tracing_context.get())]))
+
+                    if token is not None:
+                        tracing_context.reset(token)
+
+                    # logging.debug(" ".join(["<..", str(tracing_context.get())]))
+
+                    raise e
+
+            return await wrapped_func(*args, **kwargs)
 
         @wraps(func)
         def sync_wrapper(*args, **kwargs):
@@ -84,33 +122,65 @@ class instrument(BaseDecorator):
             input_dict = {name: value for name, value in zip(func_args, args)}
             input_dict.update(kwargs)
 
-            span = self.tracing.start_span(
-                name=func.__name__,
-                input=input_dict,
-                spankind=self.spankind,
-                config=self.config,
-            )
-
-            try:
-                result = func(*args, **kwargs)
-                self.tracing.update_span_status(span=span, value="OK")
-                self.tracing.end_span(
-                    outputs=(
-                        {"message": result} if not isinstance(result, dict) else result
-                    )
-                )
-                return result
-
-            except Exception as e:
-                result = {
-                    "message": str(e),
-                    "stacktrace": traceback.format_exc(),
-                }
-                self.tracing.set_span_attribute(
-                    {"traceback_exception": traceback.format_exc()}
+            def wrapped_func(*args, **kwargs):
+                # logging.debug(" ".join([">..", str(tracing_context.get())]))
+
+                token = None
+                if tracing_context.get() is None:
+                    token = tracing_context.set(TracingContext())
+
+                # logging.debug(" ".join([">>.", str(tracing_context.get())]))
+
+                span = self.tracing.start_span(
+                    name=func.__name__,
+                    input=input_dict,
+                    spankind=self.spankind,
+                    config=self.config,
                 )
-                self.tracing.update_span_status(span=span, value="ERROR")
-                self.tracing.end_span(outputs=result)
-                raise e
+
+                try:
+                    result = func(*args, **kwargs)
+
+                    self.tracing.set_status(status="OK")
+                    self.tracing.end_span(
+                        outputs=(
+                            {"message": result}
+                            if not isinstance(result, dict)
+                            else result
+                        )
+                    )
+
+                    # logging.debug(" ".join(["<<.", str(tracing_context.get())]))
+
+                    if token is not None:
+                        tracing_context.reset(token)
+
+                    # logging.debug(" ".join(["<..", str(tracing_context.get())]))
+
+                    return result
+
+                except Exception as e:
+                    result = {
+                        "message": str(e),
+                        "stacktrace": traceback.format_exc(),
+                    }
+
+                    self.tracing.set_attributes(
+                        {"traceback_exception": traceback.format_exc()}
+                    )
+
+                    self.tracing.set_status(status="ERROR")
+                    self.tracing.end_span(outputs=result)
+
+                    # logging.debug(" ".join(["<<.", str(tracing_context.get())]))
+
+                    if token is not None:
+                        tracing_context.reset(token)
+
+                    # logging.debug(" ".join(["<..", str(tracing_context.get())]))
+
+                    raise e
+
+            return wrapped_func(*args, **kwargs)
 
         return async_wrapper if is_coroutine_function else sync_wrapper

agenta/sdk/tracing/llm_tracing.py

@@ -1,9 +1,12 @@
 import os
+from uuid import uuid4
+
 from threading import Lock
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any, List
 
-from agenta.sdk.tracing.logger import llm_logger
+from agenta.sdk.tracing.tracing_context import tracing_context
+from agenta.sdk.tracing.logger import llm_logger as logging
 from agenta.sdk.tracing.tasks_manager import TaskQueue
 from agenta.client.backend.client import AsyncAgentaApi
 from agenta.client.backend.client import AsyncObservabilityClient
@@ -15,9 +18,13 @@ from agenta.client.backend.types.create_span import (
 
 from bson.objectid import ObjectId
 
-
 VARIANT_TRACKING_FEATURE_FLAG = False
 
+from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
+
+
+logging.setLevel("DEBUG")
+
 
 class SingletonMeta(type):
     """
@@ -54,7 +61,6 @@ class Tracing(metaclass=SingletonMeta):
         host (str): The URL of the backend host
         api_key (str): The API Key of the backend host
         tasks_manager (TaskQueue): The tasks manager dedicated to handling asynchronous tasks
-        llm_logger (Logger): The logger associated with the LLM tracing
         max_workers (int): The maximum number of workers to run tracing
     """
 
@@ -67,19 +73,11 @@ class Tracing(metaclass=SingletonMeta):
     ):
         self.host = host + "/api"
         self.api_key = api_key if api_key is not None else ""
-        self.llm_logger = llm_logger
         self.app_id = app_id
         self.tasks_manager = TaskQueue(
-            max_workers if max_workers else 4, logger=llm_logger
+            max_workers if max_workers else 4, logger=logging
         )
-        self.active_span: Optional[CreateSpan] = None
-        self.active_trace_id: Optional[str] = None
-        self.pending_spans: List[CreateSpan] = []
-        self.tags: List[str] = []
-        self.trace_config_cache: Dict[
-            str, Any
-        ] = {}  # used to save the trace configuration before starting the first span
-        self.span_dict: Dict[str, CreateSpan] = {}  # type: ignore
+        self.baggage = None
 
     @property
     def client(self) -> AsyncObservabilityClient:
@@ -93,30 +91,104 @@ class Tracing(metaclass=SingletonMeta):
             base_url=self.host, api_key=self.api_key, timeout=120  # type: ignore
         ).observability
 
-    def set_span_attribute(
+    ### --- API --- ###
+
+    @debug()
+    def get_context(self):
+        tracing = tracing_context.get()
+
+        return tracing
+
+    @debug()
+    def update_baggage(
         self,
         attributes: Dict[str, Any] = {},
     ):
+        if self.baggage is None:
+            self.baggage = {}
+
+        for key, value in attributes.items():
+            self.baggage[key] = value
+
+    @debug()
+    def open_trace(
+        self,
+        span: Optional[CreateSpan],
+        config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:
+        tracing = tracing_context.get()
+
+        tracing.trace_id = self._create_trace_id()
+
+        logging.info(f"Opening trace {tracing.trace_id}")
+
+        if span is not None:
+            ### --- TO BE CLEANED --- >>>
+            span.environment = (
+                self.baggage.get("environment")
+                if self.baggage is not None
+                else os.environ.get("environment", "unset")
+            )
+
+            span.config = (
+                self.baggage.get("config")
+                if not config and self.baggage is not None
+                else None
+            )
+            if VARIANT_TRACKING_FEATURE_FLAG:
+                # TODO: we should get the variant_id and variant_name (and environment) from the config object
+                span.variant_id = config.variant_id  # type: ignore
+                span.variant_name = (config.variant_name,)  # type: ignore
+            ### --- TO BE CLEANED --- <<<
+
+        logging.info(f"Opened  trace {tracing.trace_id}")
+
+    @debug()
+    def set_trace_tags(self, tags: List[str]) -> None:
+        tracing = tracing_context.get()
+
+        tracing.trace_tags.extend(tags)
+
+    @debug()
+    def close_trace(self) -> None:
         """
-        Set attributes for the active span.
+        Ends the active trace and sends the recorded spans for processing.
 
         Args:
-            attributes (Dict[str, Any], optional): A dictionary of attributes to set. Defaults to {}.
+            parent_span (CreateSpan): The parent span of the trace.
+
+        Raises:
+            RuntimeError: If there is no active trace to end.
+
+        Returns:
+            None
         """
+        tracing = tracing_context.get()
+
+        logging.info(f"Closing trace {tracing.trace_id}")
+
+        trace_id = tracing.trace_id
 
-        if (
-            self.active_span is None
-        ):  # This is the case where entrypoint wants to save the trace information but the parent span has not been initialized yet
-            for key, value in attributes.items():
-                self.trace_config_cache[key] = value
+        if tracing.trace_id is None:
+            logging.error("Cannot close trace, no trace to close")
+            return
+
+        if not self.api_key:
+            logging.error("No API key")
         else:
-            for key, value in attributes.items():
-                self.active_span.attributes[key] = value  # type: ignore
+            self._process_closed_spans()
 
-    def set_trace_tags(self, tags: List[str]):
-        self.tags.extend(tags)
+        self._clear_closed_spans()
+        self._clear_tracked_spans()
+        self._clear_active_span()
 
-    def start_span(
+        self._clear_trace_tags()
+
+        logging.info(f"Closed  trace {trace_id}")
+
+    @debug()
+    def open_span(
         self,
         name: str,
         spankind: str,
@@ -124,10 +196,13 @@ class Tracing(metaclass=SingletonMeta):
         config: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> CreateSpan:
+        tracing = tracing_context.get()
+
         span_id = self._create_span_id()
-        self.llm_logger.info(
-            f"Recording {'parent' if spankind == 'workflow' else spankind} span..."
-        )
+
+        logging.info(f"Opening span  {span_id} {spankind.upper()}")
+
+        ### --- TO BE CLEANED --- >>>
         span = CreateSpan(
             id=span_id,
             inputs=input,
@@ -148,54 +223,66 @@ class Tracing(metaclass=SingletonMeta):
             parent_span_id=None,
         )
 
-        if self.active_trace_id is None:  # This is a parent span
-            self.active_trace_id = self._create_trace_id()
-            span.environment = (
-                self.trace_config_cache.get("environment")
-                if self.trace_config_cache is not None
-                else os.environ.get("environment", "unset")
-            )
-            span.config = (
-                self.trace_config_cache.get("config")
-                if not config and self.trace_config_cache is not None
-                else None
-            )
-            if VARIANT_TRACKING_FEATURE_FLAG:
-                # TODO: we should get the variant_id and variant_name (and environment) from the config object
-                span.variant_id = config.variant_id  # type: ignore
-                span.variant_name = (config.variant_name,)  # type: ignore
-
+        if tracing.trace_id is None:
+            self.start_trace(span, config)
         else:
-            span.parent_span_id = self.active_span.id  # type: ignore
+            span.parent_span_id = tracing.active_span.id  # type: ignore
 
-        self.span_dict[span.id] = span
-        self.active_span = span
+        tracing.tracked_spans[span.id] = span
+        tracing.active_span = span
+        ### --- TO BE CLEANED --- <<<
+
+        logging.info(f"Opened  span  {span_id} {spankind.upper()}")
 
-        self.llm_logger.info(f"Recorded span and setting parent_span_id: {span.id}")
         return span
 
-    def update_span_status(self, span: CreateSpan, value: str):
-        span.status = value
+    @debug(req=True)
+    def set_attributes(
+        self,
+        attributes: Dict[str, Any] = {},
+    ) -> None:
+        """
+        Set attributes for the active span.
 
-    def _update_span_cost(self, span: CreateSpan, cost: Optional[float]):
-        if cost is not None and isinstance(cost, float):
-            if span.cost is None:
-                span.cost = cost
-            else:
-                span.cost += cost
+        Args:
+            attributes (Dict[str, Any], optional): A dictionary of attributes to set. Defaults to {}.
+        """
 
-    def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]):
-        if isinstance(tokens, LlmTokens):
-            tokens = tokens.dict()
-        if tokens is not None and isinstance(tokens, dict):
-            if span.tokens is None:
-                span.tokens = LlmTokens(**tokens)
-            else:
-                span.tokens.prompt_tokens += tokens["prompt_tokens"]
-                span.tokens.completion_tokens += tokens["completion_tokens"]
-                span.tokens.total_tokens += tokens["total_tokens"]
+        tracing = tracing_context.get()
 
-    def end_span(self, outputs: Dict[str, Any]):
+        if tracing.active_span is None:
+            logging.error(f"Cannot set attributes ({set(attributes)}), no active span")
+            return
+
+        logging.info(
+            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} attributes={attributes}"
+        )
+
+        for key, value in attributes.items():
+            tracing.active_span.attributes[key] = value  # type: ignore
+
+    @debug()
+    def set_status(self, status: str) -> None:
+        """
+        Set status for the active span.
+
+        Args:
+            status: Enum ( UNSET, OK, ERROR )
+        """
+        tracing = tracing_context.get()
+
+        if tracing.active_span is None:
+            logging.error(f"Cannot set status ({status}), no active span")
+            return
+
+        logging.info(
+            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} status={status}"
+        )
+
+        tracing.active_span.status = status
+
+    @debug()
+    def close_span(self, outputs: Dict[str, Any]) -> None:
         """
         Ends the active span, if it is a parent span, ends the trace too.
 
@@ -213,123 +300,160 @@ class Tracing(metaclass=SingletonMeta):
             None
         """
 
-        if self.active_span is None:
-            raise ValueError("There is no active span to end.")
+        tracing = tracing_context.get()
+
+        if tracing.active_span is None:
+            logging.error("Cannot close span, no active span")
+
+        span_id = tracing.active_span.id
+        spankind = tracing.active_span.spankind
+
+        logging.info(f"Closing span  {span_id} {spankind}")
+
+        ### --- TO BE CLEANED --- >>>
+        tracing.active_span.end_time = datetime.now(timezone.utc)
+
+        tracing.active_span.outputs = [outputs.get("message", "")]
 
-        self.active_span.end_time = datetime.now(timezone.utc)
-        self.active_span.outputs = [outputs.get("message", "")]
-        if self.active_span.spankind in [
+        if tracing.active_span.spankind.upper() in [
             "LLM",
             "RETRIEVER",
         ]:  # TODO: Remove this whole part. Setting the cost should be done through set_span_attribute
-            self._update_span_cost(self.active_span, outputs.get("cost", None))
-            self._update_span_tokens(self.active_span, outputs.get("usage", None))
+            self._update_span_cost(tracing.active_span, outputs.get("cost", None))
+            self._update_span_tokens(tracing.active_span, outputs.get("usage", None))
 
-        # Push span to list of recorded spans
-        self.pending_spans.append(self.active_span)
+        tracing.closed_spans.append(tracing.active_span)
 
-        active_span_parent_id = self.active_span.parent_span_id
-        if (
-            self.active_span.status == SpanStatusCode.ERROR.value
-            and active_span_parent_id is not None
-        ):
-            self.record_exception_and_end_trace(span_parent_id=active_span_parent_id)
+        active_span_parent_id = tracing.active_span.parent_span_id
 
         if active_span_parent_id is None:
-            self.end_trace(parent_span=self.active_span)
+            self.end_trace(parent_span=tracing.active_span)
 
         else:
-            parent_span = self.span_dict[active_span_parent_id]
-            self._update_span_cost(parent_span, self.active_span.cost)
-            self._update_span_tokens(parent_span, self.active_span.tokens)
-            self.active_span = parent_span
+            parent_span = tracing.tracked_spans[active_span_parent_id]
+            self._update_span_cost(parent_span, tracing.active_span.cost)
+            self._update_span_tokens(parent_span, tracing.active_span.tokens)
+            tracing.active_span = parent_span
+        ### --- TO BE CLEANED --- <<<
 
-    def record_exception_and_end_trace(self, span_parent_id: str):
-        """
-        Record an exception and end the trace.
+        logging.info(f"Closed  span  {span_id} {spankind}")
 
-        Args:
-            span_parent_id (str): The ID of the parent span.
+    ### --- Legacy API --- ###
 
-        Returns:
-            None
-        """
+    def start_trace(
+        self,
+        span: CreateSpan,
+        config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:  # Legacy
+        self.open_trace(span, config, **kwargs)
 
-        parent_span = self.span_dict.get(span_parent_id)
-        if parent_span is not None:
-            # Update parent span of active span
-            parent_span.outputs = self.active_span.outputs  # type: ignore
-            parent_span.status = "ERROR"
-            parent_span.end_time = datetime.now(timezone.utc)
+    def end_trace(self, parent_span: CreateSpan) -> None:  # Legacy
+        self.close_trace()
 
-            # Push parent span to list of recorded spans and end trace
-            self.pending_spans.append(parent_span)
-            self.end_trace(parent_span=parent_span)
+    def start_span(
+        self,
+        name: str,
+        spankind: str,
+        input: Dict[str, Any],
+        config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> CreateSpan:  # Legacy
+        return self.open_span(name, spankind, input, config, **kwargs)
 
-        # TODO: improve exception logic here.
+    def update_span_status(self, _: CreateSpan, status: str) -> None:  # Legacy
+        self.update_span_status(status)
 
-    def end_trace(self, parent_span: CreateSpan):
+    def set_span_attribute(
+        self,
+        attributes: Dict[str, Any] = {},
+    ) -> None:  # Legacy
+        self.set_attributes(attributes)
+
+    def end_span(self, outputs: Dict[str, Any]) -> None:  # Legacy
+        self.close_span(outputs)
+
+    ### --- Helper Functions --- ###
+
+    def _create_trace_id(self) -> str:
+        """Creates a 32HEXDIGL / ObjectId ID for the trace object.
+
+        Returns:
+            str: stringify oid of the trace
         """
-        Ends the active trace and sends the recorded spans for processing.
 
-        Args:
-            parent_span (CreateSpan): The parent span of the trace.
+        # return uuid4().hex
+        return str(ObjectId())
 
-        Raises:
-            RuntimeError: If there is no active trace to end.
+    def _clear_trace_tags(self) -> None:
+        tracing = tracing_context.get()
+
+        tracing.trace_tags.clear()
+
+    def _create_span_id(self) -> str:
+        """Creates a  16HEXDIGL / ObjectId ID for the span object.
 
         Returns:
-            None
+            str: stringify oid of the span
         """
 
-        if self.api_key == "":
-            return
+        # return uuid4().hex[:16]
+        return str(ObjectId())
+
+    def _process_closed_spans(self) -> None:
+        tracing = tracing_context.get()
+
+        logging.info(f"Sending spans {tracing.trace_id} #={len(tracing.closed_spans)} ")
 
-        if not self.active_trace_id:
-            raise RuntimeError("No active trace to end.")
+        # async def mock_create_traces(trace, spans):
+        #    print("trace-id", trace)
+        #    print("spans", spans)
 
-        self.llm_logger.info("Preparing to send recorded spans for processing.")
-        self.llm_logger.info(f"Recorded spans => {len(self.pending_spans)}")
         self.tasks_manager.add_task(
-            self.active_trace_id,
+            tracing.trace_id,
             "trace",
+            # mock_create_traces(
             self.client.create_traces(
-                trace=self.active_trace_id, spans=self.pending_spans  # type: ignore
+                trace=tracing.trace_id, spans=tracing.closed_spans  # type: ignore
             ),
             self.client,
         )
-        self.llm_logger.info(
-            f"Tracing for {parent_span.id} recorded successfully and sent for processing."
-        )
-        self._clear_pending_spans()
-        self.active_trace_id = None
-        self.active_span = None
-        self.trace_config_cache.clear()
 
-    def _create_trace_id(self) -> str:
-        """Creates a unique mongo id for the trace object.
+        logging.info(f"Sent    spans {tracing.trace_id} #={len(tracing.closed_spans)}")
 
-        Returns:
-            str: stringify oid of the trace
-        """
+    def _clear_closed_spans(self) -> None:
+        tracing = tracing_context.get()
 
-        return str(ObjectId())
+        tracing.closed_spans.clear()
 
-    def _create_span_id(self) -> str:
-        """Creates a unique mongo id for the span object.
+    def _clear_tracked_spans(self) -> None:
+        tracing = tracing_context.get()
 
-        Returns:
-            str: stringify oid of the span
-        """
+        tracing.tracked_spans.clear()
 
-        return str(ObjectId())
+    def _clear_active_span(self) -> None:
+        tracing = tracing_context.get()
 
-    def _clear_pending_spans(self) -> None:
-        """
-        Clear the list of recorded spans to prepare for next batch processing.
-        """
+        span_id = tracing.active_span.id
 
-        self.pending_spans = []
-        self.llm_logger.info(
-            f"Cleared all recorded spans from batch: {self.pending_spans}"
-        )
+        tracing.active_span = None
+
+        logging.debug(f"Cleared active span {span_id}")
+
+    def _update_span_cost(self, span: CreateSpan, cost: Optional[float]) -> None:
+        if span is not None and cost is not None and isinstance(cost, float):
+            if span.cost is None:
+                span.cost = cost
+            else:
+                span.cost += cost
+
+    def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]) -> None:
+        if isinstance(tokens, LlmTokens):
+            tokens = tokens.dict()
+        if span is not None and tokens is not None and isinstance(tokens, dict):
+            if span.tokens is None:
+                span.tokens = LlmTokens(**tokens)
+            else:
+                span.tokens.prompt_tokens += tokens["prompt_tokens"]
+                span.tokens.completion_tokens += tokens["completion_tokens"]
+                span.tokens.total_tokens += tokens["total_tokens"]

agenta/sdk/tracing/tracing_context.py

@@ -0,0 +1,28 @@
+import contextvars
+
+from typing import Optional, Dict, Any, List
+
+from agenta.client.backend.types.create_span import CreateSpan
+
+CURRENT_TRACING_CONTEXT_KEY = "current_tracing_context"
+
+
+class TracingContext:
+    def __init__(self):
+        ### --- TRACE --- ###
+        self.trace_id: Optional[str] = None
+        self.trace_tags: List[str] = []
+
+        ### --- SPANS --- ###
+        self.active_span: Optional[CreateSpan] = None
+        self.tracked_spans: Dict[str, CreateSpan] = {}
+        self.closed_spans: List[CreateSpan] = []
+
+    def __repr__(self) -> str:
+        return f"TracingContext(trace_id=[{self.trace_id}], active_span=[{self.active_span.id if self.active_span else None}{' ' + self.active_span.spankind if self.active_span else ''}])"
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+
+tracing_context = contextvars.ContextVar(CURRENT_TRACING_CONTEXT_KEY, default=None)

agenta/sdk/utils/debug.py

@@ -0,0 +1,68 @@
+from inspect import iscoroutinefunction
+from functools import wraps
+
+from agenta.sdk.tracing.logger import llm_logger as logging
+
+DEBUG = False
+SHIFT = 7
+
+
+def debug(shift=1, req=False, res=False, chars=[">", "<"]):
+    def log_decorator(f):
+        is_async = iscoroutinefunction(f)
+
+        @wraps(f)
+        async def async_log_wrapper(*args, **kwargs):
+            if DEBUG:
+                logging.debug(
+                    " ".join(
+                        [
+                            chars[0] * shift + " " * (SHIFT - shift),
+                            f.__name__ + " ()",
+                            str(args) if req else "",
+                            str(kwargs) if req else "",
+                        ]
+                    )
+                )
+            result = await f(*args, **kwargs)
+            if DEBUG:
+                logging.debug(
+                    " ".join(
+                        [
+                            chars[1] * shift + " " * (SHIFT - shift),
+                            f.__name__ + " <-",
+                            str(result) if res else "",
+                        ]
+                    )
+                )
+            return result
+
+        @wraps(f)
+        def log_wrapper(*args, **kwargs):
+            if DEBUG:
+                logging.debug(
+                    " ".join(
+                        [
+                            chars[0] * shift + " " * (SHIFT - shift),
+                            f.__name__ + " ()",
+                            str(args) if req else "",
+                            str(kwargs) if req else "",
+                        ]
+                    )
+                )
+            result = f(*args, **kwargs)
+            if DEBUG:
+                logging.debug(
+                    " ".join(
+                        [
+                            chars[1] * shift + " " * (SHIFT - shift),
+                            f.__name__ + " <-",
+                            str(result) if res else "",
+                        ]
+                    )
+                )
+            return result
+
+        return async_log_wrapper if is_async else log_wrapper
+
+    return log_decorator

{agenta-0.19.4.dist-info → agenta-0.19.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.19.4
+Version: 0.19.6
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Home-page: https://agenta.ai
 Keywords: LLMOps,LLM,evaluation,prompt engineering

{agenta-0.19.4.dist-info → agenta-0.19.6.dist-info}/RECORD

@@ -1,4 +1,4 @@
-agenta/__init__.py,sha256=geaQThRuZzebZncauJW2SZzKqIT7-ZkIJJf-zZlWHCs,833
+agenta/__init__.py,sha256=6ZfEnXNHoJVXO9HT_U-WuitFvE6Oq_UTrjAZGIRzaDY,887
 agenta/cli/evaluation_commands.py,sha256=fs6492tprPId9p8eGO02Xy-NCBm2RZNJLZWcUxugwd8,474
 agenta/cli/helper.py,sha256=vRxHyeNaltzNIGrfU2vO0H28_rXDzx9QqIZ_S-W6zL4,6212
 agenta/cli/main.py,sha256=Wz0ODhoeKK3Qg_CFUhu6D909szk05tc8ZVBB6H1-w7k,9763
@@ -127,20 +127,22 @@ agenta/docker/docker-assets/lambda_function.py,sha256=h4UZSSfqwpfsCgERv6frqwm_4J
 agenta/docker/docker-assets/main.py,sha256=7MI-21n81U7N7A0GxebNi0cmGWtJKcR2sPB6FcH2QfA,251
 agenta/docker/docker_utils.py,sha256=5uHMCzXkCvIsDdEiwbnnn97KkzsFbBvyMwogCsv_Z5U,3509
 agenta/sdk/__init__.py,sha256=cF0de6DiH-NZWEm0XvPN8_TeC1whPBnDf1WYYE1qK2g,762
-agenta/sdk/agenta_init.py,sha256=zhMYCH6QMDzQADrOmFBaqZRvAYrA6CC35XsePY7IMxg,9789
+agenta/sdk/agenta_init.py,sha256=8MfDuypxohd0qRTdtGjX7L17KW-1UGmzNVdiqF15_ak,9790
 agenta/sdk/client.py,sha256=trKyBOYFZRk0v5Eptxvh87yPf50Y9CqY6Qgv4Fy-VH4,2142
 agenta/sdk/context.py,sha256=q-PxL05-I84puunUAs9LGsffEXcYhDxhQxjuOz2vK90,901
 agenta/sdk/decorators/base.py,sha256=9aNdX5h8a2mFweuhdO-BQPwXGKY9ONPIdLRhSGAGMfY,217
-agenta/sdk/decorators/llm_entrypoint.py,sha256=awdd2fSjVKIz3MVmElHw8_nJxFvePI_zMdPd_QWxIM0,22827
-agenta/sdk/decorators/tracing.py,sha256=4Ta9UN3pjQhxST6lnPaNz1GwnoXz2EHjq5JGqd4g8qw,3817
+agenta/sdk/decorators/llm_entrypoint.py,sha256=umcniOOQfKVdPgHVb_jRhoGKei14Yc3PIZmEC8CU2Wg,22996
+agenta/sdk/decorators/tracing.py,sha256=c9LwQJkhJcyO7Uq-sNpDSwfwOUTAmlqNuJjz4bSx-k0,6172
 agenta/sdk/router.py,sha256=0sbajvn5C7t18anH6yNo7-oYxldHnYfwcbmQnIXBePw,269
 agenta/sdk/tracing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agenta/sdk/tracing/callbacks.py,sha256=8XYcGC4EZQ0lBn2IWmtPH_yC5Hu6Yu-NkdZRuY7CoTw,6816
 agenta/sdk/tracing/context_manager.py,sha256=HskDaiORoOhjeN375gm05wYnieQzh5UnoIsnSAHkAyc,252
-agenta/sdk/tracing/llm_tracing.py,sha256=fHtc1tKgdR2CDnFMrfDIxSvyjwFW3Qk2RE8P-y_tCM0,11361
+agenta/sdk/tracing/llm_tracing.py,sha256=b94Fx5eIOxOvB6S2zeINq8o2taeRpZgxcfNQlh3Jeq0,13689
 agenta/sdk/tracing/logger.py,sha256=GfH7V-jBHcn7h5dbdrnkDMe_ml3wkXFBeoQiqR4KVRc,474
 agenta/sdk/tracing/tasks_manager.py,sha256=ROrWIaqS2J2HHiJtRWiHKlLY8CCsqToP5VeXu7mamck,3748
+agenta/sdk/tracing/tracing_context.py,sha256=EOi1zfqpb2cBjhBtHIphUkVHi4jiWes-RRNdbgk1kMc,906
 agenta/sdk/types.py,sha256=KMnQUOdjaHSWctDLIiMHnk0o3c-C47Vm4Mn2kIZ88YI,5740
+agenta/sdk/utils/debug.py,sha256=QyuPsSoN0425UD13x_msPxSF_VT6YwHiQunZUibI-jg,2149
 agenta/sdk/utils/globals.py,sha256=JmhJcCOSbwvjQ6GDyUc2_SYR27DZk7YcrRH80ktHHOM,435
 agenta/sdk/utils/helper/openai_cost.py,sha256=1VkgvucDnNZm1pTfcVLz9icWunntp1d7zwMmnviy3Uw,5877
 agenta/sdk/utils/preinit.py,sha256=YlJL7RLfel0R7DFp-jK7OV-z4ZIQJM0oupYlk7g8b5o,1278
@@ -159,7 +161,7 @@ agenta/templates/simple_prompt/app.py,sha256=kODgF6lhzsaJPdgL5b21bUki6jkvqjWZzWR
 agenta/templates/simple_prompt/env.example,sha256=g9AE5bYcGPpxawXMJ96gh8oenEPCHTabsiOnfQo3c5k,70
 agenta/templates/simple_prompt/requirements.txt,sha256=ywRglRy7pPkw8bljmMEJJ4aOOQKrt9FGKULZ-DGkoBU,23
 agenta/templates/simple_prompt/template.toml,sha256=DQBtRrF4GU8LBEXOZ-GGuINXMQDKGTEG5y37tnvIUIE,60
-agenta-0.19.4.dist-info/METADATA,sha256=eS83a2QK24f3aAXei-SPL6Xfeqa0MKW0koqmI-aov3g,26460
-agenta-0.19.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-agenta-0.19.4.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.19.4.dist-info/RECORD,,
+agenta-0.19.6.dist-info/METADATA,sha256=9tRG-XDUvW0HDdCk7pM7SwjPRdJ5Y2_w4_ON3ikinKQ,26460
+agenta-0.19.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+agenta-0.19.6.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.19.6.dist-info/RECORD,,