feat: python async function tracing in dev mode, closing OPEN-6157

Author: whoseoyster

src/openlayer/lib/core/base_model.py

@@ -1,13 +1,13 @@
 """Base class for an Openlayer model."""
 
-import os
 import abc
+import argparse
+import inspect
 import json
+import os
 import time
-import inspect
-import argparse
+from dataclasses import dataclass, field
 from typing import Any, Dict, Tuple
-from dataclasses import field, dataclass
 
 import pandas as pd
 
@@ -42,9 +42,7 @@ class OpenlayerModel(abc.ABC):
     def run_from_cli(self) -> None:
         """Run the model from the command line."""
         parser = argparse.ArgumentParser(description="Run data through a model.")
-        parser.add_argument(
-            "--dataset-path", type=str, required=True, help="Path to the dataset"
-        )
+        parser.add_argument("--dataset-path", type=str, required=True, help="Path to the dataset")
         parser.add_argument(
             "--output-dir",
             type=str,
@@ -85,9 +83,7 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
             # Filter row_dict to only include keys that are valid parameters
             # for the 'run' method
             row_dict = row.to_dict()
-            filtered_kwargs = {
-                k: v for k, v in row_dict.items() if k in run_signature.parameters
-            }
+            filtered_kwargs = {k: v for k, v in row_dict.items() if k in run_signature.parameters}
 
             # Call the run method with filtered kwargs
             output = self.run(**filtered_kwargs)

src/openlayer/lib/integrations/langchain_callback.py

@@ -2,7 +2,7 @@
 
 # pylint: disable=unused-argument
 import time
-from typing import Any, Dict, List, Union, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from langchain import schema as langchain_schema
 from langchain.callbacks.base import BaseCallbackHandler
@@ -35,9 +35,7 @@ def __init__(self, **kwargs: Any) -> None:
         self.metatada: Dict[str, Any] = kwargs or {}
 
     # noqa arg002
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> Any:
+    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> Any:
         """Run when LLM starts running."""
         pass
 
@@ -81,45 +79,32 @@ def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
         """Run on new LLM token. Only available when streaming is enabled."""
         pass
 
-    def on_llm_end(
-        self, response: langchain_schema.LLMResult, **kwargs: Any  # noqa: ARG002, E501
-    ) -> Any:
+    def on_llm_end(self, response: langchain_schema.LLMResult, **kwargs: Any) -> Any:  # noqa: ARG002, E501
         """Run when LLM ends running."""
         self.end_time = time.time()
         self.latency = (self.end_time - self.start_time) * 1000
 
         if response.llm_output and "token_usage" in response.llm_output:
-            self.prompt_tokens = response.llm_output["token_usage"].get(
-                "prompt_tokens", 0
-            )
-            self.completion_tokens = response.llm_output["token_usage"].get(
-                "completion_tokens", 0
-            )
+            self.prompt_tokens = response.llm_output["token_usage"].get("prompt_tokens", 0)
+            self.completion_tokens = response.llm_output["token_usage"].get("completion_tokens", 0)
             self.cost = self._get_cost_estimate(
                 num_input_tokens=self.prompt_tokens,
                 num_output_tokens=self.completion_tokens,
             )
-            self.total_tokens = response.llm_output["token_usage"].get(
-                "total_tokens", 0
-            )
+            self.total_tokens = response.llm_output["token_usage"].get("total_tokens", 0)
 
         for generations in response.generations:
             for generation in generations:
                 self.output += generation.text.replace("\n", " ")
 
         self._add_to_trace()
 
-    def _get_cost_estimate(
-        self, num_input_tokens: int, num_output_tokens: int
-    ) -> float:
+    def _get_cost_estimate(self, num_input_tokens: int, num_output_tokens: int) -> float:
         """Returns the cost estimate for a given model and number of tokens."""
         if self.model not in constants.OPENAI_COST_PER_TOKEN:
             return None
         cost_per_token = constants.OPENAI_COST_PER_TOKEN[self.model]
-        return (
-            cost_per_token["input"] * num_input_tokens
-            + cost_per_token["output"] * num_output_tokens
-        )
+        return cost_per_token["input"] * num_input_tokens + cost_per_token["output"] * num_output_tokens
 
     def _add_to_trace(self) -> None:
         """Adds to the trace."""
@@ -141,56 +126,42 @@ def _add_to_trace(self) -> None:
             metadata=self.metatada,
         )
 
-    def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
         """Run when LLM errors."""
         pass
 
-    def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> Any:
+    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> Any:
         """Run when chain starts running."""
         pass
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
         """Run when chain ends running."""
         pass
 
-    def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
         """Run when chain errors."""
         pass
 
-    def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
-    ) -> Any:
+    def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
         """Run when tool starts running."""
         pass
 
     def on_tool_end(self, output: str, **kwargs: Any) -> Any:
         """Run when tool ends running."""
         pass
 
-    def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> Any:
+    def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
         """Run when tool errors."""
         pass
 
     def on_text(self, text: str, **kwargs: Any) -> Any:
         """Run on arbitrary text."""
         pass
 
-    def on_agent_action(
-        self, action: langchain_schema.AgentAction, **kwargs: Any
-    ) -> Any:
+    def on_agent_action(self, action: langchain_schema.AgentAction, **kwargs: Any) -> Any:
         """Run on agent action."""
         pass
 
-    def on_agent_finish(
-        self, finish: langchain_schema.AgentFinish, **kwargs: Any
-    ) -> Any:
+    def on_agent_finish(self, finish: langchain_schema.AgentFinish, **kwargs: Any) -> Any:
         """Run on agent end."""
         pass

src/openlayer/lib/integrations/openai_tracer.py

@@ -1,10 +1,10 @@
 """Module with methods used to trace OpenAI / Azure OpenAI LLMs."""
 
 import json
-import time
 import logging
-from typing import Any, Dict, List, Union, Iterator, Optional
+import time
 from functools import wraps
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 import openai
 

src/openlayer/lib/tracing/steps.py

@@ -4,8 +4,8 @@
 import uuid
 from typing import Any, Dict, Optional
 
-from . import enums
 from .. import utils
+from . import enums
 
 
 class Step:

src/openlayer/lib/tracing/tracer.py

@@ -1,17 +1,18 @@
 """Module with the logic to create and manage traces and steps."""
 
-import time
+import asyncio
+import contextvars
 import inspect
 import logging
-import contextvars
-from typing import Any, Dict, List, Tuple, Optional, Generator
-from functools import wraps
+import time
 from contextlib import contextmanager
+from functools import wraps
+from typing import Any, Awaitable, Dict, Generator, List, Optional, Tuple
 
-from . import enums, steps, traces
-from .. import utils
 from ..._client import Openlayer
 from ...types.inference_pipelines.data_stream_params import ConfigLlmData
+from .. import utils
+from . import enums, steps, traces
 
 logger = logging.getLogger(__name__)
 
@@ -195,6 +196,100 @@ def wrapper(*func_args, **func_kwargs):
     return decorator
 
 
+def trace_async(*step_args, **step_kwargs):
+    """Decorator to trace a function.
+
+    Examples
+    --------
+
+    To trace a function, simply decorate it with the ``@trace()`` decorator. By doing so,
+    the functions inputs, outputs, and metadata will be automatically logged to your
+    Openlayer project.
+
+    >>> import os
+    >>> from openlayer.tracing import tracer
+    >>>
+    >>> # Set the environment variables
+    >>> os.environ["OPENLAYER_API_KEY"] = "YOUR_OPENLAYER_API_KEY_HERE"
+    >>> os.environ["OPENLAYER_PROJECT_NAME"] = "YOUR_OPENLAYER_PROJECT_NAME_HERE"
+    >>>
+    >>> # Decorate all the functions you want to trace
+    >>> @tracer.trace_async()
+    >>> async def main(user_query: str) -> str:
+    >>>     context = retrieve_context(user_query)
+    >>>     answer = generate_answer(user_query, context)
+    >>>     return answer
+    >>>
+    >>> @tracer.trace_async()
+    >>> def retrieve_context(user_query: str) -> str:
+    >>>     return "Some context"
+    >>>
+    >>> @tracer.trace_async()
+    >>> def generate_answer(user_query: str, context: str) -> str:
+    >>>     return "Some answer"
+    >>>
+    >>> # Every time the main function is called, the data is automatically
+    >>> # streamed to your Openlayer project. E.g.:
+    >>> tracer.run_async_func(main("What is the meaning of life?"))
+    """
+
+    def decorator(func):
+        func_signature = inspect.signature(func)
+
+        @wraps(func)
+        async def wrapper(*func_args, **func_kwargs):
+            if step_kwargs.get("name") is None:
+                step_kwargs["name"] = func.__name__
+            with create_step(*step_args, **step_kwargs) as step:
+                output = exception = None
+                try:
+                    output = await func(*func_args, **func_kwargs)
+                # pylint: disable=broad-except
+                except Exception as exc:
+                    step.log(metadata={"Exceptions": str(exc)})
+                    exception = exc
+                end_time = time.time()
+                latency = (end_time - step.start_time) * 1000  # in ms
+
+                bound = func_signature.bind(*func_args, **func_kwargs)
+                bound.apply_defaults()
+                inputs = dict(bound.arguments)
+                inputs.pop("self", None)
+                inputs.pop("cls", None)
+
+                step.log(
+                    inputs=inputs,
+                    output=output,
+                    end_time=end_time,
+                    latency=latency,
+                )
+
+                if exception is not None:
+                    raise exception
+            return output
+
+        return wrapper
+
+    return decorator
+
+
+async def _invoke_with_context(coroutine: Awaitable[Any]) -> Tuple[contextvars.Context, Any]:
+    """Runs a coroutine and preserves the context variables set within it."""
+    result = await coroutine
+    context = contextvars.copy_context()
+    return context, result
+
+
+def run_async_func(coroutine: Awaitable[Any]) -> Any:
+    """Runs an async function while preserving the context. This is needed
+    for tracing async functions.
+    """
+    context, result = asyncio.run(_invoke_with_context(coroutine))
+    for key, value in context.items():
+        key.set(value)
+    return result
+
+
 # --------------------- Helper post-processing functions --------------------- #
 def post_process_trace(
     trace_obj: traces.Trace,

src/openlayer/lib/utils.py

@@ -2,8 +2,8 @@
 Openlayer SDK.
 """
 
-import os
 import json
+import os
 from typing import Optional