agenta 0.19.8a0__py3-none-any.whl → 0.20.0a0__py3-none-any.whl

agenta/__init__.py

@@ -17,7 +17,7 @@ from .sdk.types import (
 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
+from .sdk.decorators.llm_entrypoint import entrypoint, app, route
 from .sdk.agenta_init import Config, AgentaSingleton, init
 from .sdk.utils.helper.openai_cost import calculate_token_usage
 from .sdk.client import Agenta

agenta/client/backend/types/create_span.py

@@ -20,7 +20,10 @@ class CreateSpan(pydantic.BaseModel):
     variant_id: typing.Optional[str]
     variant_name: typing.Optional[str]
     inputs: typing.Optional[typing.Dict[str, typing.Any]]
-    outputs: typing.Optional[typing.List[str]]
+    internals: typing.Optional[typing.Dict[str, typing.Any]]
+    outputs: typing.Optional[
+        typing.Union[typing.Dict[str, typing.Any], typing.List[str]]
+    ]
     config: typing.Optional[typing.Dict[str, typing.Any]]
     environment: typing.Optional[str]
     tags: typing.Optional[typing.List[str]]

agenta/client/backend/types/llm_tokens.py

@@ -12,9 +12,9 @@ except ImportError:
 
 
 class LlmTokens(pydantic.BaseModel):
-    prompt_tokens: typing.Optional[int]
-    completion_tokens: typing.Optional[int]
-    total_tokens: typing.Optional[int]
+    prompt_tokens: typing.Optional[int] = 0
+    completion_tokens: typing.Optional[int] = 0
+    total_tokens: typing.Optional[int] = 0
 
     def json(self, **kwargs: typing.Any) -> str:
         kwargs_with_defaults: typing.Any = {

agenta/sdk/__init__.py

@@ -16,7 +16,7 @@ from .types import (
 
 from .tracing.llm_tracing import Tracing
 from .decorators.tracing import instrument
-from .decorators.llm_entrypoint import entrypoint, app
+from .decorators.llm_entrypoint import entrypoint, app, route
 from .agenta_init import Config, AgentaSingleton, init
 from .utils.helper.openai_cost import calculate_token_usage
 

agenta/sdk/decorators/llm_entrypoint.py

@@ -3,6 +3,7 @@
 import os
 import sys
 import time
+import json
 import inspect
 import argparse
 import asyncio
@@ -15,11 +16,11 @@ from typing import Any, Callable, Dict, Optional, Tuple, List
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import Body, FastAPI, UploadFile, HTTPException
 
-import agenta
+import agenta as ag
 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.tracing.tracing_context import tracing_context, TracingContext
 from agenta.sdk.decorators.base import BaseDecorator
 from agenta.sdk.types import (
     Context,
@@ -32,10 +33,12 @@ from agenta.sdk.types import (
     TextParam,
     MessagesInput,
     FileInputURL,
-    FuncResponse,
+    BaseResponse,
     BinaryParam,
 )
 
+from pydantic import BaseModel, HttpUrl
+
 app = FastAPI()
 
 origins = [
@@ -52,16 +55,56 @@ app.add_middleware(
 
 app.include_router(router, prefix="")
 
-
 from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
 
 
 logging.setLevel("DEBUG")
 
 
+class PathValidator(BaseModel):
+    url: HttpUrl
+
+
+class route(BaseDecorator):
+    # This decorator is used to expose specific stages of a workflow (embedding, retrieval, summarization, etc.)
+    # as independent endpoints. It is designed for backward compatibility with existing code that uses
+    # the @entrypoint decorator, which has certain limitations. By using @route(), we can create new
+    # routes without altering the main workflow entrypoint. This helps in modularizing the services
+    # and provides flexibility in how we expose different functionalities as APIs.
+    def __init__(self, path):
+        path = "/" + path.strip("/").strip()
+        path = "" if path == "/" else path
+
+        PathValidator(url=f"http://example.com{path}")
+
+        self.route_path = path
+
+    def __call__(self, f):
+        self.e = entrypoint(f, route_path=self.route_path)
+
+        return f
+
+
 class entrypoint(BaseDecorator):
-    """Decorator class to wrap a function for HTTP POST, terminal exposure and enable tracing.
+    """
+    Decorator class to wrap a function for HTTP POST, terminal exposure and enable tracing.
+
+    This decorator generates the following endpoints:
+
+    Playground Endpoints
+    - /generate                 with @entrypoint, @route("/"), @route(path="") # LEGACY
+    - /playground/run           with @entrypoint, @route("/"), @route(path="")
+    - /playground/run/{route}   with @route({route}), @route(path={route})
 
+    Deployed Endpoints:
+    - /generate_deployed        with @entrypoint, @route("/"), @route(path="") # LEGACY
+    - /run                      with @entrypoint, @route("/"), @route(path="")
+    - /run/{route}              with @route({route}), @route(path={route})
+
+    The rationale is:
+    - There may be multiple endpoints, based on the different routes.
+    - It's better to make it explicit that an endpoint is for the playground.
+    - Prefixing the routes with /run is more futureproof in case we add more endpoints.
 
     Example:
     ```python
@@ -73,31 +116,64 @@ class entrypoint(BaseDecorator):
     ```
     """
 
-    def __init__(self, func: Callable[..., Any]):
-        endpoint_name = "generate"
+    routes = list()
+
+    def __init__(self, func: Callable[..., Any], route_path=""):
+        DEFAULT_PATH = "generate"
+        PLAYGROUND_PATH = "/playground"
+        RUN_PATH = "/run"
+
         func_signature = inspect.signature(func)
-        config_params = agenta.config.all()
+        config_params = ag.config.all()
         ingestible_files = self.extract_ingestible_files(func_signature)
 
+        ### --- Playground  --- #
         @debug()
         @functools.wraps(func)
         async def wrapper(*args, **kwargs) -> Any:
             func_params, api_config_params = self.split_kwargs(kwargs, config_params)
             self.ingest_files(func_params, ingestible_files)
-            agenta.config.set(**api_config_params)
+            ag.config.set(**api_config_params)
 
             # Set the configuration and environment of the LLM app parent span at run-time
-            agenta.tracing.update_baggage(
+            ag.tracing.update_baggage(
                 {"config": config_params, "environment": "playground"}
             )
 
-            # Exceptions are all handled inside self.execute_function()
-            llm_result = await self.execute_function(
+            entrypoint_result = await self.execute_function(
                 func, *args, params=func_params, config_params=config_params
             )
 
-            return llm_result
+            return entrypoint_result
+
+        self.update_function_signature(
+            wrapper, func_signature, config_params, ingestible_files
+        )
 
+        #
+        if route_path == "":
+            route = f"/{DEFAULT_PATH}"
+            app.post(route, response_model=BaseResponse)(wrapper)
+            entrypoint.routes.append(
+                {
+                    "func": func.__name__,
+                    "endpoint": DEFAULT_PATH,
+                    "params": {**config_params, **func_signature.parameters},
+                }
+            )
+
+        route = f"{PLAYGROUND_PATH}{RUN_PATH}{route_path}"
+        app.post(route, response_model=BaseResponse)(wrapper)
+        entrypoint.routes.append(
+            {
+                "func": func.__name__,
+                "endpoint": route[1:].replace("/", "_"),
+                "params": {**config_params, **func_signature.parameters},
+            }
+        )
+        ### ---------------------------- #
+
+        ### --- Deployed / Published --- #
         @debug()
         @functools.wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
@@ -106,44 +182,51 @@ class entrypoint(BaseDecorator):
             }
 
             if "environment" in kwargs and kwargs["environment"] is not None:
-                agenta.config.pull(environment_name=kwargs["environment"])
+                ag.config.pull(environment_name=kwargs["environment"])
             elif "config" in kwargs and kwargs["config"] is not None:
-                agenta.config.pull(config_name=kwargs["config"])
+                ag.config.pull(config_name=kwargs["config"])
             else:
-                agenta.config.pull(config_name="default")
+                ag.config.pull(config_name="default")
 
             # Set the configuration and environment of the LLM app parent span at run-time
-            agenta.tracing.update_baggage(
+            ag.tracing.update_baggage(
                 {"config": config_params, "environment": kwargs["environment"]}
             )
 
-            llm_result = await self.execute_function(
+            entrypoint_result = await self.execute_function(
                 func, *args, params=func_params, config_params=config_params
             )
 
-            return llm_result
-
-        self.update_function_signature(
-            wrapper, func_signature, config_params, ingestible_files
-        )
-        route = f"/{endpoint_name}"
-        app.post(route, response_model=FuncResponse)(wrapper)
+            return entrypoint_result
 
         self.update_deployed_function_signature(
             wrapper_deployed,
             func_signature,
             ingestible_files,
         )
-        route_deployed = f"/{endpoint_name}_deployed"
-        app.post(route_deployed, response_model=FuncResponse)(wrapper_deployed)
-        self.override_schema(
-            openapi_schema=app.openapi(),
-            func_name=func.__name__,
-            endpoint=endpoint_name,
-            params={**config_params, **func_signature.parameters},
-        )
 
-        if self.is_main_script(func):
+        if route_path == "/":
+            route_deployed = f"/{DEFAULT_PATH}_deployed"
+            app.post(route_deployed, response_model=BaseResponse)(wrapper_deployed)
+
+        route_deployed = f"{RUN_PATH}{route_path}"
+        app.post(route_deployed, response_model=BaseResponse)(wrapper_deployed)
+        ### ---------------------------- #
+
+        ### --- Update OpenAPI --- #
+        app.openapi_schema = None  # Forces FastAPI to re-generate the schema
+        openapi_schema = app.openapi()
+
+        for route in entrypoint.routes:
+            self.override_schema(
+                openapi_schema=openapi_schema,
+                func=route["func"],
+                endpoint=route["endpoint"],
+                params=route["params"],
+            )
+        ### ---------------------- #
+
+        if self.is_main_script(func) and route_path == "":
             self.handle_terminal_run(
                 func,
                 func_signature.parameters,  # type: ignore
@@ -198,51 +281,55 @@ class entrypoint(BaseDecorator):
             For synchronous functions, it calls them directly, while for asynchronous functions,
             it awaits their execution.
             """
+            data = None
+            trace = None
+
+            token = None
+            if tracing_context.get() is None:
+                token = tracing_context.set(TracingContext())
+
             is_coroutine_function = inspect.iscoroutinefunction(func)
-            start_time = time.perf_counter()
 
             if is_coroutine_function:
                 result = await func(*args, **func_params["params"])
             else:
                 result = func(*args, **func_params["params"])
 
-            end_time = time.perf_counter()
-            latency = end_time - start_time
+            if token is not None:
+                trace = ag.tracing.dump_trace()
+                tracing_context.reset(token)
 
             if isinstance(result, Context):
                 save_context(result)
+
             if isinstance(result, Dict):
-                return FuncResponse(**result, latency=round(latency, 4))
-            if isinstance(result, str):
-                return FuncResponse(
-                    message=result, usage=None, cost=None, latency=round(latency, 4)
-                )
-            if isinstance(result, int) or isinstance(result, float):
-                return FuncResponse(
-                    message=str(result),
-                    usage=None,
-                    cost=None,
-                    latency=round(latency, 4),
-                )
-            if result is None:
-                return FuncResponse(
-                    message="Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
-                    usage=None,
-                    cost=None,
-                    latency=round(latency, 4),
+                data = result
+            elif isinstance(result, str):
+                data = {"message": result}
+            elif isinstance(result, int) or isinstance(result, float):
+                data = {"message": str(result)}
+
+            if data is None:
+                warning = (
+                    "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
                 )
+
+                data = {"message": warning}
+
+            return BaseResponse(data=data, trace=trace)
+
         except Exception as e:
             self.handle_exception(e)
-        return FuncResponse(message="Unexpected error occurred when calling the @entrypoint decorated function", latency=0)  # type: ignore
 
     def handle_exception(self, e: Exception):
-        """Handle exceptions."""
+        status_code = e.status_code if hasattr(e, "status_code") else 500
+        message = str(e)
+        stacktrace = traceback.format_exception(e, value=e, tb=e.__traceback__)  # type: ignore
+        detail = {"message": message, "stacktrace": stacktrace}
 
-        status_code: int = e.status_code if hasattr(e, "status_code") else 500
-        traceback_str = traceback.format_exception(e, value=e, tb=e.__traceback__)  # type: ignore
         raise HTTPException(
             status_code=status_code,
-            detail={"error": str(e), "traceback": "".join(traceback_str)},
+            detail=detail,
         )
 
     def update_wrapper_signature(
@@ -418,26 +505,29 @@ class entrypoint(BaseDecorator):
                 file_path=args_func_params[name],
             )
 
-        agenta.config.set(**args_config_params)
+        ag.config.set(**args_config_params)
 
         # Set the configuration and environment of the LLM app parent span at run-time
-        agenta.tracing.update_baggage(
-            {"config": agenta.config.all(), "environment": "bash"}
-        )
+        ag.tracing.update_baggage({"config": ag.config.all(), "environment": "bash"})
 
         loop = asyncio.get_event_loop()
+
         result = loop.run_until_complete(
             self.execute_function(
                 func,
                 **{"params": args_func_params, "config_params": args_config_params},
             )
         )
-        print(
-            f"\n========== Result ==========\n\nMessage: {result.message}\nCost: {result.cost}\nToken Usage: {result.usage}"
-        )
+
+        print(f"\n========== Result ==========\n")
+
+        print("-> data")
+        print(json.dumps(result.data, indent=2))
+        print("-> trace")
+        print(json.dumps(result.trace, indent=2))
 
     def override_schema(
-        self, openapi_schema: dict, func_name: str, endpoint: str, params: dict
+        self, openapi_schema: dict, func: str, endpoint: str, params: dict
     ):
         """
         Overrides the default openai schema generated by fastapi with additional information about:
@@ -452,7 +542,7 @@ class entrypoint(BaseDecorator):
 
         Args:
             openapi_schema (dict): The openapi schema generated by fastapi
-            func_name (str): The name of the function to override
+            func (str): The name of the function to override
             endpoint (str): The name of the endpoint to override
             params (dict(param_name, param_val)): The dictionary of the parameters for the function
         """
@@ -484,8 +574,29 @@ class entrypoint(BaseDecorator):
                         # value = {'temperature': { "type": "number", "title": "Temperature", "x-parameter": "float" }}
                     return value
 
+        def get_type_from_param(param_val):
+            param_type = "string"
+            annotation = param_val.annotation
+
+            if annotation == int:
+                param_type = "integer"
+            elif annotation == float:
+                param_type = "number"
+            elif annotation == dict:
+                param_type = "object"
+            elif annotation == bool:
+                param_type = "boolean"
+            elif annotation == list:
+                param_type = "list"
+            elif annotation == str:
+                param_type = "string"
+            else:
+                print("ERROR, unhandled annotation:", annotation)
+
+            return param_type
+
         schema_to_override = openapi_schema["components"]["schemas"][
-            f"Body_{func_name}_{endpoint}_post"
+            f"Body_{func}_{endpoint}_post"
         ]["properties"]
         for param_name, param_val in params.items():
             if isinstance(param_val, GroupedMultipleChoiceParam):
@@ -501,7 +612,7 @@ class entrypoint(BaseDecorator):
                 subschema["choices"] = param_val.choices  # type: ignore
                 subschema["default"] = param_val.default  # type: ignore
 
-            if isinstance(param_val, MultipleChoiceParam):
+            elif isinstance(param_val, MultipleChoiceParam):
                 subschema = find_in_schema(
                     param_val.__schema_type_properties__(),
                     schema_to_override,
@@ -520,7 +631,7 @@ class entrypoint(BaseDecorator):
                     default if default in param_choices else choices[0]
                 )
 
-            if isinstance(param_val, FloatParam):
+            elif isinstance(param_val, FloatParam):
                 subschema = find_in_schema(
                     param_val.__schema_type_properties__(),
                     schema_to_override,
@@ -531,7 +642,7 @@ class entrypoint(BaseDecorator):
                 subschema["maximum"] = param_val.maxval  # type: ignore
                 subschema["default"] = param_val
 
-            if isinstance(param_val, IntParam):
+            elif isinstance(param_val, IntParam):
                 subschema = find_in_schema(
                     param_val.__schema_type_properties__(),
                     schema_to_override,
@@ -542,7 +653,7 @@ class entrypoint(BaseDecorator):
                 subschema["maximum"] = param_val.maxval  # type: ignore
                 subschema["default"] = param_val
 
-            if (
+            elif (
                 isinstance(param_val, inspect.Parameter)
                 and param_val.annotation is DictInput
             ):
@@ -554,7 +665,7 @@ class entrypoint(BaseDecorator):
                 )
                 subschema["default"] = param_val.default["default_keys"]
 
-            if isinstance(param_val, TextParam):
+            elif isinstance(param_val, TextParam):
                 subschema = find_in_schema(
                     param_val.__schema_type_properties__(),
                     schema_to_override,
@@ -563,7 +674,7 @@ class entrypoint(BaseDecorator):
                 )
                 subschema["default"] = param_val
 
-            if (
+            elif (
                 isinstance(param_val, inspect.Parameter)
                 and param_val.annotation is MessagesInput
             ):
@@ -575,7 +686,7 @@ class entrypoint(BaseDecorator):
                 )
                 subschema["default"] = param_val.default
 
-            if (
+            elif (
                 isinstance(param_val, inspect.Parameter)
                 and param_val.annotation is FileInputURL
             ):
@@ -587,7 +698,7 @@ class entrypoint(BaseDecorator):
                 )
                 subschema["default"] = "https://example.com"
 
-            if isinstance(param_val, BinaryParam):
+            elif isinstance(param_val, BinaryParam):
                 subschema = find_in_schema(
                     param_val.__schema_type_properties__(),
                     schema_to_override,
@@ -595,3 +706,11 @@ class entrypoint(BaseDecorator):
                     "bool",
                 )
                 subschema["default"] = param_val.default  # type: ignore
+            else:
+                subschema = {
+                    "title": str(param_name).capitalize(),
+                    "type": get_type_from_param(param_val),
+                }
+                if param_val.default != inspect._empty:
+                    subschema["default"] = param_val.default  # type: ignore
+                schema_to_override[param_name] = subschema

agenta/sdk/decorators/tracing.py

@@ -8,7 +8,6 @@ from typing import Any, Callable, Optional
 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
 
 
@@ -38,149 +37,55 @@ class instrument(BaseDecorator):
     """
 
     def __init__(
-        self, config: Optional[dict] = None, spankind: str = "workflow"
+        self,
+        config: Optional[dict] = None,
+        spankind: str = "workflow",
     ) -> None:
         self.config = config
         self.spankind = spankind
-        self.tracing = ag.tracing
 
     def __call__(self, func: Callable[..., Any]):
         is_coroutine_function = inspect.iscoroutinefunction(func)
 
-        @debug()
-        @wraps(func)
-        async def async_wrapper(*args, **kwargs):
-            result = None
+        def get_inputs(*args, **kwargs):
             func_args = inspect.getfullargspec(func).args
             input_dict = {name: value for name, value in zip(func_args, args)}
             input_dict.update(kwargs)
 
-            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())]))
+            return input_dict
 
-                self.tracing.start_span(
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            async def wrapped_func(*args, **kwargs):
+                with ag.tracing.Context(
                     name=func.__name__,
-                    input=input_dict,
+                    input=get_inputs(*args, **kwargs),
                     spankind=self.spankind,
                     config=self.config,
-                )
-
-                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())]))
+                    ag.tracing.store_outputs(result)
 
                     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):
-            result = None
-            func_args = inspect.getfullargspec(func).args
-            input_dict = {name: value for name, value in zip(func_args, args)}
-            input_dict.update(kwargs)
-
             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(
+                with ag.tracing.Context(
                     name=func.__name__,
-                    input=input_dict,
+                    input=get_inputs(*args, **kwargs),
                     spankind=self.spankind,
                     config=self.config,
-                )
-
-                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())]))
+                    ag.tracing.store_outputs(result)
 
                     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,11 +1,16 @@
 import os
+import copy
+import json
 from uuid import uuid4
 
+import traceback
 from threading import Lock
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any, List
 
-from agenta.sdk.tracing.tracing_context import tracing_context
+from contextlib import contextmanager
+
+from agenta.sdk.tracing.tracing_context import tracing_context, TracingContext
 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
@@ -91,6 +96,46 @@ class Tracing(metaclass=SingletonMeta):
             base_url=self.host, api_key=self.api_key, timeout=120  # type: ignore
         ).observability
 
+    ### --- Context Manager --- ###
+
+    @contextmanager
+    def Context(self, **kwargs):
+        # This will evolve as be work towards OTel compliance
+
+        token = None
+
+        try:
+            if tracing_context.get() is None:
+                token = tracing_context.set(TracingContext())
+
+            self.open_span(**kwargs)
+
+            yield
+
+            self.set_status(status="OK")
+
+        except Exception as e:
+            logging.error(e)
+
+            result = {
+                "message": str(e),
+                "stacktrace": traceback.format_exc(),
+            }
+
+            self.set_status(status="ERROR")
+            self.set_attributes({"traceback_exception": traceback.format_exc()})
+            self.store_outputs(result)
+
+            raise
+
+        finally:
+            self.close_span()
+
+            if token is not None:
+                self.flush_spans()
+
+                tracing_context.reset(token)
+
     ### --- API --- ###
 
     @debug()
@@ -177,11 +222,7 @@ class Tracing(metaclass=SingletonMeta):
         if not self.api_key:
             logging.error("No API key")
         else:
-            self._process_closed_spans()
-
-        self._clear_closed_spans()
-        self._clear_tracked_spans()
-        self._clear_active_span()
+            self._process_spans()
 
         self._clear_trace_tags()
 
@@ -213,6 +254,7 @@ class Tracing(metaclass=SingletonMeta):
             attributes={},
             status=SpanStatusCode.UNSET.value,
             start_time=datetime.now(timezone.utc),
+            internals=None,
             outputs=None,
             tags=None,
             user=None,
@@ -224,12 +266,11 @@ class Tracing(metaclass=SingletonMeta):
         )
 
         if tracing.trace_id is None:
-            self.start_trace(span, config)
+            self.open_trace(span, config)
         else:
             span.parent_span_id = tracing.active_span.id  # type: ignore
 
-        tracing.tracked_spans[span.id] = span
-        tracing.active_span = span
+        tracing.push(span)
         ### --- TO BE CLEANED --- <<<
 
         logging.info(f"Opened  span  {span_id} {spankind.upper()}")
@@ -282,7 +323,7 @@ class Tracing(metaclass=SingletonMeta):
         tracing.active_span.status = status
 
     @debug()
-    def close_span(self, outputs: Dict[str, Any]) -> None:
+    def close_span(self) -> None:
         """
         Ends the active span, if it is a parent span, ends the trace too.
 
@@ -313,24 +354,19 @@ class Tracing(metaclass=SingletonMeta):
         ### --- TO BE CLEANED --- >>>
         tracing.active_span.end_time = datetime.now(timezone.utc)
 
-        tracing.active_span.outputs = [outputs.get("message", "")]
-
-        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(tracing.active_span, outputs.get("cost", None))
-            self._update_span_tokens(tracing.active_span, outputs.get("usage", None))
-
-        tracing.closed_spans.append(tracing.active_span)
+        # TODO: Remove this whole part. Setting the cost should be done through set_span_attribute
+        if isinstance(tracing.active_span.outputs, dict):
+            self._update_span_cost(
+                tracing.active_span, tracing.active_span.outputs.get("cost", None)
+            )
+            self._update_span_tokens(
+                tracing.active_span, tracing.active_span.outputs.get("usage", None)
+            )
 
         active_span_parent_id = tracing.active_span.parent_span_id
 
-        if active_span_parent_id is None:
-            self.end_trace(parent_span=tracing.active_span)
-
-        else:
-            parent_span = tracing.tracked_spans[active_span_parent_id]
+        if active_span_parent_id is not None:
+            parent_span = tracing.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
@@ -338,6 +374,95 @@ class Tracing(metaclass=SingletonMeta):
 
         logging.info(f"Closed  span  {span_id} {spankind}")
 
+    @debug()
+    def store_internals(self, internals: Dict[str, Any] = {}) -> None:
+        """
+        Set internals for the active span.
+
+        Args:
+            internals (Dict[str, Any], optional): A dictionary of local variables to set. Defaults to {}.
+        """
+
+        tracing = tracing_context.get()
+
+        if tracing.active_span is None:
+            logging.error(f"Cannot set internals ({set(internals)}), no active span")
+            return
+
+        logging.info(
+            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} internals={internals}"
+        )
+
+        if tracing.active_span.internals is None:
+            tracing.active_span.internals = dict()
+
+        for key, value in internals.items():
+            tracing.active_span.internals[key] = value  # type: ignore
+
+    @debug()
+    def store_outputs(self, outputs: Dict[str, Any] = {}) -> None:
+        """
+        Set outputs for the active span.
+
+        Args:
+            outputs (Dict[str, Any], optional): A dictionary of output variables to set. Defaults to {}.
+        """
+
+        tracing = tracing_context.get()
+
+        if tracing.active_span is None:
+            logging.error(f"Cannot set outputs ({set(outputs)}), no active span")
+            return
+
+        logging.info(
+            f"Setting span  {tracing.active_span.id} {tracing.active_span.spankind.upper()} outputs={outputs}"
+        )
+
+        tracing.active_span.outputs = outputs
+
+    def dump_trace(self):
+        """
+        Collects and organizes tracing information into a dictionary.
+        This function retrieves the current tracing context and extracts relevant data such as `trace_id`, `cost`, `tokens`, and `latency` for the whole trace.
+        It also dumps detailed span information using the `dump_spans` method and includes it in the trace dictionary.
+        If an error occurs during the process, it logs the error message and stack trace.
+
+        Returns:
+            dict: A dictionary containing the trace information.
+        """
+        try:
+            trace = dict()
+
+            tracing = tracing_context.get()
+
+            trace["trace_id"] = tracing.trace_id
+
+            for span in tracing.spans.values():
+                if span.parent_span_id is None:
+                    trace["cost"] = span.cost
+                    trace["usage"] = (
+                        None if span.tokens is None else json.loads(span.tokens.json())
+                    )
+                    trace["latency"] = (span.end_time - span.start_time).total_seconds()
+
+            spans = (
+                []
+                if len(tracing.spans) == 0
+                else [json.loads(span.json()) for span in tracing.spans.values()]
+            )
+
+            if spans is not None:
+                trace["spans"] = spans
+
+        except Exception as e:
+            logging.error(e)
+            logging.error(traceback.format_exc())
+
+        return trace
+
+    def flush_spans(self) -> None:
+        self.close_trace()
+
     ### --- Legacy API --- ###
 
     def start_trace(
@@ -348,7 +473,7 @@ class Tracing(metaclass=SingletonMeta):
     ) -> None:  # Legacy
         self.open_trace(span, config, **kwargs)
 
-    def end_trace(self, parent_span: CreateSpan) -> None:  # Legacy
+    def end_trace(self, _: CreateSpan) -> None:  # Legacy
         self.close_trace()
 
     def start_span(
@@ -371,7 +496,8 @@ class Tracing(metaclass=SingletonMeta):
         self.set_attributes(attributes)
 
     def end_span(self, outputs: Dict[str, Any]) -> None:  # Legacy
-        self.close_span(outputs)
+        self.store_outputs(outputs)
+        self.close_span()
 
     ### --- Helper Functions --- ###
 
@@ -400,46 +526,24 @@ class Tracing(metaclass=SingletonMeta):
         # return uuid4().hex[:16]
         return str(ObjectId())
 
-    def _process_closed_spans(self) -> None:
+    def _process_spans(self) -> None:
         tracing = tracing_context.get()
 
-        logging.info(f"Sending spans {tracing.trace_id} #={len(tracing.closed_spans)} ")
+        spans = list(tracing.spans.values())
 
-        # async def mock_create_traces(trace, spans):
-        #    print("trace-id", trace)
-        #    print("spans", spans)
+        logging.info(f"Sending trace {tracing.trace_id} spans={len(spans)} ")
 
         self.tasks_manager.add_task(
             tracing.trace_id,
-            "trace",
+            "send-trace",
             # mock_create_traces(
             self.client.create_traces(
-                trace=tracing.trace_id,
-                spans=tracing.closed_spans,  # type: ignore
+                trace=tracing.trace_id, spans=spans  # type: ignore
             ),
             self.client,
         )
 
-        logging.info(f"Sent    spans {tracing.trace_id} #={len(tracing.closed_spans)}")
-
-    def _clear_closed_spans(self) -> None:
-        tracing = tracing_context.get()
-
-        tracing.closed_spans.clear()
-
-    def _clear_tracked_spans(self) -> None:
-        tracing = tracing_context.get()
-
-        tracing.tracked_spans.clear()
-
-    def _clear_active_span(self) -> None:
-        tracing = tracing_context.get()
-
-        span_id = tracing.active_span.id
-
-        tracing.active_span = None
-
-        logging.debug(f"Cleared active span {span_id}")
+        logging.info(f"Sent    trace {tracing.trace_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):
@@ -455,6 +559,16 @@ class Tracing(metaclass=SingletonMeta):
             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"]
+                span.tokens.prompt_tokens += (
+                    tokens["prompt_tokens"]
+                    if tokens["prompt_tokens"] is not None
+                    else 0
+                )
+                span.tokens.completion_tokens += (
+                    tokens["completion_tokens"]
+                    if tokens["completion_tokens"] is not None
+                    else 0
+                )
+                span.tokens.total_tokens += (
+                    tokens["total_tokens"] if tokens["total_tokens"] is not None else 0
+                )

agenta/sdk/tracing/tasks_manager.py

@@ -105,8 +105,9 @@ class TaskQueue(object):
                 future = self._thread_pool.submit(asyncio.run, task.run())
                 future.result()
             except Exception as exc:
-                self._logger.error(f"Error running task: {str(exc)}")
-                self._logger.error(f"Recording {task.coroutine_type} status to ERROR.")
+                self._logger.error(
+                    f"Task '{task.coroutine_type}' failed with error: {str(exc)}"
+                )
                 break
             finally:
                 self.tasks.task_done()

agenta/sdk/tracing/tracing_context.py

@@ -1,6 +1,6 @@
-import contextvars
+from contextvars import ContextVar
 
-from typing import Optional, Dict, Any, List
+from typing import Optional, Dict, List
 
 from agenta.client.backend.types.create_span import CreateSpan
 
@@ -15,14 +15,17 @@ class TracingContext:
 
         ### --- SPANS --- ###
         self.active_span: Optional[CreateSpan] = None
-        self.tracked_spans: Dict[str, CreateSpan] = {}
-        self.closed_spans: List[CreateSpan] = []
+        self.spans: Dict[str, 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 ''}])"
+        return f"TracingContext(trace='{self.trace_id}', spans={[f'{span.id} {span.spankind}' for span in self.spans.values()]})"
 
     def __str__(self) -> str:
         return self.__repr__()
 
+    def push(self, span) -> None:
+        self.active_span = span
+        self.spans[span.id] = span
 
-tracing_context = contextvars.ContextVar(CURRENT_TRACING_CONTEXT_KEY, default=None)
+
+tracing_context = ContextVar(CURRENT_TRACING_CONTEXT_KEY, default=None)

agenta/sdk/types.py

@@ -1,5 +1,5 @@
 import json
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 
 from pydantic import ConfigDict, BaseModel, HttpUrl
 
@@ -16,11 +16,10 @@ class LLMTokenUsage(BaseModel):
     total_tokens: int
 
 
-class FuncResponse(BaseModel):
-    message: str
-    usage: Optional[LLMTokenUsage]
-    cost: Optional[float]
-    latency: float
+class BaseResponse(BaseModel):
+    version: Optional[str] = "2.0"
+    data: Optional[Dict[str, Any]]
+    trace: Optional[Dict[str, Any]]
 
 
 class DictInput(dict):

{agenta-0.19.8a0.dist-info → agenta-0.20.0a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agenta
-Version: 0.19.8a0
+Version: 0.20.0a0
 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.8a0.dist-info → agenta-0.20.0a0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-agenta/__init__.py,sha256=6ZfEnXNHoJVXO9HT_U-WuitFvE6Oq_UTrjAZGIRzaDY,887
+agenta/__init__.py,sha256=nL_Fogmx_64xS--WjhE4gE_0h5LIJDbDGKpjZoNJ9oI,894
 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
@@ -54,7 +54,7 @@ agenta/client/backend/types/base_output.py,sha256=ynXhDBQKrkR6Lnkx-yv6Q8xW4wXmzX
 agenta/client/backend/types/body_import_testset.py,sha256=7dVF3mv3VO0Co8F0qxLAgu4jabqDPjebK4mYvcd_TuA,1061
 agenta/client/backend/types/config_db.py,sha256=P0cSYvVOn0ZxpYMIdvhWpQVjRuBS5APe6qlc69AXaF4,1028
 agenta/client/backend/types/create_app_output.py,sha256=pgnTnfZx35Q-8wZ1yTZBQ0ydYacGzFC9kyLug_UvymM,986
-agenta/client/backend/types/create_span.py,sha256=Ldb2zclVtVyBRKxM2Ap3YjE-FP3jbaOya96_ZsYw1cg,1794
+agenta/client/backend/types/create_span.py,sha256=CfTLCHRnDyicJS3-cioDM5ZWZQq8SEi6z3BEDfnQ9Hs,1913
 agenta/client/backend/types/create_trace_response.py,sha256=FO-Ii9JEn2AQ1nmZYmjnKRbACsNxRvY_-xn7Ys7Yo8A,1012
 agenta/client/backend/types/docker_env_vars.py,sha256=altCvA1k-zdAkKNYLwaCnmV48HZg9cwe2cHu_BGImac,986
 agenta/client/backend/types/environment_output.py,sha256=dl0GKodeqB7kWK5mH6Y4iBppkpwRzSTmtkXH1II4L6w,1257
@@ -85,7 +85,7 @@ agenta/client/backend/types/image.py,sha256=p7Vmp7HlMV3YyXe8SFdXYJjCbPNIypW6NfVG
 agenta/client/backend/types/invite_request.py,sha256=1nJTUHspzw2WYpUSd4UUtRnjDHM-dqDBvYewgw-hCQE,993
 agenta/client/backend/types/list_api_keys_response.py,sha256=ZNh7jKwHEMKNp8OV5WJ5XxtKn39DwqK1f8vlFKl54x4,1097
 agenta/client/backend/types/llm_run_rate_limit.py,sha256=mfT4lTczPxrJvd8ZCOAjPvw58QoM151p_uZT0PWNOJ4,1045
-agenta/client/backend/types/llm_tokens.py,sha256=J236Fgmz5TeFO0MQA1ZA1QozvR6d3kt8aEUyWkq3jLI,1070
+agenta/client/backend/types/llm_tokens.py,sha256=nDaiJRZs1jvsXVDDoc-CXGNLUP8hRhioY37P1CGEav0,1082
 agenta/client/backend/types/new_human_evaluation.py,sha256=lIgMjVccSp22RRfMGGLH4-yKjMtJeQvjhlwX9EtAxmY,1150
 agenta/client/backend/types/new_testset.py,sha256=9NOC1-f_UZASy4ptzidLNcRU6Odq609ayvSQxEva-40,1009
 agenta/client/backend/types/organization.py,sha256=vJf6Gbz8WCnqabPQmt_t_gfrWPpuvTXgTxKCJKJsrmc,1218
@@ -126,22 +126,22 @@ agenta/docker/docker-assets/entrypoint.sh,sha256=29XK8VQjQsx4hN2j-4JDy-6kQb5y4LC
 agenta/docker/docker-assets/lambda_function.py,sha256=h4UZSSfqwpfsCgERv6frqwm_4JrYu9rLz3I-LxCfeEg,83
 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/__init__.py,sha256=ewYNjm6AHlqkIrPfX2D_pXZMwShOdhEUcWXb7xGA2bk,769
 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=umcniOOQfKVdPgHVb_jRhoGKei14Yc3PIZmEC8CU2Wg,22996
-agenta/sdk/decorators/tracing.py,sha256=c9LwQJkhJcyO7Uq-sNpDSwfwOUTAmlqNuJjz4bSx-k0,6172
+agenta/sdk/decorators/llm_entrypoint.py,sha256=cyXUHJ-cldbXmBOo3CF5XTSz_Pr2vAA5Inv5LHdigyE,26910
+agenta/sdk/decorators/tracing.py,sha256=BzXa_2b2VvbO5ZNTh126OkL3LFq8dMK-QFug8BHUBKA,2687
 agenta/sdk/router.py,sha256=0sbajvn5C7t18anH6yNo7-oYxldHnYfwcbmQnIXBePw,269
 agenta/sdk/tracing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agenta/sdk/tracing/callbacks.py,sha256=0rqkW-PGZSaNSMp_t4bGI7R9HQRgTyIy0gxmpGVJWpE,6915
 agenta/sdk/tracing/context_manager.py,sha256=HskDaiORoOhjeN375gm05wYnieQzh5UnoIsnSAHkAyc,252
-agenta/sdk/tracing/llm_tracing.py,sha256=OL9OZ9sKv58hAsSxf33E-VYFx2ZeJpr9jqe3Hvz1CX8,13706
+agenta/sdk/tracing/llm_tracing.py,sha256=VZskBPxa4S-0810KhsTnSy7XdrG3ovcpsKjfxQEGUx8,17167
 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/tracing/tasks_manager.py,sha256=FBSFOWIKBycyA4ShB2ZVMzrzYQ8pWGWWBClFX8nlZFA,3726
+agenta/sdk/tracing/tracing_context.py,sha256=dTbsBMbIAmSOaWIxSfhKHOB1JjrAZ0eSHAlfzZs08Z4,891
+agenta/sdk/types.py,sha256=1rVy8ob-rTOrIFcSSseXrt0JQtqqhlJfVgVxCB2ErCk,5754
 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
@@ -161,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.8a0.dist-info/METADATA,sha256=SJjoIarkgSM03m_L3eZ0_xLTZwYNRHr6JbO0bj4r4WA,26462
-agenta-0.19.8a0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-agenta-0.19.8a0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.19.8a0.dist-info/RECORD,,
+agenta-0.20.0a0.dist-info/METADATA,sha256=xMd6YzEzMtRW5i6uO62QRRePYN0r91gLpIgaF8QrzIE,26462
+agenta-0.20.0a0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+agenta-0.20.0a0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.20.0a0.dist-info/RECORD,,