Completes OPEN-5899 Add tracing documentation to reference.openlayer.com

Author: gustavocidornelas

docs/source/reference/monitoring.rst

@@ -31,6 +31,17 @@ Openlayer project is what enables the monitoring mode.
    Project.create_inference_pipeline
    Project.load_inference_pipeline
 
+Tracing
+-------
+If you have a multi-step system (e.g., RAG), you can trace all the steps in the system
+by decorating the functions with the ``@trace()`` decorator.
+
+.. autosummary::
+   :toctree: api/
+   :template: class.rst
+
+   openlayer.tracing.tracer.trace
+
 Publishing production data
 ----------------------------
 
@@ -45,6 +56,7 @@ single line of code.
    :template: class.rst
 
    openlayer.llm_monitors.OpenAIMonitor
+   openlayer.llm_monitors.AzureOpenAIMonitor
 
 Traditional ML models
 ^^^^^^^^^^^^^^^^^^^^^

openlayer/llm_monitors.py

@@ -19,7 +19,7 @@ class OpenAIMonitor:
 
     Parameters
     ----------
-    client : openai.api_client.Client, optional
+    client : openai.api_client.Client
         The OpenAI client. It is required if you are using openai>=1.0.0.
 
     Examples
@@ -43,7 +43,7 @@ class OpenAIMonitor:
     >>> from openai import OpenAI
     >>>
     >>> openai_client = OpenAI()
-    >>> monitor = llm_monitors.OpenAIMonitor(publish=True, client=openai_client)
+    >>> monitor = llm_monitors.OpenAIMonitor(client=openai_client)
 
     3. Use the OpenAI model as you normally would:
 
@@ -76,7 +76,7 @@ def __init__(
             )
 
     def start_monitoring(self) -> None:
-        """Start monitoring the OpenAI assistant."""
+        """(Deprecated) Start monitoring the OpenAI assistant."""
         warnings.warn(
             "The `start_monitoring` method is deprecated and will be removed in a future"
             " version. Monitoring is now automatically enabled once the OpenAIMonitor"
@@ -86,7 +86,7 @@ def start_monitoring(self) -> None:
         )
 
     def stop_monitoring(self) -> None:
-        """Stop monitoring the OpenAI assistant."""
+        """(Deprecated) Stop monitoring the OpenAI assistant."""
         warnings.warn(
             "The `stop_monitoring` method is deprecated and will be removed in a future"
             " version. Monitoring is now automatically enabled once the OpenAIMonitor"
@@ -417,7 +417,7 @@ def monitor_thread_run(self, run: "openai.types.beta.threads.run.Run") -> None:
             messages = self.openai_client.beta.threads.messages.list(
                 thread_id=run.thread_id, order="asc"
             )
-            prompt = self.thread_messages_to_prompt(messages)
+            prompt = self._thread_messages_to_prompt(messages)
 
             # Add step to the trace
             tracer.add_openai_chat_completion_step_to_trace(
@@ -466,7 +466,7 @@ def _extract_run_metadata(
         }
 
     @staticmethod
-    def thread_messages_to_prompt(
+    def _thread_messages_to_prompt(
         messages: List["openai.types.beta.threads.thread_message.ThreadMessage"],
     ) -> List[Dict[str, str]]:
         """Given list of ThreadMessage, return its contents in the `prompt` format,
@@ -493,7 +493,57 @@ def thread_messages_to_prompt(
 
 
 class AzureOpenAIMonitor(OpenAIMonitor):
-    """Monitor inferences from Azure OpenAI LLMs and upload traces to Openlayer."""
+    """Monitor inferences from Azure OpenAI LLMs and upload traces to Openlayer.
+
+    Parameters
+    ----------
+    client : openai.AzureOpenAI
+        The AzureOpenAI client.
+
+    Examples
+    --------
+
+    Let's say that you have a GPT model you want to monitor. You can turn on monitoring
+    with Openlayer by simply doing:
+
+    1. Set the environment variables:
+
+    .. code-block:: bash
+
+        export AZURE_OPENAI_ENDPOINT=<your-azure-openai-endpoint>
+        export AZURE_OPENAI_API_KEY=<your-azure-openai-api-key>
+        export AZURE_OPENAI_DEPLOYMENT_NAME=<your-azure-openai-deployment-name>
+
+        export OPENLAYER_API_KEY=<your-openlayer-api-key>
+        export OPENLAYER_PROJECT_NAME=<your-project-name>
+
+    2. Instantiate the monitor:
+
+    >>> from opemlayer import llm_monitors
+    >>> from openai import AzureOpenAI
+    >>>
+    >>> azure_client = AzureOpenAI(
+    >>>     api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
+    >>>     api_version="2024-02-01",
+    >>>     azure_endpoint=os.environ.get("AZURE_OPENAI_ENDPOINT"),
+    >>> )
+    >>> monitor = llm_monitors.OpenAIMonitor(client=azure_client)
+
+    3. Use the Azure OpenAI model as you normally would:
+
+    From this point onwards, you can continue making requests to your model normally:
+
+    >>> completion = azure_client.chat.completions.create(
+    >>>     model=os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME"),
+    >>>     messages=[
+    >>>         {"role": "system", "content": "You are a helpful assistant."},
+    >>>         {"role": "user", "content": "How are you doing today?"},
+    >>>     ]
+    >>> )
+
+    The trace of this inference request is automatically uploaded to your Openlayer
+    project.
+    """
 
     def __init__(
         self,

openlayer/tracing/tracer.py

@@ -110,6 +110,42 @@ def add_openai_chat_completion_step_to_trace(**kwargs) -> None:
 
 # ----------------------------- Tracing decorator ---------------------------- #
 def trace(*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()
+    >>> def main(user_query: str) -> str:
+    >>>     context = retrieve_context(user_query)
+    >>>     answer = generate_answer(user_query, context)
+    >>>     return answer
+    >>>
+    >>> @tracer.trace()
+    >>> def retrieve_context(user_query: str) -> str:
+    >>>     return "Some context"
+    >>>
+    >>> @tracer.trace()
+    >>> 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.:
+    >>> main("What is the meaning of life?")
+    """
+
     def decorator(func):
         func_signature = inspect.signature(func)