Merge docs changes and code style

Makefile

@@ -0,0 +1,19 @@
+.PHONY: help doc doc-deploy
+
+help:
+	@echo "Makefile for managing Jupyter Book documentation and deployment"
+	@echo ""
+	@echo "Usage:"
+	@echo "  make doc            - Build the documentation"
+	@echo "  make doc-deploy     - Deploy the documentation to GitHub Pages"
+	@echo "  make help           - Display this help message"
+	@echo ""
+	@echo "For more information, refer to the README or script documentation."
+
+doc:
+	@echo "Building documentation..."
+	@bash ./docs/jupyter_build.sh
+
+doc-deploy:
+	@echo "Deploying documentation to GitHub Pages..."
+	@ghp-import -n -p -f docs/_build/html

docs/_config.yml

@@ -4,7 +4,7 @@
 title: Trace
 author: Ching-An Cheng, Allen Nie, Adith Swaminathan
 copyright: "2024 Trace Team"
-# logo: logo.png
+logo: trace_logo.png
 only_build_toc_files: true
 
 # Force re-execution of notebooks on each build.
@@ -45,24 +45,26 @@ sphinx:
   html_context:
     default_mode: light
   extra_extensions:
-  - sphinx_plausible
-  - autodoc2
+  - 'sphinx_plausible'
   - 'sphinx.ext.autodoc'
   - 'sphinx.ext.napoleon'
+  - 'sphinx.ext.autosummary'
   - 'sphinx.ext.viewcode'
   config:
     add_module_names: false
     plausible_domain: microsoft.github.io/trace
     nb_merge_streams: true
-    autodoc2_index_template: null
-    autodoc2_output_dir: api
-    autodoc2_module_all_regexes:
-      - "opto.*(trace|optimizers)"
-    autodoc2_packages:
-      - path: "../opto"
-    autodoc2_hidden_objects:
-      - inherited
-      - private
-      - dunder
-    autodoc2_skip_module_regexes:
-      - .*test.*
+    templates_path: ["_templates"]
+    autosummary_generate: True
+    autodoc_mock_imports: ['autogen']
+    suppress_warnings: ["etoc.toctree"]
+    # autodoc settings
+    # ref: https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#configuration
+    autoclass_content:                 both
+    autodoc_class_signature:           separated
+    autodoc_member_order:              groupwise
+    autodoc_docstring_signature:       True
+    autodoc_typehints:                 signature
+    autodoc_typehints_format:          short
+    autosummary_filename_map:
+      opto.trace.nodes.node: "opto.trace.nodes.node-function"

docs/_toc.yml

@@ -11,12 +11,6 @@ parts:
     - file: quickstart/quick_start
     - file: quickstart/quick_start_2
     - file: quickstart/virtualhome
-
-  - caption: FAQ
-    numbered: false
-    chapters:
-      - file: faq/faq
-
   - caption: 📚Tutorials
     chapters:
     - file: tutorials/basic_tutorial
@@ -41,6 +35,11 @@ parts:
     chapters:
     - file: examples/robotics/metaworld
 
+  - caption: FAQ
+    numbered: false
+    chapters:
+      - file: faq/faq
+
   - caption: 📖 API Reference
     chapters:
-    - glob: api/opto/*
+    - file: api

docs/api.rst

@@ -0,0 +1,6 @@
+.. autosummary::
+   :toctree: api/
+   :template: module.rst_t
+   :recursive:
+
+   opto

docs/faq/faq.md

@@ -1,6 +1,6 @@
 # FAQ
 
-### Difference to Libraries like TextGrad
+## Difference to Libraries like TextGrad
 
 TextGrad is both a library and an optimizer algorithm. Currently, we support three optimizers:
 
@@ -43,5 +43,5 @@ We provide a comparison to validate our implementation of TextGrad in Trace:
 To produce this table, we ran the TextGrad pip-installed repo on 2024-10-30, and we also include the numbers reported in the TextGrad paper.
 The LLM APIs are called around the same time to ensure a fair comparison. TextGrad paper's result was reported in 2024-06.
 
-### Difference to Libraries like AutoGen, AG2, OpenAI Swarm, Llama Stack
+## Difference to Libraries like AutoGen, AG2, OpenAI Swarm, Llama Stack
 

docs/intro.md

@@ -3,6 +3,9 @@
 **Trace is a Python library for tracing and optimizing workflows end-to-end by using LLM-powered generative optimizers.**
 **It can record *traces* of operations on any Python objects and functions, and automatically construct an execution graph that is useful when LLMs are used as optimizers.**
 
+
+<a href="https://colab.research.google.com/github/microsoft/Trace/blob/experimental/docs/examples/basic/greeting.ipynb" rel="nofollow" target="_blank"><img src="https://camo.githubusercontent.com/96889048f8a9014fdeba2a891f97150c6aac6e723f5190236b10215a97ed41f3/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667" alt="Open In Colab" data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" style="width: 120px;"></a>
+
 Our implementation is minimal and purely based on Python. It does not involve any API calls or library-specific dependencies, so it is composable with other libraries and tools. 
 Trace features an API design inspired by PyTorch Autograd's gradient tape mechanism, which we adopted to reduce the learning curve of using Trace. 
 These features make Trace an intuitive and flexible framework for building self-adapting AI agents.
@@ -25,14 +28,59 @@ This step is the **declare** phase where a user chooses how to represent the age
 After the user has declared the inputs and operations, Trace captures the execution flow of the program as a graph. This step is the **forward** phase.
 Finally, the user can optimize the entire program, such as by updating the LLM instructions, using Trace. This step is the **optimize** phase.
 
+```python
+@trace.model
+class Agent:
+
+    def __init__(self, system_prompt):
+        self.system_prompt = system_prompt
+        self.instruct1 = trace.node("Decide the language", trainable=True)
+        self.instruct2 = trace.node("Extract name",  trainable=True)
+
+    def __call__(self, user_query):
+        # First LLM 
+        response = call_llm(self.system_prompt, self.instruct1, user_query)
+        en_or_es = self.decide_lang(response)
+        # Second LLM 
+        user_name = call_llm(self.system_prompt, self.instruct2, user_query)
+        greeting = self.greet(en_or_es, user_name)
+        return greeting
+
+    @trace.bundle(trainable=True)
+    def decide_lang(self, response):
+        """Map the language into a variable"""
+        return
+
+    @trace.bundle(trainable=True)
+    def greet(self, lang, user_name):
+        """Produce a greeting based on the language"""
+        greeting = "Hola"
+        return f"{greeting}, {user_name}!"
+```
+
 Each application of Trace is defined by an **agent**, a source of **feedback**, and an **optimizer**.
 Enabling traces of operations on Python objects allows us to capture the execution flow of an agent, including AI systems that involve LLMs.
 In the example below, we show how Trace can optimize an entire AI system end-to-end.
 
+```python
+agent = Agent("You are a sales assistant.")
+optimizer = OptoPrime(agent.parameters())
 
+try:
+    greeting = agent("Hola, soy Juan.")
+    feedback = feedback_fn(greeting.data, 'es')
+    # feedback = "Correct" or "Incorrect"
+except ExecutionError as e:
+    greeting = e.exception_node
+    feedback = greeting.data, 
+
+optimizer.zero_feedback()
+optimizer.backward(greeting, feedback)
+optimizer.step()
+```
 
 ----
 
 
 <!-- ```{tableofcontents}
-``` -->
+``` -->

docs/jupyter_build.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+cd "$(dirname "$0")/.." || exit
+rm -r docs/_build docs/api
+ORIGINAL_PYTHONPATH=$PYTHONPATH
+export PYTHONPATH=$(pwd)/..:$PYTHONPATH
+
+jupyter-book build docs
+
+# clean up sphinx-autosummary generated files
+rm -r docs/api
+
+# Restored PYTHONPATH
+export PYTHONPATH=$ORIGINAL_PYTHONPATH
+
+# move all files associated with the landing page into the `_build/html` folder
+python docs/post_build_script.py

docs/quickstart/installation.md

@@ -12,10 +12,10 @@ then we require `autogen` package to make LLM API calls.
 To install Trace, run: 
 
 ```{admonition} Installation Command
+
 ```bash
 pip install trace-opt
 ```
-```
 
 To contribute to the development, you can clone the repository and install the package in editable mode:
 

docs/readme.md

@@ -2,19 +2,16 @@ Steps of deployment:
 
 IMPORTANT: checkout the `website` branch.
 
-1. Run `jupyter-book build docs` under the root directory to build the book. This will create a folder `_build/html` that has the static webpages.
-2. Run `python docs/post_build_script.py` to move all files associated with the landing page into the `_build/html` folder.
-3. Run `ghp-import -n -p -f docs/_build/html` to deploy the book to GitHub Pages (it creates a branch in the repo)
-
-Or simply run `docs/publish.sh` to run all the above commands.
+1. Run `make doc` under the root directory to build the book. This will create a folder `docs/_build/html` that has the static webpages.
+2. Run `make doc-deploy` to deploy the book to GitHub Pages (it creates a branch in the repo)
 
 References:
 
-https://jupyterbook.org/en/stable/start/publish.html 
+https://jupyterbook.org/en/stable/start/publish.html
 
 A few notes:
 1. There is no direct way to add an HTML page to Jupyter book.
-2. Run `pip install ghp-import` before step 3.
+2. Run `pip install -r requirements.txt` to install dependencies.
 3. Do not manually modify `gh-pages` branch.
 
 

docs/requirements.txt

@@ -1,3 +1,8 @@
 jupyter-book
 matplotlib
 numpy
+sphinx
+sphinx-plausible
+# sphinx-autodoc2
+sphinx-autoapi
+ghp-import

docs/tutorials/custom_optimizers.ipynb

@@ -13,7 +13,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Basic back-propagation and gradient descent with PyTorch\n",
+    "## Basic back-propagation and gradient descent with PyTorch\n",
     "\n",
     "To start, let's define a simple objective and run vanilla gradient descent to optimize the variable in pytorch. This code will be used as the reference of desired behaviors. We make the code below transparent for tutorial purppose, so we use the `torch.autograd.grad` api and write down the gradient descent update rule manually."
    ]
@@ -73,7 +73,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Set up the objective in Trace\n",
+    "## Set up the objective in Trace\n",
     "\n",
     "After seeing how ideally basic gradient descent + back-propagation behaves, next we show how it can be implemented it in Trace. To this end, we need to turn each math ops used in the above loss as a `bundle`, and define the parameter as a `node`. In this way, Trace can create a computational graph (DAG) of the workflow of computing the objective. We visualize the DAG below."
    ]
@@ -216,7 +216,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Version 1 Trace Implementation based on Optimizer\n",
+    "## Version 1 Trace Implementation based on Optimizer\n",
     "\n",
     "The first way is to implement the back-propagation algorithm as part of the optimizer in Trace. By default, optimzers in Trace receive the propagated Trace graph at the parameter nodes. Trace graph is a generalization of gradient. Here we show how we can implement back-propagation on the Trace graph to recover the propagated gradient and use it for gradient descent. We can see the loss sequence here matches what we had above implemented by PyTorch.\n"
    ]
@@ -286,7 +286,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Version 2 Trace Implementation based on Propagator + Optimizer\n",
+    "## Version 2 Trace Implementation based on Propagator + Optimizer\n",
     "\n",
     "Another way is to override the what's propagated in the `backward` call of Trace. Trace has a generic backward routine performed on the computational graph that can support designing new end-to-end optimization algorithms. While by default Trace propagates Trace graphes in `backward` for generality, for the differentiable problems here we can override the behavior and let it directly propagate gradients. In this way, the optimizer would receive directly the propagted gradient instead of Trace graphs.\n"
    ]

docs/tutorials/minibatch.ipynb

@@ -267,7 +267,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Batching Non-Commutative Feedbacks\n",
+    "## Batching Non-Commutative Feedbacks\n",
     "\n",
     "In the earlier numerical example, the loss function was commutative so that we can do  `batch_loss += loss(each_input)`. What if the feedbacks received are not commutative? This can happen often with non-numeric (e.g. text) feedbacks. Here we will see a simple design pattern for using `trace` and `OptoPrime` for batch optimization in such cases."
    ]

opto/optimizers/buffers.py

@@ -7,7 +7,7 @@ def __init__(self, size: int):
     def add(self, item):
         if self.size > 0:
             self.buffer.append(item)
-            self.buffer = self.buffer[-self.size :]
+            self.buffer = self.buffer[-self.size:]
 
     def __iter__(self):
         return iter(self.buffer)

opto/optimizers/opro.py

@@ -19,7 +19,8 @@ class OPRO(OptoPrime):
 
     output_format_prompt = dedent(
         """
-        Output_format: Your output should be in the following json format, satisfying the json syntax:
+        Output_format: Your output should be in the following json format, satisfying
+        the json syntax:
 
         {{
         "suggestion": {{
@@ -28,7 +29,10 @@ class OPRO(OptoPrime):
         }}
         }}
 
-        When suggestion variables, write down the suggested values in "suggestion".  When <type> of a variable is (code), you should write the new definition in the format of python code without syntax errors, and you should not change the function name or the function signature.
+        When suggestion variables, write down the suggested values in "suggestion".
+        When <type> of a variable is (code), you should write the new definition in the
+        format of python code without syntax errors, and you should not change the
+        function name or the function signature.
 
         If no changes or answer are needed, just output TERMINATE.
         """
@@ -57,5 +61,7 @@ def construct_prompt(self, summary, mask=None, *args, **kwargs):
             )
         examples = "\n".join(examples)
 
-        user_prompt = self.user_prompt_template.format(examples=examples, instruction=self.objective)
-        return self.output_format_prompt, user_prompt
+        user_prompt = self.user_prompt_template.format(
+            examples=examples, instruction=self.objective
+        )
+        return self.output_format_prompt, user_prompt

opto/optimizers/optimizer.py

@@ -29,9 +29,15 @@ def propagator(self):
 
 
 class Optimizer(AbstractOptimizer):
-    """ Optimizer based on Trace graph. """
-
-    def __init__(self, parameters: List[ParameterNode], *args, propagator: Propagator = None, **kwargs):
+    """Optimizer based on Trace graph."""
+
+    def __init__(
+        self,
+        parameters: List[ParameterNode],
+        *args,
+        propagator: Propagator = None,
+        **kwargs
+    ):
         super().__init__(parameters)
         propagator = propagator if propagator is not None else self.default_propagator()
         assert isinstance(propagator, Propagator)
@@ -43,7 +49,7 @@ def propagator(self):
 
     @property
     def trace_graph(self):
-        """ Aggregate the graphs of all the parameters. """
+        """Aggregate the graphs of all the parameters."""
         return sum_feedback(self.parameters)
 
     def step(self, *args, **kwargs):