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

agenta/sdk/decorators/{llm_entrypoint.py → routing.py}

@@ -1,31 +1,25 @@
-"""The code for the Agenta SDK"""
-
-from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
-import os
-import sys
-import time
-import json
-import inspect
-import argparse
-import asyncio
-import traceback
-import functools
+from typing import Type, Any, Callable, Dict, Optional, Tuple, List
+from annotated_types import Ge, Le, Gt, Lt
+from pydantic import BaseModel, HttpUrl, ValidationError
+from json import dumps
+from inspect import signature, iscoroutinefunction, Signature, Parameter, _empty
+from argparse import ArgumentParser
+from functools import wraps
+from asyncio import sleep, get_event_loop
+from traceback import format_exc, format_exception
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Any, Callable, Dict, Optional, Tuple, List
-from importlib.metadata import version
+from os import environ
 
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import Body, FastAPI, UploadFile, HTTPException
 
-import agenta as ag
-from agenta.sdk.context import save_context
+from agenta.sdk.context.routing import routing_context_manager, routing_context
+from agenta.sdk.context.tracing import tracing_context
 from agenta.sdk.router import router as router
-from agenta.sdk.tracing.logger import llm_logger as logging
-from agenta.sdk.tracing.tracing_context import tracing_context, TracingContext
-from agenta.sdk.decorators.base import BaseDecorator
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.utils.logging import log
 from agenta.sdk.types import (
-    Context,
     DictInput,
     FloatParam,
     InFile,
@@ -39,16 +33,8 @@ from agenta.sdk.types import (
     BaseResponse,
     BinaryParam,
 )
-import pydantic
-
-from pydantic import BaseModel
-from typing import Type
-from annotated_types import Ge, Le, Gt, Lt
-
-from pydantic import BaseModel, HttpUrl
-import contextvars
-from contextlib import contextmanager
 
+import agenta as ag
 
 app = FastAPI()
 
@@ -67,36 +53,14 @@ app.add_middleware(
 app.include_router(router, prefix="")
 
 
-logging.setLevel("DEBUG")
-
-route_context = contextvars.ContextVar("route_context", default={})
-
-
-@contextmanager
-def route_context_manager(
-    config: Optional[Dict[str, Any]] = None,
-    environment: Optional[str] = None,
-    version: Optional[str] = None,
-    variant: Optional[str] = None,
-):
-    context = {
-        "config": config,
-        "environment": environment,
-        "version": version,
-        "variant": variant,
-    }
-    token = route_context.set(context)
-    try:
-        yield
-    finally:
-        route_context.reset(token)
+log.setLevel("DEBUG")
 
 
 class PathValidator(BaseModel):
     url: HttpUrl
 
 
-class route(BaseDecorator):
+class route:
     # 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
@@ -118,7 +82,7 @@ class route(BaseDecorator):
         return f
 
 
-class entrypoint(BaseDecorator):
+class entrypoint:
     """
     Decorator class to wrap a function for HTTP POST, terminal exposure and enable tracing.
 
@@ -152,19 +116,20 @@ class entrypoint(BaseDecorator):
     routes = list()
 
     def __init__(
-        self, func: Callable[..., Any], route_path="", config_schema: BaseModel = None
+        self,
+        func: Callable[..., Any],
+        route_path="",
+        config_schema: Optional[BaseModel] = None,
     ):
-        logging.info(f"Using Agenta Python SDK version {version('agenta')}")
-
         DEFAULT_PATH = "generate"
         PLAYGROUND_PATH = "/playground"
         RUN_PATH = "/run"
-        func_signature = inspect.signature(func)
+        func_signature = signature(func)
         try:
             config = (
                 config_schema() if config_schema else None
             )  # we initialize the config object to be able to use it
-        except pydantic.ValidationError as e:
+        except ValidationError as e:
             raise ValueError(
                 f"Error initializing config_schema. Please ensure all required fields have default values: {str(e)}"
             ) from e
@@ -176,20 +141,19 @@ class entrypoint(BaseDecorator):
         config_params = config.dict() if config else ag.config.all()
         ingestible_files = self.extract_ingestible_files(func_signature)
 
+        self.route_path = route_path
+
         ### --- Playground  --- #
-        @debug()
-        @functools.wraps(func)
+        @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)
             if not config_schema:
                 ag.config.set(**api_config_params)
 
-            # Set the configuration and environment of the LLM app parent span at run-time
-            ag.tracing.update_baggage(
-                {"config": config_params, "environment": "playground"}
-            )
-            with route_context_manager(config=api_config_params):
+            with routing_context_manager(
+                config=api_config_params,
+            ):
                 entrypoint_result = await self.execute_function(
                     func,
                     True,  # inline trace: True
@@ -241,12 +205,13 @@ class entrypoint(BaseDecorator):
         )
         ### ---------------------------- #
 
-        ### --- Deployed / Published --- #
-        @debug()
-        @functools.wraps(func)
+        ### --- Deployed --- #
+        @wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
             func_params = {
-                k: v for k, v in kwargs.items() if k not in ["config", "environment"]
+                k: v
+                for k, v in kwargs.items()
+                if k not in ["config", "environment", "app"]
             }
             if not config_schema:
                 if "environment" in kwargs and kwargs["environment"] is not None:
@@ -256,12 +221,19 @@ class entrypoint(BaseDecorator):
                 else:
                     ag.config.pull(config_name="default")
 
-            # Set the configuration and environment of the LLM app parent span at run-time
-            ag.tracing.update_baggage(
-                {"config": config_params, "environment": kwargs["environment"]}
-            )
-            with route_context_manager(
-                variant=kwargs["config"], environment=kwargs["environment"]
+            app_id = environ.get("AGENTA_APP_ID")
+
+            with routing_context_manager(
+                application={
+                    "id": app_id,
+                    "slug": kwargs["app"],
+                },
+                variant={
+                    "slug": kwargs.get("config"),
+                },
+                environment={
+                    "slug": kwargs.get("environment"),
+                },
             ):
                 entrypoint_result = await self.execute_function(
                     func,
@@ -284,7 +256,7 @@ class entrypoint(BaseDecorator):
 
         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
@@ -315,8 +287,8 @@ class entrypoint(BaseDecorator):
 
     def extract_ingestible_files(
         self,
-        func_signature: inspect.Signature,
-    ) -> Dict[str, inspect.Parameter]:
+        func_signature: Signature,
+    ) -> Dict[str, Parameter]:
         """Extract parameters annotated as InFile from function signature."""
 
         return {
@@ -343,7 +315,7 @@ class entrypoint(BaseDecorator):
     def ingest_files(
         self,
         func_params: Dict[str, Any],
-        ingestible_files: Dict[str, inspect.Parameter],
+        ingestible_files: Dict[str, Parameter],
     ) -> None:
         """Ingest files specified in function parameters."""
 
@@ -352,96 +324,134 @@ class entrypoint(BaseDecorator):
                 func_params[name] = self.ingest_file(func_params[name])
 
     async def execute_function(
-        self, func: Callable[..., Any], inline_trace, *args, **func_params
+        self,
+        func: Callable[..., Any],
+        inline_trace,
+        *args,
+        **func_params,
     ):
-        """Execute the function and handle any exceptions."""
+        log.info(f"---------------------------")
+        log.info(f"Agenta SDK - running route: {repr(self.route_path or '/')}")
+        log.info(f"---------------------------")
+
+        tracing_context.set(routing_context.get())
 
         try:
-            """Note: The following block is for backward compatibility.
-            It allows functions to work seamlessly whether they are synchronous or asynchronous.
-            For synchronous functions, it calls them directly, while for asynchronous functions,
-            it awaits their execution.
-            """
-            logging.info(f"Using Agenta Python SDK version {version('agenta')}")
+            result = (
+                await func(*args, **func_params["params"])
+                if iscoroutinefunction(func)
+                else func(*args, **func_params["params"])
+            )
 
-            WAIT_FOR_SPANS = True
-            TIMEOUT = 1
-            TIMESTEP = 0.1
-            NOFSTEPS = TIMEOUT / TIMESTEP
+            return await self.handle_success(result, inline_trace)
 
-            data = None
-            trace = None
+        except Exception as error:
+            self.handle_failure(error)
 
-            token = None
-            if tracing_context.get() is None:
-                token = tracing_context.set(TracingContext())
+    async def handle_success(self, result: Any, inline_trace: bool):
+        data = None
+        trace = dict()
 
-            is_coroutine_function = inspect.iscoroutinefunction(func)
+        with suppress():
+            data = self.patch_result(result)
 
-            if is_coroutine_function:
-                result = await func(*args, **func_params["params"])
-            else:
-                result = func(*args, **func_params["params"])
+            if inline_trace:
+                trace = await self.fetch_inline_trace(inline_trace)
 
-            if token is not None:
-                if WAIT_FOR_SPANS:
-                    remaining_steps = NOFSTEPS
+        log.info(f"----------------------------------")
+        log.info(f"Agenta SDK - exiting with success: 200")
+        log.info(f"----------------------------------")
 
-                    while not ag.tracing.is_trace_ready() and remaining_steps > 0:
-                        await asyncio.sleep(TIMESTEP)
-                        remaining_steps -= 1
+        return BaseResponse(data=data, trace=trace)
 
-                trace = ag.tracing.dump_trace()
+    def handle_failure(self, error: Exception):
+        log.error("--------------------------------------------------")
+        log.error("Agenta SDK - handling application exception below:")
+        log.error("--------------------------------------------------")
+        log.error(format_exc().strip("\n"))
+        log.error("--------------------------------------------------")
 
-                if not inline_trace:
-                    trace = {"trace_id": trace["trace_id"]}
+        status_code = error.status_code if hasattr(error, "status_code") else 500
+        message = str(error)
+        stacktrace = format_exception(error, value=error, tb=error.__traceback__)  # type: ignore
+        detail = {"message": message, "stacktrace": stacktrace}
 
-                ag.tracing.flush_spans()
-                tracing_context.reset(token)
+        log.error(f"----------------------------------")
+        log.error(f"Agenta SDK - exiting with failure: {status_code}")
+        log.error(f"----------------------------------")
 
-            if isinstance(result, Context):
-                save_context(result)
+        raise HTTPException(status_code=status_code, detail=detail)
 
-            data = result
+    def patch_result(self, result: Any):
+        """
+        Patch the result to only include the message if the result is a FuncResponse-style dictionary with message, cost, and usage keys.
 
-            # PATCH : if result is not a dict, make it a dict
-            if not isinstance(result, dict):
-                data = str(result)
-            else:
-                # PATCH : if result is a legacy dict, clean it up
-                if (
-                    "message" in result.keys()
-                    and "cost" in result.keys()
-                    and "usage" in result.keys()
-                ):
-                    data = str(result["message"])
+        Example:
+        ```python
+        result = {
+            "message": "Hello, world!",
+            "cost": 0.5,
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 20,
+                "total_tokens": 30
+            }
+        }
+        result = patch_result(result)
+        print(result)
+        # Output: "Hello, world!"
+        ```
+        """
+        data = (
+            result["message"]
+            if isinstance(result, dict)
+            and all(key in result for key in ["message", "cost", "usage"])
+            else result
+        )
 
-            # END OF PATH
+        if data is None:
+            data = (
+                "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
+            )
 
-            if data is None:
-                data = (
-                    "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
-                )
+        if not isinstance(result, dict):
+            data = str(data)
 
-            response = BaseResponse(data=data, trace=trace)
+        return data
 
-            # logging.debug(response)
+    async def fetch_inline_trace(self, inline_trace):
+        WAIT_FOR_SPANS = True
+        TIMEOUT = 1
+        TIMESTEP = 0.1
+        FINALSTEP = 0.001
+        NOFSTEPS = TIMEOUT / TIMESTEP
 
-            return response
+        trace = None
 
-        except Exception as e:
-            self.handle_exception(e)
+        root_context: Dict[str, Any] = tracing_context.get().get("root")
 
-    def handle_exception(self, e: Exception):
-        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}
+        trace_id = root_context.get("trace_id") if root_context else None
 
-        raise HTTPException(
-            status_code=status_code,
-            detail=detail,
-        )
+        if trace_id is not None:
+            if inline_trace:
+                if WAIT_FOR_SPANS:
+                    remaining_steps = NOFSTEPS
+
+                    while (
+                        not ag.tracing.is_inline_trace_ready(trace_id)
+                        and remaining_steps > 0
+                    ):
+                        await sleep(TIMESTEP)
+
+                        remaining_steps -= 1
+
+                    await sleep(FINALSTEP)
+
+                trace = ag.tracing.get_inline_trace(trace_id)
+            else:
+                trace = {"trace_id": trace_id}
+
+        return trace
 
     def update_wrapper_signature(
         self, wrapper: Callable[..., Any], updated_params: List
@@ -451,25 +461,25 @@ class entrypoint(BaseDecorator):
 
         Args:
             wrapper (callable): A callable object, such as a function or a method, that requires a signature update.
-            updated_params (List[inspect.Parameter]): A list of `inspect.Parameter` objects representing the updated parameters
+            updated_params (List[Parameter]): A list of `Parameter` objects representing the updated parameters
                 for the wrapper function.
         """
 
-        wrapper_signature = inspect.signature(wrapper)
+        wrapper_signature = signature(wrapper)
         wrapper_signature = wrapper_signature.replace(parameters=updated_params)
         wrapper.__signature__ = wrapper_signature  # type: ignore
 
     def update_function_signature(
         self,
         wrapper: Callable[..., Any],
-        func_signature: inspect.Signature,
+        func_signature: Signature,
         config_class: Type[BaseModel],  # TODO: change to our type
         config_dict: Dict[str, Any],
-        ingestible_files: Dict[str, inspect.Parameter],
+        ingestible_files: Dict[str, Parameter],
     ) -> None:
         """Update the function signature to include new parameters."""
 
-        updated_params: List[inspect.Parameter] = []
+        updated_params: List[Parameter] = []
         if config_class:
             self.add_config_params_to_parser(updated_params, config_class)
         else:
@@ -480,21 +490,21 @@ class entrypoint(BaseDecorator):
     def update_deployed_function_signature(
         self,
         wrapper: Callable[..., Any],
-        func_signature: inspect.Signature,
-        ingestible_files: Dict[str, inspect.Parameter],
+        func_signature: Signature,
+        ingestible_files: Dict[str, Parameter],
     ) -> None:
         """Update the function signature to include new parameters."""
 
-        updated_params: List[inspect.Parameter] = []
+        updated_params: List[Parameter] = []
         self.add_func_params_to_parser(updated_params, func_signature, ingestible_files)
         for param in [
             "config",
             "environment",
         ]:  # we add the config and environment parameters
             updated_params.append(
-                inspect.Parameter(
+                Parameter(
                     name=param,
-                    kind=inspect.Parameter.KEYWORD_ONLY,
+                    kind=Parameter.KEYWORD_ONLY,
                     default=Body(None),
                     annotation=str,
                 )
@@ -508,9 +518,9 @@ class entrypoint(BaseDecorator):
         for name, field in config_class.__fields__.items():
             assert field.default is not None, f"Field {name} has no default value"
             updated_params.append(
-                inspect.Parameter(
+                Parameter(
                     name=name,
-                    kind=inspect.Parameter.KEYWORD_ONLY,
+                    kind=Parameter.KEYWORD_ONLY,
                     annotation=field.annotation.__name__,
                     default=Body(field.default),
                 )
@@ -525,9 +535,9 @@ class entrypoint(BaseDecorator):
                 len(param.__class__.__bases__) == 1
             ), f"Inherited standard type of {param.__class__} needs to be one."
             updated_params.append(
-                inspect.Parameter(
+                Parameter(
                     name=name,
-                    kind=inspect.Parameter.KEYWORD_ONLY,
+                    kind=Parameter.KEYWORD_ONLY,
                     default=Body(param),
                     annotation=param.__class__.__bases__[
                         0
@@ -540,23 +550,23 @@ class entrypoint(BaseDecorator):
     def add_func_params_to_parser(
         self,
         updated_params: list,
-        func_signature: inspect.Signature,
-        ingestible_files: Dict[str, inspect.Parameter],
+        func_signature: Signature,
+        ingestible_files: Dict[str, Parameter],
     ) -> None:
         """Add function parameters to function signature."""
         for name, param in func_signature.parameters.items():
             if name in ingestible_files:
                 updated_params.append(
-                    inspect.Parameter(name, param.kind, annotation=UploadFile)
+                    Parameter(name, param.kind, annotation=UploadFile)
                 )
             else:
                 assert (
                     len(param.default.__class__.__bases__) == 1
                 ), f"Inherited standard type of {param.default.__class__} needs to be one."
                 updated_params.append(
-                    inspect.Parameter(
+                    Parameter(
                         name,
-                        inspect.Parameter.KEYWORD_ONLY,
+                        Parameter.KEYWORD_ONLY,
                         default=Body(..., embed=True),
                         annotation=param.default.__class__.__bases__[
                             0
@@ -585,7 +595,7 @@ class entrypoint(BaseDecorator):
     def handle_terminal_run(
         self,
         func: Callable,
-        func_params: Dict[str, inspect.Parameter],
+        func_params: Dict[str, Parameter],
         config_params: Dict[str, Any],
         ingestible_files: Dict,
     ):
@@ -599,7 +609,7 @@ class entrypoint(BaseDecorator):
         """
 
         # For required parameters, we add them as arguments
-        parser = argparse.ArgumentParser()
+        parser = ArgumentParser()
         for name, param in func_params.items():
             if name in ingestible_files:
                 parser.add_argument(name, type=str)
@@ -644,28 +654,46 @@ class entrypoint(BaseDecorator):
             }
         )
 
-        # Set the configuration and environment of the LLM app parent span at run-time
-        ag.tracing.update_baggage({"config": args_config_params, "environment": "bash"})
+        loop = get_event_loop()
 
-        loop = asyncio.get_event_loop()
-
-        result = loop.run_until_complete(
-            self.execute_function(
-                func,
-                True,  # inline trace: True
-                **{"params": args_func_params, "config_params": args_config_params},
+        with routing_context_manager(
+            config=args_config_params,
+            environment="terminal",
+        ):
+            result = loop.run_until_complete(
+                self.execute_function(
+                    func,
+                    True,  # inline trace: True
+                    **{"params": args_func_params, "config_params": args_config_params},
+                )
             )
-        )
 
-        print("\n========== Result ==========\n")
+        SHOW_DETAILS = True
+        SHOW_DATA = False
+        SHOW_TRACE = False
+
+        if result.trace:
+            log.info("\n========= Result =========\n")
+
+            log.info(f"trace_id: {result.trace['trace_id']}")
+            if SHOW_DETAILS:
+                log.info(f"latency:  {result.trace.get('latency')}")
+                log.info(f"cost:     {result.trace.get('cost')}")
+                log.info(f"usage:   {list(result.trace.get('usage', {}).values())}")
 
-        print("-> data")
-        print(json.dumps(result.data, indent=2))
-        print("-> trace")
-        print(json.dumps(result.trace, indent=2))
+            if SHOW_DATA:
+                log.info(" ")
+                log.info(f"data:")
+                log.info(dumps(result.data, indent=2))
 
-        with open("trace.json", "w") as trace_file:
-            json.dump(result.trace, trace_file, indent=4)
+            if SHOW_TRACE:
+                log.info(" ")
+                log.info(f"trace:")
+                log.info(f"----------------")
+                log.info(dumps(result.trace.get("spans", []), indent=2))
+                log.info(f"----------------")
+
+            log.info("\n==========================\n")
 
     def override_config_in_schema(
         self,
@@ -872,10 +900,7 @@ class entrypoint(BaseDecorator):
                 subschema["maximum"] = param_val.maxval  # type: ignore
                 subschema["default"] = param_val
 
-            elif (
-                isinstance(param_val, inspect.Parameter)
-                and param_val.annotation is DictInput
-            ):
+            elif isinstance(param_val, Parameter) and param_val.annotation is DictInput:
                 subschema = find_in_schema(
                     param_val.annotation.__schema_type_properties__(),
                     schema_to_override,
@@ -894,7 +919,7 @@ class entrypoint(BaseDecorator):
                 subschema["default"] = param_val
 
             elif (
-                isinstance(param_val, inspect.Parameter)
+                isinstance(param_val, Parameter)
                 and param_val.annotation is MessagesInput
             ):
                 subschema = find_in_schema(
@@ -906,7 +931,7 @@ class entrypoint(BaseDecorator):
                 subschema["default"] = param_val.default
 
             elif (
-                isinstance(param_val, inspect.Parameter)
+                isinstance(param_val, Parameter)
                 and param_val.annotation is FileInputURL
             ):
                 subschema = find_in_schema(
@@ -930,6 +955,6 @@ class entrypoint(BaseDecorator):
                     "title": str(param_name).capitalize(),
                     "type": get_type_from_param(param_val),
                 }
-                if param_val.default != inspect._empty:
+                if param_val.default != _empty:
                     subschema["default"] = param_val.default  # type: ignore
                 schema_to_override[param_name] = subschema