agenta 0.25.4a2__py3-none-any.whl → 0.25.4a4__py3-none-any.whl

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

@@ -1,5 +1,6 @@
 """The code for the Agenta SDK"""
 
+from agenta.sdk.utils.debug import debug, DEBUG, SHIFT
 import os
 import sys
 import time
@@ -13,17 +14,18 @@ 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.routing import routing_context_manager, routing_context
-from agenta.sdk.context.tracing import tracing_context
+from agenta.sdk.context import save_context
 from agenta.sdk.router import router as router
-from agenta.sdk.utils.exceptions import suppress
-from agenta.sdk.utils.logging import log
+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.types import (
+    Context,
     DictInput,
     FloatParam,
     InFile,
@@ -65,14 +67,36 @@ app.add_middleware(
 app.include_router(router, prefix="")
 
 
-log.setLevel("DEBUG")
+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)
 
 
 class PathValidator(BaseModel):
     url: HttpUrl
 
 
-class route:
+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
@@ -94,7 +118,7 @@ class route:
         return f
 
 
-class entrypoint:
+class entrypoint(BaseDecorator):
     """
     Decorator class to wrap a function for HTTP POST, terminal exposure and enable tracing.
 
@@ -128,11 +152,10 @@ class entrypoint:
     routes = list()
 
     def __init__(
-        self,
-        func: Callable[..., Any],
-        route_path="",
-        config_schema: Optional[BaseModel] = None,
+        self, func: Callable[..., Any], route_path="", config_schema: BaseModel = None
     ):
+        logging.info(f"Using Agenta Python SDK version {version('agenta')}")
+
         DEFAULT_PATH = "generate"
         PLAYGROUND_PATH = "/playground"
         RUN_PATH = "/run"
@@ -153,9 +176,8 @@ class entrypoint:
         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)
@@ -163,10 +185,11 @@ class entrypoint:
             if not config_schema:
                 ag.config.set(**api_config_params)
 
-            with routing_context_manager(
-                config=api_config_params,
-                environment="playground",
-            ):
+            # 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):
                 entrypoint_result = await self.execute_function(
                     func,
                     True,  # inline trace: True
@@ -219,6 +242,7 @@ class entrypoint:
         ### ---------------------------- #
 
         ### --- Deployed / Published --- #
+        @debug()
         @functools.wraps(func)
         async def wrapper_deployed(*args, **kwargs) -> Any:
             func_params = {
@@ -232,10 +256,12 @@ class entrypoint:
                 else:
                     ag.config.pull(config_name="default")
 
-            with routing_context_manager(
-                config=config_params,
-                variant=kwargs["config"],
-                environment=kwargs["environment"],
+            # 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"]
             ):
                 entrypoint_result = await self.execute_function(
                     func,
@@ -325,108 +351,86 @@ class entrypoint:
             if name in func_params and func_params[name] is not None:
                 func_params[name] = self.ingest_file(func_params[name])
 
-    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.
+    async def execute_function(
+        self, func: Callable[..., Any], inline_trace, *args, **func_params
+    ):
+        """Execute the function and handle any exceptions."""
 
-        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
-        )
+        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')}")
 
-        if data is None:
-            data = (
-                "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
-            )
+            WAIT_FOR_SPANS = True
+            TIMEOUT = 1
+            TIMESTEP = 0.1
+            NOFSTEPS = TIMEOUT / TIMESTEP
 
-        if not isinstance(result, dict):
-            data = str(data)
+            data = None
+            trace = None
 
-        return data
+            token = None
+            if tracing_context.get() is None:
+                token = tracing_context.set(TracingContext())
 
-    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")
+            is_coroutine_function = inspect.iscoroutinefunction(func)
 
-        tracing_context.set(routing_context.get())
+            if is_coroutine_function:
+                result = await func(*args, **func_params["params"])
+            else:
+                result = func(*args, **func_params["params"])
 
-        WAIT_FOR_SPANS = True
-        TIMEOUT = 1
-        TIMESTEP = 0.1
-        FINALSTEP = 0.001
-        NOFSTEPS = TIMEOUT / TIMESTEP
+            if token is not None:
+                if WAIT_FOR_SPANS:
+                    remaining_steps = NOFSTEPS
 
-        data = None
-        trace = {}
+                    while not ag.tracing.is_trace_ready() and remaining_steps > 0:
+                        await asyncio.sleep(TIMESTEP)
+                        remaining_steps -= 1
 
-        try:
-            result = (
-                await func(*args, **func_params["params"])
-                if inspect.iscoroutinefunction(func)
-                else func(*args, **func_params["params"])
-            )
-            data = self.patch_result(result)
-        except Exception as e:
-            log.error(f"Agenta SDK - Routing Exception")
+                trace = ag.tracing.dump_trace()
 
-            traceback.print_exc()
+                if not inline_trace:
+                    trace = {"trace_id": trace["trace_id"]}
 
-            self.handle_exception(e)
+                ag.tracing.flush_spans()
+                tracing_context.reset(token)
 
-        with suppress():
-            root_context: Dict[str, Any] = tracing_context.get().get("root")
+            if isinstance(result, Context):
+                save_context(result)
 
-            trace_id = root_context.get("trace_id") if root_context else None
+            data = result
 
-            if trace_id is not None:
-                if inline_trace:
-                    if WAIT_FOR_SPANS:
-                        remaining_steps = NOFSTEPS
+            # 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"])
 
-                        while (
-                            not ag.tracing.is_inline_trace_ready(trace_id)
-                            and remaining_steps > 0
-                        ):
-                            await asyncio.sleep(TIMESTEP)
+            # END OF PATH
 
-                            remaining_steps -= 1
+            if data is None:
+                data = (
+                    "Function executed successfully, but did return None. \n Are you sure you did not forget to return a value?",
+                )
 
-                        await asyncio.sleep(FINALSTEP)
+            response = BaseResponse(data=data, trace=trace)
 
-                    trace = ag.tracing.get_inline_trace(trace_id)
-                else:
-                    trace = {"trace_id": trace_id}
+            # logging.debug(response)
 
-        response = BaseResponse(data=data, trace=trace)
+            return response
 
-        return response
+        except Exception as e:
+            self.handle_exception(e)
 
     def handle_exception(self, e: Exception):
         status_code = e.status_code if hasattr(e, "status_code") else 500
@@ -640,45 +644,28 @@ class entrypoint:
             }
         )
 
+        # 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()
 
-        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},
-                )
+        result = loop.run_until_complete(
+            self.execute_function(
+                func,
+                True,  # inline trace: True
+                **{"params": args_func_params, "config_params": args_config_params},
             )
+        )
 
-        SHOW_DETAILS = True
-        SHOW_DATA = False
-        SHOW_TRACE = 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"usage:   {list(result.trace.get('usage', {}).values())}")
-
-        if SHOW_DATA:
-            log.info(" ")
-            log.info(f"data:")
-            log.info(json.dumps(result.data, indent=2))
+        print("\n========== Result ==========\n")
 
-        if SHOW_TRACE:
-            log.info(" ")
-            log.info(f"trace:")
-            log.info(f"----------------")
-            log.info(json.dumps(result.trace.get("spans", []), indent=2))
-            log.info(f"----------------")
+        print("-> data")
+        print(json.dumps(result.data, indent=2))
+        print("-> trace")
+        print(json.dumps(result.trace, indent=2))
 
-        log.info("\n==========================\n")
+        with open("trace.json", "w") as trace_file:
+            json.dump(result.trace, trace_file, indent=4)
 
     def override_config_in_schema(
         self,