Merge pull request #40 from microsoft/0.1.3.5

Update to 0.1.3.5. Change default to use LiteLLM backend. (Autogen backend is still supported). Add custom openai backend support too.

Author: chinganc

README.md

@@ -31,14 +31,13 @@ Or for development, clone the repo and run the following.
 
     pip install -e .
 
-The library requires Python >= 3.9. The installation script will git
-clone [AutoGen](https://github.com/microsoft/autogen). You may require [Git Large File Storage](https://git-lfs.com/) if
-git is unable to clone the repository otherwise.
+The library requires Python >= 3.9. By default (starting with v0.1.3.5), we use [LiteLLM](https://github.com/BerriAI/litellm) as the backend of LLMs. For backward compatibility, we provide backend-support with [AutoGen](https://github.com/microsoft/autogen); when installing, users can add `[autogen]` tag to install a compatible AutoGen version (e.g., `pip install trace-opt[autogen]`). You may require [Git Large File Storage](https://git-lfs.com/) if
+git is unable to clone the repository.
 
 ## Updates
 - **2025.2.7** Trace was featured in the [G-Research NeurIPS highlight](https://www.gresearch.com/news/neurips-paper-reviews-2024-8/) by the Science Director Hugh Salimbeni.
 - **2024.12.10** Trace was demoed in person at NeurIPS 2024 Expo.
-- **2024.11.05** Ching-An Cheng gave a talk at UW Robotics Colloquium on Trace: [video](https://www.youtube.com/watch?v=T2g1Vo3u_9g). 
+- **2024.11.05** Ching-An Cheng gave a talk at UW Robotics Colloquium on Trace: [video](https://www.youtube.com/watch?v=T2g1Vo3u_9g).
 - **2024.10.21**    New [paper](https://arxiv.org/abs/2410.15625) by Nvidia, Stanford, Visa, & Intel applies Trace to
   optimize for mapper code of parallel programming (for scientific computing and matrix multiplication). Trace (OptoPrime) learns code achieving 1.3X speed up under 10
   minutes, compared to the code optimized by a system engineer expert.
@@ -103,7 +102,6 @@ Now, after declaring what is trainable and what isn't, and use `node` and `bundl
 can use the optimizer to optimize the computation graph.
 
 ```python
-import autogen
 from opto.optimizers import OptoPrime
 
 
@@ -121,8 +119,7 @@ test_input = [1, 2, 3, 4]
 
 epoch = 2
 
-optimizer = OptoPrime(strange_sort_list.parameters(),
-                      config_list=autogen.config_list_from_json("OAI_CONFIG_LIST"))
+optimizer = OptoPrime(strange_sort_list.parameters())
 
 for i in range(epoch):
     print(f"Training Epoch {i}")
@@ -215,16 +212,16 @@ def train():
 agent = train()
 ```
 
-Defining and training an agent through Trace will give you more flexibility and control over what the agent learns. 
+Defining and training an agent through Trace will give you more flexibility and control over what the agent learns.
 
 ## Tutorials
 
 | **Level** | **Tutorial**                                                                              | **Run in Colab**                                                                                                                                                                                       | **Description**                                                                                                                                                                       |
-| --- |-------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| 
+| --- |-------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | Beginner | [Getting Started](https://microsoft.github.io/Trace/quickstart/quick_start.html)          | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/quickstart/quick_start.ipynb)       | Introduces basic primitives like `node` and `bundle`. Showcases a code optimization pipeline.                                                                                         |
-| Beginner | [Adaptive AI Agent](https://microsoft.github.io/Trace/quickstart/quick_start_2.html)      | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/quickstart/quick_start_2.ipynb)      | Introduce primitive `model` that allows anyone to build self-improving agents that react to environment feedback. Shows how an LLM agent learns to place a shot in a Battleship game. 
-| Intermediate | [Multi-Agent Collaboration](https://microsoft.github.io/Trace/quickstart/virtualhome.html) | N/A                                                                                                                                                                                                    | Demonstrates how Trace can be used for multi-agent collaboration environment in Virtualhome.                                                                                          
-| Intermediate | [NLP Prompt Optimization](https://microsoft.github.io/Trace/examples/nlp/bigbench_hard.html) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/examples/nlp/bigbench_hard.ipynb) | Shows how Trace can optimizes prompt and code together jointly for BigBench-Hard 23 tasks.                                                                                            
+| Beginner | [Adaptive AI Agent](https://microsoft.github.io/Trace/quickstart/quick_start_2.html)      | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/quickstart/quick_start_2.ipynb)      | Introduce primitive `model` that allows anyone to build self-improving agents that react to environment feedback. Shows how an LLM agent learns to place a shot in a Battleship game.
+| Intermediate | [Multi-Agent Collaboration](https://microsoft.github.io/Trace/quickstart/virtualhome.html) | N/A                                                                                                                                                                                                    | Demonstrates how Trace can be used for multi-agent collaboration environment in Virtualhome.
+| Intermediate | [NLP Prompt Optimization](https://microsoft.github.io/Trace/examples/nlp/bigbench_hard.html) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/examples/nlp/bigbench_hard.ipynb) | Shows how Trace can optimizes prompt and code together jointly for BigBench-Hard 23 tasks.
 | Advanced | [Robotic Arm Control](https://microsoft.github.io/Trace/examples/robotics/metaworld.html) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/microsoft/Trace/blob/website/docs/examples/robotics/metaworld.ipynb)                                     | Trace can optimize code to control a robotic arm after observing a full trajectory of interactions.                                                                                   |
 
 
@@ -276,8 +273,9 @@ with TraceGraph coming soon).
 
 ## LLM API Setup
 
-Currently we rely on AutoGen for LLM caching and API-Key management. 
-AutoGen relies on `OAI_CONFIG_LIST`, which is a file you put in your working directory. It has the format of:
+Currently we rely on LiteLLM or AutoGen for LLM caching and API-Key management.
+By default, LiteLLM is used. Please the documentation there to set the right environment variables for keys and end-point urls.
+On the other hand, AutoGen relies on `OAI_CONFIG_LIST`, which is a file you put in your working directory. It has the format of:
 
 ```json lines
 [

docs/examples/game/battleship.ipynb

@@ -923,7 +923,7 @@
     }
    ],
    "source": [
-    "optimizer = OptoPrime(policy.parameters(), memory_size=0, config_list=config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime(policy.parameters(), memory_size=0)\n",
     "\n",
     "\n",
     "feedback = \"\"\n",

docs/quickstart/quick_start.ipynb

@@ -381,8 +381,7 @@
     "\n",
     "epoch = 2\n",
     "\n",
-    "optimizer = OptoPrime(strange_sort_list.parameters(),\n",
-    "                      config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime(strange_sort_list.parameters())\n",
     "\n",
     "for i in range(epoch):\n",
     "    print(f\"Training Epoch {i}\")\n",
@@ -933,10 +932,8 @@
     "epoch = 2\n",
     "\n",
     "code_optimizer = OptoPrime(strange_sort_list.parameters(),\n",
-    "                          config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"),\n",
     "                          ignore_extraction_error=True)\n",
     "verifier_optimizer = OptoPrime(verifier.parameters(),\n",
-    "                              config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"),\n",
     "                              ignore_extraction_error=True)\n",
     "\n",
     "for i in range(epoch):\n",

docs/quickstart/quick_start_2.ipynb

@@ -127,14 +127,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys\n",
     "import os\n",
+    "import urllib.request\n",
+    "import importlib.util\n",
     "\n",
-    "# Get the absolute path of the examples folder\n",
-    "examples_path = os.path.abspath(os.path.join('..', '..', 'examples'))\n",
+    "# Define the raw URL for downloading\n",
+    "raw_url = \"https://raw.githubusercontent.com/microsoft/Trace/main/examples/battleship.py\"\n",
     "\n",
-    "# Add the examples folder to the Python path\n",
-    "sys.path.append(examples_path)"
+    "# Define the local file path\n",
+    "local_file = \"battleship.py\"\n",
+    "\n",
+    "# Download the file\n",
+    "urllib.request.urlretrieve(raw_url, local_file)\n",
+    "print(f\"Downloaded {local_file}\")\n",
+    "\n",
+    "# Load the module dynamically\n",
+    "spec = importlib.util.spec_from_file_location(\"battleship\", local_file)\n",
+    "battleship = importlib.util.module_from_spec(spec)\n",
+    "spec.loader.exec_module(battleship)\n",
+    "\n",
+    "print(\"battleship module loaded successfully!\")"
    ]
   },
   {
@@ -2189,7 +2201,7 @@
     "\n",
     "agent = Agent()\n",
     "obs = node(board.get_shots(), trainable=False)\n",
-    "optimizer = OptoPrime(agent.parameters(), config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime(agent.parameters())\n",
     "\n",
     "feedback, terminal, cum_reward = \"\", False, 0\n",
     "\n",

docs/tutorials/optimization_tutorial.ipynb

@@ -177,7 +177,7 @@
     "\n",
     "# One-step optimization example\n",
     "x = node(-1.0, trainable=True)\n",
-    "optimizer = OptoPrime([x], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([x])\n",
     "output = foobar(x)\n",
     "feedback = user(output.data)\n",
     "optimizer.zero_feedback()\n",
@@ -415,7 +415,7 @@
     "# A small example of how to use the optimizer in a loop\n",
     "GRAPH.clear()\n",
     "x = node(-1.0, trainable=True)\n",
-    "optimizer = OptoPrime([x], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([x])\n",
     "\n",
     "history = [x.data]\n",
     "feedback = \"\"\n",
@@ -466,7 +466,7 @@
     "# A small example of how to include constraints on parameters\n",
     "GRAPH.clear()\n",
     "x = node(-1.0, trainable=True, constraint=\"The value should be greater than 2.0\")\n",
-    "optimizer = OptoPrime([x], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([x])\n",
     "\n",
     "history = [x.data]\n",
     "feedback = \"\"\n",
@@ -600,7 +600,7 @@
    "source": [
     "GRAPH.clear()\n",
     "x = node(\"negative point one\", trainable=True)\n",
-    "optimizer = OptoPrime([x], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([x])\n",
     "\n",
     "history = [x.data]\n",
     "feedback = \"\"\n",
@@ -675,7 +675,7 @@
     "\n",
     "\n",
     "x = node(-1, trainable=False)\n",
-    "optimizer = OptoPrime([my_fun.parameter], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([my_fun.parameter])\n",
     "\n",
     "feedback = \"\"\n",
     "while feedback != \"Success.\":\n",
@@ -917,7 +917,7 @@
    ],
    "source": [
     "x = node(0.005, trainable=True)\n",
-    "optimizer = OptoPrime([x], config_list=autogen.config_list_from_json(\"OAI_CONFIG_LIST\"))\n",
+    "optimizer = OptoPrime([x])\n",
     "\n",
     "history = []\n",
     "bestScore = None\n",

opto/optimizers/optoprime.py

@@ -1,7 +1,6 @@
 from typing import Any, List, Dict, Union, Tuple
 from dataclasses import dataclass, asdict
 from textwrap import dedent, indent
-import autogen
 import warnings
 import json
 import re
@@ -11,7 +10,7 @@
 from opto.trace.propagators.propagators import Propagator
 from opto.optimizers.optimizer import Optimizer
 from opto.optimizers.buffers import FIFOBuffer
-from opto.utils.llm import AutoGenLLM
+from opto.utils.llm import AbstractModel, LLM
 
 
 def get_fun_name(node: MessageNode):
@@ -250,7 +249,7 @@ class OptoPrime(Optimizer):
     def __init__(
         self,
         parameters: List[ParameterNode],
-        llm: AutoGenLLM = None,
+        llm: AbstractModel = None,
         *args,
         propagator: Propagator = None,
         objective: Union[None, str] = None,
@@ -260,12 +259,11 @@ def __init__(
         max_tokens=4096,
         log=True,
         prompt_symbols=None,
-        filter_dict: Dict = None,  # autogen filter_dict
         **kwargs,
     ):
         super().__init__(parameters, *args, propagator=propagator, **kwargs)
         self.ignore_extraction_error = ignore_extraction_error
-        self.llm = llm or AutoGenLLM()
+        self.llm = llm or LLM()
         self.objective = objective or self.default_objective
         self.example_problem = ProblemInstance.problem_template.format(
             instruction=self.default_objective,

opto/optimizers/textgrad.py

@@ -1,12 +1,11 @@
 import json
 from dataclasses import dataclass
-import autogen
 from typing import Any, List, Dict, Union, Tuple, Optional
 from opto.optimizers.optimizer import Optimizer
 from opto.trace.nodes import ParameterNode, Node, MessageNode
 from opto.trace.propagators import TraceGraph, GraphPropagator, Propagator
 from opto.trace.utils import escape_json_nested_quotes, remove_non_ascii
-from opto.utils.llm import AutoGenLLM
+from opto.utils.llm import LLM, AbstractModel
 
 from copy import copy
 import re
@@ -309,7 +308,7 @@ class TextGrad(Optimizer):
     def __init__(
         self,
         parameters: List[ParameterNode],
-        llm: AutoGenLLM = None,
+        llm: AbstractModel = None,
         *args,
         propagator: Propagator = None,
         objective: Union[None, str] = None,
@@ -318,7 +317,7 @@ def __init__(
         **kwargs,
     ):
         super().__init__(parameters, *args, **kwargs)
-        self.llm = llm or AutoGenLLM()
+        self.llm = llm or LLM()
         self.print_limit = 100
         self.max_tokens = max_tokens
         self.new_variable_tags = ["<IMPROVED_VARIABLE>", "</IMPROVED_VARIABLE>"]

opto/trace/nodes.py

@@ -2181,7 +2181,7 @@ def __init__(
     def create_feedback(self, style="simple"):
         assert style in ("simple", "full")
         feedback = self._data
-        if style in ("line", "full"):
+        if style == "full":
             if type(self.info) is dict and self.info.get("error_comment") is not None:
                 feedback = self.info["error_comment"]
         return feedback

opto/trace/operators.py

@@ -594,8 +594,7 @@ def call_llm(system_prompt, *user_prompts, **kwargs):
     messages = [{"role": "system", "content": system_prompt}]
     for user_prompt in user_prompts:
         messages.append({"role": "user", "content": user_prompt})
-    from opto.utils.llm import AutoGenLLM
-
-    llm = AutoGenLLM()
+    from opto.utils.llm import LLM
+    llm = LLM()
     response = llm(messages=messages, **kwargs)
     return response.choices[0].message.content