From ac4381cae9b493f95dd67d30ef17d1771615983a Mon Sep 17 00:00:00 2001
From: zachary62 <zhuang333@wisc.edu>
Date: Fri, 23 May 2025 01:15:49 -0400
Subject: [PATCH] code generator utils

---
 .../pocketflow-code-generator/doc/design.md   |  6 +-
 .../requirements.txt                          |  3 +
 .../utils/__init__.py                         |  0
 .../utils/call_llm.py                         | 20 +++++++
 .../utils/code_executor.py                    | 56 +++++++++++++++++++
 5 files changed, 82 insertions(+), 3 deletions(-)
 create mode 100644 cookbook/pocketflow-code-generator/requirements.txt
 create mode 100644 cookbook/pocketflow-code-generator/utils/__init__.py
 create mode 100644 cookbook/pocketflow-code-generator/utils/call_llm.py
 create mode 100644 cookbook/pocketflow-code-generator/utils/code_executor.py

diff --git a/cookbook/pocketflow-code-generator/doc/design.md b/cookbook/pocketflow-code-generator/doc/design.md
index 95cfd25..5392c29 100644
--- a/cookbook/pocketflow-code-generator/doc/design.md
+++ b/cookbook/pocketflow-code-generator/doc/design.md
@@ -65,9 +65,9 @@ flowchart TD
    - Used by all LLM-powered nodes for generating tests, code, and analysis
 
 2. **Execute Python Code** (`utils/code_executor.py`)
-   - *Input*: function_code (str), test_case (dict)
-   - *Output*: test_result (dict with passed, failed, error details)
-   - Used by RunTests batch node to safely execute generated code against individual test cases
+   - *Input*: function_code (str), input (dict/list/any)
+   - *Output*: output (any), error (str)
+   - Used by RunTests batch node to safely execute generated code against individual input
 
 ## Node Design
 
diff --git a/cookbook/pocketflow-code-generator/requirements.txt b/cookbook/pocketflow-code-generator/requirements.txt
new file mode 100644
index 0000000..d9c598a
--- /dev/null
+++ b/cookbook/pocketflow-code-generator/requirements.txt
@@ -0,0 +1,3 @@
+anthropic
+pocketflow
+pyyaml 
\ No newline at end of file
diff --git a/cookbook/pocketflow-code-generator/utils/__init__.py b/cookbook/pocketflow-code-generator/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/cookbook/pocketflow-code-generator/utils/call_llm.py b/cookbook/pocketflow-code-generator/utils/call_llm.py
new file mode 100644
index 0000000..df57be2
--- /dev/null
+++ b/cookbook/pocketflow-code-generator/utils/call_llm.py
@@ -0,0 +1,20 @@
+from anthropic import Anthropic
+import os
+
+def call_llm(prompt):
+    client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", "your-api-key"))
+    response = client.messages.create(
+        model="claude-sonnet-4-20250514",
+        max_tokens=6000,
+        messages=[
+            {"role": "user", "content": prompt}
+        ]
+    )
+    return response.content[0].text
+
+if __name__ == "__main__":
+    print("## Testing call_llm")
+    prompt = "In a few words, what is the meaning of life?"
+    print(f"## Prompt: {prompt}")
+    response = call_llm(prompt)
+    print(f"## Response: {response}")
\ No newline at end of file
diff --git a/cookbook/pocketflow-code-generator/utils/code_executor.py b/cookbook/pocketflow-code-generator/utils/code_executor.py
new file mode 100644
index 0000000..5e376e3
--- /dev/null
+++ b/cookbook/pocketflow-code-generator/utils/code_executor.py
@@ -0,0 +1,56 @@
+import sys
+import io
+import traceback
+from contextlib import redirect_stdout, redirect_stderr
+
+def execute_python(function_code, input):
+    try:
+        namespace = {"__builtins__": __builtins__}
+        stdout_capture = io.StringIO()
+        stderr_capture = io.StringIO()
+        
+        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
+            exec(function_code, namespace)
+            
+            if "run_code" not in namespace:
+                return None, "Function 'run_code' not found"
+            
+            run_code = namespace["run_code"]
+            
+            if isinstance(input, dict):
+                result = run_code(**input)
+            elif isinstance(input, (list, tuple)):
+                result = run_code(*input)
+            else:
+                result = run_code(input)
+            
+            return result, None
+                
+    except Exception as e:
+        return None, f"{type(e).__name__}: {str(e)}"
+
+if __name__ == "__main__":
+    # Test 1: Working function
+    function_code = """
+def run_code(nums, target):
+    for i in range(len(nums)):
+        for j in range(i + 1, len(nums)):
+            if nums[i] + nums[j] == target:
+                return [i, j]
+    return []
+"""
+    
+    input = {"nums": [2, 7, 11, 15], "target": 9}
+    output, error = execute_python(function_code, input)
+    print(f"Output: {output}")
+    print(f"Error: {error}")
+    
+    # Test 2: Function with error
+    broken_function_code = """
+def run_code(nums, target):
+    return nums[100]  # Index error
+"""
+    
+    output2, error2 = execute_python(broken_function_code, input)
+    print(f"Output: {output2}")
+    print(f"Error: {error2}") 
\ No newline at end of file