agenta 0.25.2__py3-none-any.whl → 0.25.3a1__py3-none-any.whl

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

@@ -1,6 +1,5 @@
 """The code for the Agenta SDK"""
 
-from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
 import os
 import sys
 import time
@@ -14,18 +13,17 @@ from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import Any, Callable, Dict, Optional, Tuple, List
 from importlib.metadata import version
-
 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,
@@ -67,36 +65,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 +94,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,10 +128,11 @@ 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"
@@ -176,8 +153,9 @@ 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)
         async def wrapper(*args, **kwargs) -> Any:
             func_params, api_config_params = self.split_kwargs(kwargs, config_params)
@@ -185,11 +163,10 @@ class entrypoint(BaseDecorator):
             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,
+                environment="playground",
+            ):
                 entrypoint_result = await self.execute_function(
                     func,
                     True,  # inline trace: True
@@ -242,7 +219,6 @@ class entrypoint(BaseDecorator):
         ### ---------------------------- #
 
         ### --- Deployed / Published --- #
-        @debug()
         @functools.wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
             func_params = {
@@ -256,12 +232,10 @@ 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"]
+            with routing_context_manager(
+                config=config_params,
+                variant=kwargs["config"],
+                environment=kwargs["environment"],
             ):
                 entrypoint_result = await self.execute_function(
                     func,
@@ -351,86 +325,105 @@ class entrypoint(BaseDecorator):
             if name in func_params and func_params[name] is not None:
                 func_params[name] = self.ingest_file(func_params[name])
 
-    async def execute_function(
-        self, func: Callable[..., Any], inline_trace, *args, **func_params
-    ):
-        """Execute the function and handle any exceptions."""
-
-        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')}")
+    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.
 
-            WAIT_FOR_SPANS = True
-            TIMEOUT = 1
-            TIMESTEP = 0.1
-            NOFSTEPS = TIMEOUT / TIMESTEP
+        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
+        )
 
-            data = None
-            trace = None
+        if data is None:
+            data = (
+                "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
+            )
 
-            token = None
-            if tracing_context.get() is None:
-                token = tracing_context.set(TracingContext())
+        if not isinstance(result, dict):
+            data = str(data)
 
-            is_coroutine_function = inspect.iscoroutinefunction(func)
+        return data
 
-            if is_coroutine_function:
-                result = await func(*args, **func_params["params"])
-            else:
-                result = func(*args, **func_params["params"])
+    async def execute_function(
+        self,
+        func: Callable[..., Any],
+        inline_trace,
+        *args,
+        **func_params,
+    ):
+        log.info(f"\n--------------------------")
+        log.info(
+            f"Running application route: {repr(self.route_path if self.route_path != '' else '/')}"
+        )
+        log.info(f"--------------------------\n")
 
-            if token is not None:
-                if WAIT_FOR_SPANS:
-                    remaining_steps = NOFSTEPS
+        tracing_context.set(routing_context.get())
 
-                    while not ag.tracing.is_trace_ready() and remaining_steps > 0:
-                        await asyncio.sleep(TIMESTEP)
-                        remaining_steps -= 1
+        WAIT_FOR_SPANS = True
+        TIMEOUT = 1
+        TIMESTEP = 0.1
+        FINALSTEP = 0.001
+        NOFSTEPS = TIMEOUT / TIMESTEP
 
-                trace = ag.tracing.dump_trace()
+        data = None
+        trace = {}
 
-                if not inline_trace:
-                    trace = {"trace_id": trace["trace_id"]}
+        try:
+            result = (
+                await func(*args, **func_params["params"])
+                if inspect.iscoroutinefunction(func)
+                else func(*args, **func_params["params"])
+            )
 
-                ag.tracing.flush_spans()
-                tracing_context.reset(token)
+            data = self.patch_result(result)
+        except Exception as e:
+            self.handle_exception(e)
 
-            if isinstance(result, Context):
-                save_context(result)
+        with suppress():
+            root_context: Dict[str, Any] = tracing_context.get().get("root")
 
-            data = result
+            trace_id = root_context.get("trace_id") if root_context else None
 
-            # 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"])
+            if trace_id is not None:
+                if inline_trace:
+                    if WAIT_FOR_SPANS:
+                        remaining_steps = NOFSTEPS
 
-            # END OF PATH
+                        while (
+                            not ag.tracing.is_inline_trace_ready(trace_id)
+                            and remaining_steps > 0
+                        ):
+                            await asyncio.sleep(TIMESTEP)
 
-            if data is None:
-                data = (
-                    "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
-                )
+                            remaining_steps -= 1
 
-            response = BaseResponse(data=data, trace=trace)
+                        await asyncio.sleep(FINALSTEP)
 
-            # logging.debug(response)
+                    trace = ag.tracing.get_inline_trace(trace_id)
+                else:
+                    trace = {"trace_id": trace_id}
 
-            return response
+        response = BaseResponse(data=data, trace=trace)
 
-        except Exception as e:
-            self.handle_exception(e)
+        return response
 
     def handle_exception(self, e: Exception):
         status_code = e.status_code if hasattr(e, "status_code") else 500
@@ -644,28 +637,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 = 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_SPANS = True
+        SHOW_SPAN_ATTRIBUTES = False
+
+        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"tokens:   {list(result.trace.get('tokens', {}).values())}")
+
+        if SHOW_DATA:
+            log.info(" ")
+            log.info(f"data:")
+            log.info(json.dumps(result.data, indent=2))
 
-        print("-> data")
-        print(json.dumps(result.data, indent=2))
-        print("-> trace")
-        print(json.dumps(result.trace, indent=2))
+        if SHOW_SPANS:
+            log.info(" ")
+            log.info(f"trace:")
+            log.info(f"----------------")
+            log.info(json.dumps(result.trace.get("spans", []), indent=2))
+            log.info(f"----------------")
 
-        with open("trace.json", "w") as trace_file:
-            json.dump(result.trace, trace_file, indent=4)
+        log.info("\n==========================\n")
 
     def override_config_in_schema(
         self,