Jayr/appworld

.gitignore

@@ -207,4 +207,5 @@ __marimo__/
 
 # Experiments
 runs/
-gepa_terminus/
+gepa_terminus/
+data/

src/gepa/adapters/app_world_adapter/app_world_adapter.py

@@ -0,0 +1,107 @@
+import json
+import os
+import io
+import contextlib
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Any
+
+from pydantic import BaseModel
+from appworld_experiments.code.gepa.gepa_agent import GEPAAgent
+
+from gepa import EvaluationBatch, GEPAAdapter
+from appworld import AppWorld
+from appworld.evaluator import TestTracker
+
+class AppWorldTask(BaseModel):
+    task_id: str
+
+class AppWorldAdapter(GEPAAdapter):
+
+    def __init__(
+        self,
+        agent_config: Dict[str, Any],
+        experiment_name: str,
+    ):
+        self.agent = GEPAAgent.from_dict(agent_config)
+        self.experiment_name = experiment_name
+
+    def evaluate(
+        self,
+        batch: list[AppWorldTask],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch:
+        outputs = []
+        scores = []
+        trajectories = []
+
+        example_run_id = "_temp_gepa_run" + "_" + datetime.now().strftime("%Y%m%d%H%M%S")
+        instruct_prompt = candidate["instruction_prompt"]
+        num_tasks = len(batch)
+        self.agent.logger.initialize(
+            experiment_name=self.experiment_name + example_run_id,
+            num_tasks=num_tasks,
+            num_processes=1,
+            process_index=0,
+        )
+        self.agent.gepa_prompt_replace = instruct_prompt
+
+        with contextlib.redirect_stdout(io.StringIO()):
+            for example in batch:
+                task_id = example.task_id
+                test_tracker = self.agent.solve_task(task_id, self.experiment_name + example_run_id)
+                try:
+                    success = test_tracker.success
+                    score = int(success)
+                    #score = len(test_tracker.passes) / test_tracker._num_tests
+                    failed_reason_list = []
+                    for failure in test_tracker.failures:
+                        failed_reason_list.append(json.dumps(failure, indent=2))
+                    failed_reason = ','.join(failed_reason_list)
+                except Exception as e:
+                    #TODO: need to handle case for failed code execution
+                    success = False
+                    score = 0
+                    failed_reason = "\n\n".join(test_tracker)
+                outputs.append(
+                    f"App World outputs are omitted. Please see directory for detailed logging."
+                )
+                scores.append(score)
+                trajectories.append(
+                    {
+                        "messages": self.agent.messages,
+                        "instruction_prompt": instruct_prompt,
+                        "failed_reason": str(failed_reason),
+                        "success": success,
+                    }
+                )
+        return EvaluationBatch(
+            outputs=outputs,
+            scores=scores,
+            trajectories=trajectories,
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch,
+        components_to_update: list[str],
+    ):
+        reflective_dataset = {"instruction_prompt": []}
+        for score, trajectory in zip(eval_batch.scores, eval_batch.trajectories, strict=False):
+            if trajectory["success"]:
+                feedback = "Successfully solved the task!"
+            else:
+                feedback = (
+                    f"Failed to solve the task. Reason: {trajectory['failed_reason']}"
+                )
+            reflective_dataset["instruction_prompt"].append(
+                {
+                    "Message History": trajectory["messages"],
+                    "Instruction Prompt": candidate["instruction_prompt"],
+                    "Feedback": feedback,
+                }
+            )
+        return reflective_dataset

src/gepa/adapters/bird_adapter/bird_adapter.py

@@ -0,0 +1,90 @@
+import json
+import os
+import subprocess
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Any
+
+from pydantic import BaseModel
+
+from gepa import EvaluationBatch, GEPAAdapter
+import io
+import contextlib
+import sys
+from llm.src.gpt_request import collect_response_from_gpt
+from llm.src.evaluation import execute_model
+
+class BirdTask(BaseModel):
+    question: str
+    db_path: str
+    knowledge: str
+    ground_truth: str
+
+class BirdAdapter(GEPAAdapter):
+
+    def __init__(
+        self,
+        model_name: str,
+    ):
+        self.model_name = model_name
+
+    def evaluate(
+        self,
+        batch: list[BirdTask],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch:
+        outputs = []
+        scores = []
+        trajectories = []
+
+        instruct_prompt = candidate["instruction_prompt"]
+
+        for example in batch:
+            responses, message_list = collect_response_from_gpt(db_path_list=[example.db_path], question_list=[example.question], api_key="", engine=self.model_name, knowledge_list=[example.knowledge], system_message=instruct_prompt)
+            predict_sql, predict_db_name = responses[0].split('\t----- bird -----\t')
+            predict_db_path = ...
+            ground_truth_sql = example.ground_truth.strip()
+            result = execute_model(predict_sql, ground_truth_sql, predict_db_path, 0, 30.0)
+            score = result['res']
+            if score == 1:
+                success = True
+            else:
+                success = False
+            outputs.append("No logging here...")
+            scores.append(score)
+            trajectories.append(
+                {
+                    "messages": message_list[0],
+                    "instruction_prompt": instruct_prompt,
+                    "success": success,
+                }
+            )
+        return EvaluationBatch(
+            outputs=outputs,
+            scores=scores,
+            trajectories=trajectories,
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch,
+        components_to_update: list[str],
+    ):
+        reflective_dataset = {"instruction_prompt": []}
+        for score, trajectory in zip(eval_batch.scores, eval_batch.trajectories, strict=False):
+            if trajectory["success"]:
+                feedback = "Successfully solved the task!"
+            else:
+                feedback = (
+                    f"Failed to solve the task."
+                )
+            reflective_dataset["instruction_prompt"].append(
+                {
+                    "Message History": trajectory["messages"],
+                    "Instruction Prompt": candidate["instruction_prompt"],
+                    "Feedback": feedback,
+                }
+            )
+        return reflective_dataset

src/gepa/examples/app-world/test_appworld.py

@@ -0,0 +1,42 @@
+import os
+import argparse
+from appworld.common.path_store import path_store
+from gepa.core.state import GEPAState
+from gepa.core.result import GEPAResult
+from appworld.common.utils import jsonnet_load, read_file
+from appworld.task import Task, load_task_ids
+import litellm
+from gepa.adapters.app_world_adapter.app_world_adapter import (
+    AppWorldTask,
+    AppWorldAdapter,
+)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment_name", type=str, required=True)
+    parser.add_argument("--appworld_bin_dir_path", type=str, required=True)
+    parser.add_argument("--test_type", type=str, choices=["normal", "challenge"], required=True, help="Choose either normal or challenge mode.")
+    args = parser.parse_args()
+    experiment_name = args.experiment_name
+    gepa_state_path = args.appworld_bin_dir_path
+    test_type = args.test_type
+    experiment_file_path = os.path.join(path_store.experiment_configs, experiment_name + ".jsonnet")
+    experiment_config = jsonnet_load(
+        experiment_file_path,
+        APPWORLD_EXPERIMENT_PROMPTS_PATH=path_store.experiment_prompts,
+        APPWORLD_EXPERIMENT_CONFIGS_PATH=path_store.experiment_configs,
+        APPWORLD_EXPERIMENT_CODE_PATH=path_store.experiment_code,
+    )
+    runner_config = experiment_config.pop("config")
+    agent_config = runner_config.pop("agent")
+
+    gepa_state = GEPAState.load(gepa_state_path)
+    result = GEPAResult.from_state(gepa_state)
+
+    test_task_ids = load_task_ids(f'test_{test_type}')
+    for task_id in test_task_ids:
+        Task.load(task_id=task_id)
+    print(f"Length of original test {test_type} dataset: {len(test_task_ids)}")
+    testset = [AppWorldTask(task_id=task_id) for task_id in test_task_ids[:]]
+
+    adapter = AppWorldAdapter(agent_config, experiment_name=experiment_name)
+    testset_results_before_opt = adapter.evaluate(testset, {"instruction_prompt": result.best_candidate["instruction_prompt"]}, capture_traces=True)

src/gepa/examples/app-world/train_appworld.py

@@ -0,0 +1,139 @@
+import os
+import argparse
+from appworld.common.path_store import path_store
+from gepa import optimize
+import litellm
+from gepa.adapters.app_world_adapter.app_world_adapter import (
+    AppWorldTask,
+    AppWorldAdapter,
+)
+from appworld.common.utils import jsonnet_load, read_file
+from appworld.task import Task, load_task_ids
+import json
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment_name", type=str, default=None)
+    parser.add_argument("--save-price-path", type=str, default='appworld_gepa_prompt_modification_pricing.jsonl')
+    args = parser.parse_args()
+    experiment_name = args.experiment_name
+    price_path_file = args.save_price_path
+    experiment_file_path = os.path.join(path_store.experiment_configs, experiment_name + ".jsonnet")
+    experiment_config = jsonnet_load(
+        experiment_file_path,
+        APPWORLD_EXPERIMENT_PROMPTS_PATH=path_store.experiment_prompts,
+        APPWORLD_EXPERIMENT_CONFIGS_PATH=path_store.experiment_configs,
+        APPWORLD_EXPERIMENT_CODE_PATH=path_store.experiment_code,
+    )
+    runner_config = experiment_config.pop("config")
+    agent_config = runner_config.pop("agent")
+    print(f"Running with agent config: {agent_config}")
+
+    initial_prompt_from_cleaned = """I am your supervisor and you are a super intelligent AI Assistant whose job is to achieve my day-to-day tasks completely autonomously.
+
+To do this, you will need to interact with app/s (e.g., spotify, venmo etc) using their associated APIs on my behalf. For this you will undertake a *multi-step conversation* using a python REPL environment. That is, you will write the python code and the environment will execute it and show you the result, based on which, you will write python code for the next step and so on, until you've achieved the goal. This environment will let you interact with app/s using their associated APIs on my behalf.
+
+Here are three key APIs that you need to know to get more information
+
+# To get a list of apps that are available to you.
+
+```python
+print(apis.api_docs.show_app_descriptions())
+```
+
+# To get the list of apis under any app listed above, e.g. spotify
+
+```python
+print(apis.api_docs.show_api_descriptions(app_name='spotify'))
+```
+
+# To get the specification of a particular api, e.g. spotify app's login api
+
+```python
+print(apis.api_docs.show_api_doc(app_name='spotify', api_name='login'))
+```
+
+Each code execution will produce an output that you can use in subsequent calls. Using these APIs, you can now generate code, that I will execute, to solve the task. 
+
+You are also provided with a curated cheatsheet of strategies, API-specific information, common mistakes, and proven solutions to help you solve the task effectively.
+
+**Cheatsheet**: - Read the **Cheatsheet** first, then execute the task by explicitly leveraging each relevant section:
+### CHEATSHEET BEGIN
+## STRATEGIES AND HARD RULES
+[shr-00001] Make sure to end code blocks with ``` followed by a newline(\\n).
+[shr-00005] Always look at API specifications (using apis.api_docs.show_api_doc) before calling an API.
+[shr-00006] Write small chunks of code and only one chunk of code in every step. Make sure everything is working correctly before making any irreversible change.
+
+## APIs TO USE FOR SPECIFIC INFORMATION
+[api-00004] You can use the "supervisor" app to get information about my accounts and use the "phone" app to get information about friends and family.
+
+## USEFUL CODE SNIPPETS AND TEMPLATES
+
+## COMMON MISTAKES AND CORRECT STRATEGIES
+
+## PROBLEM-SOLVING HEURISTICS AND WORKFLOWS
+[psw-00002] Remember you can use the variables in your code in subsequent code blocks.
+[psw-00007] Many APIs return items in "pages". Make sure to run through all the pages by looping over `page_index`.
+
+## VERIFICATION CHECKLIST
+
+## TROUBLESHOOTING AND PITFALLS:
+
+## OTHERS
+[misc-00003] Remember that the email addresses, access tokens and variables (e.g. spotify_password) in the example above are not valid anymore.
+[misc-00008] Once you have completed the task, make sure to call apis.supervisor.complete_task(). If the task asked for some information, return it as the answer argument, i.e. call apis.supervisor.complete_task(answer=<answer>). Many tasks do not require an answer, so in those cases, just call apis.supervisor.complete_task() i.e. do not pass any argument.
+
+### CHEATSHEET END"""
+
+    train_task_ids = load_task_ids('train')
+    val_task_ids = load_task_ids('dev')
+    test_task_ids = load_task_ids('test_normal')
+    for task_id in train_task_ids:
+        Task.load(task_id=task_id)
+    for task_id in val_task_ids:
+        Task.load(task_id=task_id)
+    for task_id in test_task_ids:
+        Task.load(task_id=task_id)
+    print(f"Length of original train dataset: {len(train_task_ids)}")
+    print(f"Length of original val dataset: {len(val_task_ids)}")
+    print(f"Length of original test dataset: {len(test_task_ids)}")
+    trainset = [
+        AppWorldTask(task_id=task_id) for task_id in train_task_ids[:]
+    ]
+    valset = [AppWorldTask(task_id=task_id) for task_id in val_task_ids[:]]
+
+    gepa_prompt_gen_file = price_path_file
+    with open(gepa_prompt_gen_file, "w"):
+        pass
+    reflection_lm_name = "sambanova/DeepSeek-V3.1"
+    print(f"Running with reflector model: {reflection_lm_name}")
+    def call_lm(prompt):
+        response = litellm.completion(
+            model=reflection_lm_name,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        input_tokens = response.usage.prompt_tokens
+        output_tokens = response.usage.completion_tokens
+        with open(gepa_prompt_gen_file, "a") as f:
+            f.write(json.dumps({'input_tokens': input_tokens, 'output_tokens': output_tokens}) + "\n")
+        return response.choices[0].message.content
+    reflection_lm = call_lm
+
+    adapter = AppWorldAdapter(agent_config, experiment_name=experiment_name)
+
+    run_dir = "gepa_app_world_deepseek-v3-1"
+    optimized_results = optimize(
+        seed_candidate={"instruction_prompt": initial_prompt_from_cleaned},
+        trainset=trainset,
+        valset=valset,
+        adapter=adapter,
+        reflection_lm=reflection_lm,
+        max_metric_calls=1434,
+        display_progress_bar=True,
+        run_dir=run_dir,
+    )
+
+    optimized_instruction_prompt = optimized_results.best_candidate["instruction_prompt"]
+
+    with open(f"{run_dir}/best_prompt.txt", 'w') as f:
+        f.write(optimized_instruction_prompt)