Add PyHealth training environments, local vLLM backend, and script ru…

.env.example

@@ -1,7 +1,16 @@
 # Example Environment Variables for Codegen Classroom
 # Copy this file to .env and fill in your values.
 
-# Gemini API Key (Required for student agents)
+# Default LLM backend: "gemini" or "openai_compat" (alias "vllm").
+# Run configs (configs/*.yaml) can set llm.backend per run or per student.
+# CODEGEN_BACKEND=vllm
+
+# OpenAI-compatible server (vLLM) URL, see scripts/serve_vllm.sh.
+# VLLM_BASE_URL=http://localhost:8000/v1
+# Only needed if the server was started with --api-key.
+# VLLM_API_KEY=
+
+# Gemini API Key (only for students using the gemini backend)
 GEMINI_API_KEY=your_api_key_here
 
 # Optional: Default output directory for experiment artifacts

.gitignore

@@ -6,4 +6,6 @@
 .pytest_cache/
 .venv/
 __pycache__/
-data/
+data/
+runs/
+logs/

codegen_classroom/__init__.py

@@ -16,8 +16,10 @@
 
 from codegen_classroom import agent_config
 from codegen_classroom import agents
+from codegen_classroom import backends
 from codegen_classroom import code_execution
 from codegen_classroom import data_utils
+from codegen_classroom import environments
 from codegen_classroom import metrics
 from codegen_classroom import plotting
 from codegen_classroom import reporting
@@ -29,8 +31,10 @@
 __all__ = [
     "__version__",
     "agents",
+    "backends",
     "code_execution",
     "data_utils",
+    "environments",
     "agent_config",
     "metrics",
     "plotting",

codegen_classroom/agent_config.py

@@ -16,13 +16,39 @@
 
 from __future__ import annotations
 
-import importlib
 import os
 
 import dotenv
 from etils import epath
 
 
+def load_env(env_path: str | epath.PathLike | None = None) -> None:
+  """Loads API keys and settings from .env files into os.environ."""
+  if env_path is not None:
+    dotenv.load_dotenv(env_path)
+  dotenv.load_dotenv(epath.Path(__file__).resolve().parent / ".env")
+  dotenv.load_dotenv()
+
+
+# Load .env before reading environment-driven defaults below.
+load_env()
+
+# Default LLM backend: "gemini" (Gemini API) or "openai_compat" (a local vLLM
+# server or any other OpenAI-compatible endpoint).
+DEFAULT_BACKEND = os.environ.get("CODEGEN_BACKEND", "gemini")
+
+# Default base URL of the OpenAI-compatible (vLLM) server.
+DEFAULT_VLLM_BASE_URL = os.environ.get(
+    "VLLM_BASE_URL", "http://localhost:8000/v1"
+)
+
+# Example open-weight models to serve with vLLM. The served model name must
+# match what `vllm serve` was started with (see scripts/serve_vllm.sh).
+DEFAULT_VLLM_MODELS = (
+    "unsloth/Qwen3.8-27B-unsloth-bnb-4bit",
+    "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8",
+)
+
 # Default model used for standard agents and generic tasks.
 # See available models at: https://ai.google.dev/gemini-api/docs/models
 DEFAULT_MODEL = "gemini-3.5-flash"
@@ -69,21 +95,3 @@
 # Combined tuple of default models used when spawning a classroom of students.
 DEFAULT_STUDENT_MODELS = DEFAULT_GEMMA_MODELS + DEFAULT_GEMINI_MODELS
 
-
-def load_env(env_path: str | epath.PathLike | None = None) -> None:
-  """Loads GEMINI_API_KEY from package .env or Colab userdata into os.environ."""
-  if env_path is not None:
-    dotenv.load_dotenv(env_path)
-  dotenv.load_dotenv(epath.Path(__file__).resolve().parent / ".env")
-  dotenv.load_dotenv()
-  if not os.environ.get("GEMINI_API_KEY"):
-    try:
-      userdata = importlib.import_module("google.colab.userdata")
-
-      if key := userdata.get("GEMINI_API_KEY"):
-        os.environ["GEMINI_API_KEY"] = key
-    except Exception:  # pylint: disable=broad-exception-caught
-      pass
-
-
-load_env()

codegen_classroom/agents.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Gemini agents and their configs: Student and Solution Analysis."""
+"""LLM agents and their configs: Student and Solution Analysis."""
 
 from __future__ import annotations
 
@@ -27,8 +27,8 @@
 from typing import Any
 
 from codegen_classroom import agent_config
+from codegen_classroom import backends
 from codegen_classroom import data_utils
-from google import genai
 import numpy as np
 
 logger = logging.getLogger(__name__)
@@ -77,10 +77,17 @@ class GenerationError(Exception):
 
 @dataclasses.dataclass
 class AgentConfig:
-  """Configuration for a Gemini/Gemma agent."""
+  """Configuration for an LLM agent (Gemini API or OpenAI-compatible server)."""
 
   name: str
   api_key: str | None = dataclasses.field(default=None, repr=False)
+  # "gemini" or "openai_compat" (aliases: "openai", "vllm").
+  backend: str = agent_config.DEFAULT_BACKEND
+  # Server URL for OpenAI-compatible backends, e.g. http://localhost:8000/v1.
+  base_url: str | None = None
+  # Extra request fields for OpenAI-compatible backends (e.g. top_p,
+  # chat_template_kwargs).
+  extra_body: Mapping[str, Any] | None = None
   system_instruction: str = ""
   model_name: str = agent_config.DEFAULT_MODEL
   temperature: float = agent_config.DEFAULT_TEMPERATURE
@@ -123,20 +130,18 @@ def flatten(self, prefix: str = "metrics_") -> Mapping[str, Any]:
     return flat_dict
 
 
-class GeminiAgent:
-  """Base class for the Student and Solution Analysis Gemini agents."""
+class LLMAgent:
+  """Base class for the Student and Solution Analysis LLM agents."""
 
   def __init__(self, config: AgentConfig, run_base_dir: str):
     self.config = config
     self.name = self.config.name
     self.run_base_dir = run_base_dir
-    api_key = self.config.api_key or os.environ.get("GEMINI_API_KEY")
-    self.client = genai.Client(
-        api_key=api_key,
-        http_options=genai.types.HttpOptions(
-            timeout=self.config.http_timeout_ms
-        ),
-    )
+    if backends.is_openai_compat(self.config.backend):
+      self.config.base_url = (
+          self.config.base_url or agent_config.DEFAULT_VLLM_BASE_URL
+      )
+    self.backend = backends.make_backend(self.config)
     self.system_instruction = self.config.system_instruction
     self.model_name = self.config.model_name
     self.temperature = self.config.temperature
@@ -170,40 +175,27 @@ def generate(self, prompt: str) -> str:
     Returns:
         The generated response text, or an empty string if generation failed.
     """
-    contents = [
-        genai.types.Content(
-            role=role,
-            parts=[genai.types.Part.from_text(text=text)],
-        )
-        for role, text in self.history
-    ]
-    contents.append(
-        genai.types.Content(
-            role="user",
-            parts=[genai.types.Part.from_text(text=prompt)],
-        )
+    response_text = self.backend.generate(
+        self.history,
+        prompt,
+        system_instruction=self.system_instruction,
+        temperature=self.temperature,
+        max_output_tokens=self.max_output_tokens,
+        json_mode=self.config.response_mime_type == "application/json",
     )
 
-    response = self.client.models.generate_content(
-        model=self.model_name,
-        contents=contents,
-        config=genai.types.GenerateContentConfig(
-            system_instruction=self.system_instruction,
-            max_output_tokens=self.max_output_tokens,
-            temperature=self.temperature,
-            response_mime_type=self.config.response_mime_type,
-        ),
-    )
-    response_text = (response.text or "") if response else ""
-
     if response_text.strip():
       self.history.append(("user", prompt))
       self.history.append(("model", response_text))
     return response_text
 
 
+# Backwards-compatible alias for the upstream class name.
+GeminiAgent = LLMAgent
+
+
 @functools.total_ordering
-class StudentAgent(GeminiAgent):
+class StudentAgent(LLMAgent):
   """Student agent that generates solutions and learns from feedback."""
 
   def __init__(
@@ -219,7 +211,10 @@ def __init__(
         run_base_dir: The base directory for the current run.
         seed_solution: Initial model response (solution) to start with.
     """
-    print(f"Initializing Student Agent ({config.name}) using Gemini...")
+    print(
+        f"Initializing Student Agent ({config.name}) using"
+        f" {config.backend}:{config.model_name}..."
+    )
     super().__init__(config, run_base_dir)
     self.current_solution: str = seed_solution or ""  # Full model output.
     self.student_dir: str = os.path.join(self.run_base_dir, "students")
@@ -310,11 +305,7 @@ def generate_solution(self, prompt: str) -> str:
         solution = super().generate(prompt)
         self.current_solution = solution
         return solution
-      except (
-          genai.errors.ServerError,
-          genai.errors.ClientError,
-          genai.errors.httpx.HTTPError,
-      ) as e:
+      except self.backend.retryable_exceptions as e:
         last_exception = e
         error_type = type(e).__name__
         if attempt < self.config.max_code_generation_retries - 1:
@@ -338,6 +329,9 @@ def generate_solution(self, prompt: str) -> str:
               e,
               exc_info=True,
           )
+      except self.backend.fatal_exceptions as e:
+        logger.error("[%s] Non-retryable LLM API error: %s", self.name, e)
+        raise GenerationError(f"Non-retryable LLM API error: {e}") from e
     raise GenerationError(
         "Failed to generate solution after"
         f" {self.config.max_code_generation_retries} attempts due to API"
@@ -522,7 +516,7 @@ def mutate(
     self.generate_solution(prompt + "\n" + final_mutation_prompt)
 
 
-class SolutionAnalysisAgent(GeminiAgent):
+class SolutionAnalysisAgent(LLMAgent):
   """Analyzes the best student model solution using an LLM agent.
 
   Extracts metadata (tags, description, insights, etc.), writes a JSON
@@ -666,11 +660,7 @@ def analyze(
             " requested schema."
         )
 
-      except (
-          genai.errors.ServerError,
-          genai.errors.ClientError,
-          genai.errors.httpx.HTTPError,
-      ) as e:
+      except self.backend.retryable_exceptions as e:
         last_exception = e
         error_type = type(e).__name__
         if attempt < max_retries - 1:
@@ -692,6 +682,8 @@ def analyze(
               e,
               exc_info=True,
           )
+      except self.backend.fatal_exceptions as e:
+        raise ValueError(f"Non-retryable LLM API error: {e}") from e
 
     raise ValueError(
         "Failed to obtain and parse valid JSON from model."

codegen_classroom/agents_test.py

@@ -1009,16 +1009,6 @@ def test_load_env_with_custom_path(self):
       agent_config.load_env(env_path=env_file.full_path)
       self.assertEqual(os.environ.get('GEMINI_API_KEY'), 'custom_path_key')
 
-  def test_load_env_from_colab_userdata(self):
-    mock_userdata = mock.Mock()
-    mock_userdata.get.return_value = 'colab_test_key'
-    with mock.patch.dict(os.environ, {}, clear=True):
-      with mock.patch.dict(
-          sys.modules, {'google.colab.userdata': mock_userdata}
-      ):
-        agent_config.load_env()
-        self.assertEqual(os.environ.get('GEMINI_API_KEY'), 'colab_test_key')
-
 
 if __name__ == '__main__':
   absltest.main()