Merge pull request #221 from xrc10/develop/v0.2.4

add RAP workflow and example

Author: panregedit

examples/rap/README.md

@@ -0,0 +1,83 @@
+# Reasoning via Planning (RAP) Example
+
+This example demonstrates how to use the RAP (Reasoning via Planning) operator for complex reasoning tasks. The RAP operator treats reasoning as a planning problem and uses Monte Carlo Tree Search (MCTS) to explore reasoning paths.
+
+> **Note**: RAP currently only supports the lite version because Conductor cannot handle nested loops (loops inside loops) which are required for the MCTS algorithm.
+
+## Overview
+
+The example implements a RAP workflow with three main components:
+
+1. **Input Interface**
+   - Handles user input for math problems
+   - Provides a simple CLI interface
+
+2. **RAP Workflow**
+   - Uses MCTS to explore reasoning paths
+   - Leverages a language model for reasoning
+   - Finds optimal solution paths
+
+3. **Conclude Task**
+   - Extracts and formats the final answer
+   - Presents results to the user
+
+## Prerequisites
+
+- Access to OpenAI API or compatible endpoint
+
+## Setup
+
+Configure environment variables:
+```bash
+export custom_openai_key="your_openai_api_key"
+export custom_openai_endpoint="your_openai_endpoint"  # Optional: defaults to https://api.openai.com/v1
+```
+
+## Running the Example
+
+Only the lightweight version is supported:
+```bash
+python run_cli_lite.py
+```
+
+The CLI will prompt you to:
+1. Select a task type (currently supports 'math')
+2. Enter your problem/question
+
+## Configuration Files
+
+- `configs/llms/gpt.yml`: LLM settings
+  ```yaml
+  name: OpenaiGPTLLM
+  model_id: gpt-4o-mini
+  api_key: ${env| custom_openai_key}
+  endpoint: ${env| custom_openai_endpoint, https://api.openai.com/v1}
+  temperature: 0.0
+  vision: false
+  ```
+
+- `configs/workers/rap_workflow.yml`: Worker configurations for the RAP pipeline
+
+## Example Usage in Code
+
+```python
+from omagent_core.advanced_components.workflow.rap import RAPWorkflow
+
+# Initialize RAP workflow
+workflow = RAPWorkflow()
+
+# Set input query
+workflow.set_input(query="If 4 friends share 7 pizzas equally, and each pizza has 8 slices, how many slices does each friend get?")
+
+# Get results
+final_answer = workflow.final_answer
+```
+
+## Troubleshooting
+
+Common issues:
+- Invalid API key: Verify your OpenAI API key is set correctly
+- Connection errors: Check your internet connection and API endpoint
+- Import errors: Ensure all dependencies are installed
+
+For more detailed documentation, visit the [OmAgent documentation](https://docs.omagent.ai).

examples/rap/agent/conclude/conclude.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+from omagent_core.engine.worker.base import BaseWorker
+from omagent_core.utils.registry import registry
+import re
+
+@registry.register_worker()
+class Conclude(BaseWorker):
+    """Worker that formats and presents the final reasoning output"""
+
+    def _run(self, final_answer: str, *args, **kwargs):
+        """Format and present the final reasoning results.
+        
+        Args:
+            final_answer: The final answer from RAP
+            
+        Returns:
+            dict: Contains the extracted short answer
+        """
+        
+        # Extract the short answer using regex pattern
+        matches = re.findall(r'the answer is\s+([^.]+)\.', final_answer.lower())
+        short_answer = matches[-1] if matches else final_answer
+        
+        # Send formatted answer through callback
+        self.callback.send_answer(
+            agent_id=self.workflow_instance_id,
+            msg=f"Final conclusion: {short_answer}"
+        )
+
+        return {"short_answer": short_answer} 

examples/rap/agent/input_interface/input_interface.py

@@ -0,0 +1,27 @@
+from omagent_core.engine.worker.base import BaseWorker
+from omagent_core.utils.registry import registry
+
+# Define supported tasks
+SUPPORTED_TASKS = {
+    'math': "Please input your math problem:"
+}
+
+@registry.register_worker()
+class InputInterface(BaseWorker):
+    """Input interface processor that handles user queries."""
+
+    def _run(self, query=None, *args, **kwargs):
+        # Get task selection from user
+        task = input(f'\nWelcome to OmAgent RAP! Please select a task type: {list(SUPPORTED_TASKS.keys())} ')
+        
+        if task not in SUPPORTED_TASKS:
+            raise ValueError(f"Unsupported task type: {task}. Must be one of {list(SUPPORTED_TASKS.keys())}")
+        
+        self.stm(self.workflow_instance_id)['task'] = task 
+
+        # Use query parameter if provided, otherwise get user input
+        data_input = query if query else input(SUPPORTED_TASKS[task])
+        
+        # Store input in STM and return
+        self.stm(self.workflow_instance_id)['data_input'] = data_input
+        return {"query": data_input} 

examples/rap/configs/llms/gpt.yml

@@ -0,0 +1,6 @@
+name: OpenaiGPTLLM
+model_id: gpt-4o-mini
+api_key: ${env| custom_openai_key}
+endpoint: ${env| custom_openai_endpoint, https://api.openai.com/v1}
+temperature: 0.0
+vision: false

examples/rap/configs/workers/rap_workflow.yml

@@ -0,0 +1,9 @@
+- name: InputInterface
+- name: Selection 
+- name: Expansion
+  llm: ${sub|gpt}
+- name: SimulationPreProcess
+- name: SimulationPostProcess
+- name: BackPropagation
+- name: MCTSCompletionCheck
+- name: OutputInterface

examples/rap/run_cli_lite.py

@@ -0,0 +1,64 @@
+# Import required modules and components
+from pathlib import Path
+from omagent_core.utils.container import container
+from omagent_core.engine.workflow.conductor_workflow import ConductorWorkflow
+from omagent_core.engine.workflow.task.simple_task import simple_task
+from omagent_core.advanced_components.workflow.rap import RAPWorkflow
+from omagent_core.clients.devices.cli import DefaultClient
+from omagent_core.utils.registry import registry
+from omagent_core.utils.logger import logging
+from agent.conclude.conclude import Conclude
+from agent.input_interface.input_interface import InputInterface
+
+import os
+os.environ["OMAGENT_MODE"] = "lite"
+
+# Initialize logging
+logging.init_logger("omagent", "omagent", level="INFO")
+
+# Set current working directory path
+CURRENT_PATH = Path(__file__).parents[0]
+
+# Import registered modules
+registry.import_module(CURRENT_PATH.joinpath("agent"))
+
+# Load container configuration and set STM
+container.register_stm("SharedMemSTM")
+
+# Initialize RAP workflow with lite_version=True
+workflow = ConductorWorkflow(name="rap_workflow", lite_version=True)
+
+# Configure workflow tasks:
+# 1. Input interface for user interaction
+client_input_task = simple_task(
+    task_def_name=InputInterface, 
+    task_reference_name="input_interface"
+)
+
+# 2. RAP workflow for reasoning
+rap_workflow = RAPWorkflow()
+rap_workflow.set_input(query=client_input_task.output("query"))
+
+# 3. Conclude task for final answer
+conclude_task = simple_task(
+    task_def_name=Conclude,
+    task_reference_name="task_conclude",
+    inputs={
+        "final_answer": rap_workflow.final_answer,
+    },
+)
+
+# Configure workflow execution flow: Input -> RAP -> Conclude
+workflow >> client_input_task >> rap_workflow >> conclude_task
+
+# Register workflow with overwrite=True for lite version
+workflow.register(overwrite=True)
+
+# Initialize and start CLI client
+config_path = CURRENT_PATH.joinpath("configs")
+cli_client = DefaultClient(
+    interactor=workflow, 
+    config_path=config_path, 
+    workers=[InputInterface(), Conclude()]
+)
+cli_client.start_interactor() 

omagent-core/src/omagent_core/advanced_components/workflow/rap/README.md

@@ -0,0 +1,43 @@
+# Reasoning via Planning (RAP) Operator
+
+RAP is a workflow operator that performs reasoning by treating it as a planning problem with a world model. It uses Monte Carlo Tree Search (MCTS) to explore the reasoning space and find high-reward reasoning paths.
+
+Refer to the example in the `examples/rap` directory to understand how to use this operator.
+
+## Overview
+
+RAP repurposes the language model as both:
+1. A world model - to predict states and simulate outcomes
+2. A reasoning agent - to explore and evaluate reasoning paths
+
+The algorithm uses MCTS to strategically explore the vast reasoning space with a proper balance between exploration and exploitation.
+
+## Inputs, Outputs and Configs
+
+### Inputs:
+| Name | Type | Required | Description |
+| ---- | ---- | -------- | ----------- |
+| query | str | true | The reasoning query/problem to solve |
+
+### Outputs:
+| Name | Type | Description |
+| ---- | ---- | ----------- |
+| final_answer | str | The final answer/solution to the query |
+
+### Configs:
+The config of the RAP operator should be defined in a `rap_workflow.yml` file:
+
+```yml
+- name: InputInterface
+- name: Selection 
+- name: Expansion
+  llm: ${sub|gpt}
+- name: SimulationPreProcess
+- name: SimulationPostProcess
+- name: BackPropagation
+- name: MCTSCompletionCheck
+```
+
+## References
+
+[1] Hao, S., Gu, Y., Ma, H., Hong, J. J., Wang, Z., Wang, D. Z., & Hu, Z. (2023). Reasoning with Language Model is Planning with World Model. arXiv preprint arXiv:2305.14992. 

omagent-core/src/omagent_core/advanced_components/workflow/rap/__init__.py

@@ -0,0 +1,3 @@
+from .workflow import RAPWorkflow
+
+__all__ = ["RAPWorkflow"]

omagent-core/src/omagent_core/advanced_components/workflow/rap/agent/back_propagation/back_propagation.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+from typing import List
+import math
+from omagent_core.models.llms.base import BaseLLMBackend
+from omagent_core.engine.worker.base import BaseWorker
+from omagent_core.models.llms.prompt import PromptTemplate
+from omagent_core.utils.registry import registry
+from pydantic import Field
+
+CURRENT_PATH = Path(__file__).parents[0]
+
+@registry.register_worker()
+class BackPropagation(BaseWorker):
+    """Back propagation worker that updates node values"""
+
+    def _run(self, *args, **kwargs):
+        path = self.stm(self.workflow_instance_id)['selected_path']
+        rewards = []
+        cum_reward = -math.inf
+
+        # Update rewards for each node in path
+        for node in reversed(path):
+            rewards.append(node.reward)
+            cum_reward = self.cum_reward(rewards[::-1])
+            node.cum_rewards.append(cum_reward)
+
+        return {}
+
+    def cum_reward(self, rewards):
+        """Calculate cumulative reward"""
+        return sum(rewards)