aparulsarma/openenv-pathway-analysis-env

1#!/usr/bin/env python32# Copyright (c) Meta Platforms, Inc. and affiliates.3# All rights reserved.4#5# OpenAI tool-calling agent for pathway_analysis_env (in-process orchestrator).6#7# Usage:8#   export OPENAI_API_KEY=...9#   PYTHONPATH=src:envs uv run python examples/pathway_agent_loop.py \10#       --case toy_case_001.json11 12from __future__ import annotations13 14import argparse15import asyncio16import json17import os18import sys19 20from pathway_analysis_env.agent_openai_tools import (21    OPENAI_TOOLS,22    observation_to_tool_result_content,23    tool_call_to_pathway_action,24)25from pathway_analysis_env.models import PathwayAction26from pathway_analysis_env.server.pathway_environment import PathwayEnvironment27 28 29SYSTEM_PROMPT = """You are a computational biologist agent operating a pathway analysis environment.30 31Required workflow (eval mode):321. understand_experiment_design and/or inspect_dataset — learn groups and sample layout.332. run_differential_expression — set reference (baseline) vs alternate (treatment) conditions.343. run_pathway_enrichment — ORA on DE genes (do not pass a custom gene_list).354. Optionally compare_pathways between two top pathway names.365. submit_answer — one pathway hypothesis string supported by ORA.37 38Rules:39- Never guess without running DE and ORA first.40- Use condition names exactly as returned in available_conditions.41- For submit_answer, name a specific pathway (e.g. from top_pathways), not a long essay.42"""43 44 45async def run_episode(46    case_file: str,47    model: str,48    max_turns: int,49    *,50    strict: bool,51) -> dict:52    try:53        from openai import AsyncOpenAI54    except ImportError as exc:55        raise SystemExit("Install openai: uv add openai") from exc56 57    if not os.environ.get("OPENAI_API_KEY"):58        print("Warning: OPENAI_API_KEY not set", file=sys.stderr)59 60    client = AsyncOpenAI()61    env = PathwayEnvironment(case_file=case_file)62    obs = env.reset(orchestrator_mode=True, strict=strict)63    messages = [64        {"role": "system", "content": SYSTEM_PROMPT},65        {66            "role": "user",67            "content": (68                f"Episode started for case {case_file}. "69                f"Conditions: {obs.available_conditions}. "70                f"{obs.message}"71            ),72        },73    ]74 75    for turn in range(max_turns):76        response = await client.chat.completions.create(77            model=model,78            messages=messages,79            tools=OPENAI_TOOLS,80            tool_choice="auto",81        )82        msg = response.choices[0].message83        if not msg.tool_calls:84            messages.append({"role": "assistant", "content": msg.content or ""})85            if env.state.is_done:86                break87            continue88 89        messages.append(msg.model_dump())90        for tc in msg.tool_calls:91            action = tool_call_to_pathway_action(92                name=tc.function.name,93                arguments_json=tc.function.arguments,94            )95            step_obs = env.step(action)96            messages.append(97                {98                    "role": "tool",99                    "tool_call_id": tc.id,100                    "content": observation_to_tool_result_content(step_obs),101                }102            )103            if step_obs.done:104                return {105                    "turns": turn + 1,106                    "done": True,107                    "episode_outcome": env.episode_outcome,108                    "last_message": step_obs.message,109                    "steps": env.state.step_count,110                }111 112    return {113        "turns": max_turns,114        "done": env.state.is_done,115        "episode_outcome": env.episode_outcome,116        "steps": env.state.step_count,117    }118 119 120def main() -> None:121    parser = argparse.ArgumentParser(description="LLM agent on pathway_analysis_env")122    parser.add_argument("--case", default="toy_case_001.json")123    parser.add_argument("--model", default="gpt-4o-mini")124    parser.add_argument("--max-turns", type=int, default=24)125    parser.add_argument("--strict", action="store_true")126    args = parser.parse_args()127    result = asyncio.run(128        run_episode(args.case, args.model, args.max_turns, strict=args.strict)129    )130    print(json.dumps(result, indent=2))131    outcome = result.get("episode_outcome") or {}132    if outcome.get("correct"):133        sys.exit(0)134    sys.exit(1 if result.get("done") else 2)135 136 137if __name__ == "__main__":138    main()139