Make model configurable, and get API key from runtime config (#32)

.env.example

@@ -8,4 +8,4 @@ OPENAI_API_KEY=""
 
 SUPABASE_URL=""
 # Ensure this is your Supabase Service Role key
-SUPABASE_KEY=""
+SUPABASE_KEY=""

README.md

@@ -104,3 +104,8 @@ For creation methods, it auto-injects the user's ID into the metadata. This is t
 By using custom authentication, we can call this LangGraph server directly from a frontend application, without having to worry about exposing API keys/secrets, since you only need a JWT token from Supabase to authenticate.
 
 For more info, see our [LangGraph custom auth docs](https://langchain-ai.github.io/langgraph/tutorials/auth/getting_started/).
+
+## Model API Keys
+This agent is configured to accept model API keys set by a user in the OAP settings page. The default OAP implementation allows users to add their own OpenAI, Anthropic, and Google API keys, and these can easily be extended in OAP.
+
+If a user sets their API keys in the OAP settings page, those keys will be passed to the agent through the config in an 'apiKeys' field and used by default. Otherwise, we will automatically fall back on the environment variables set in the agent deployed on LangGraph platform.

oap_supervisor/agent.py

@@ -1,9 +1,10 @@
+import os
 from langgraph.pregel.remote import RemoteGraph
-from langchain_openai import ChatOpenAI
 from langgraph_supervisor import create_supervisor
 from pydantic import BaseModel, Field
 from typing import List, Optional
 from langchain_core.runnables import RunnableConfig
+from langchain.chat_models import init_chat_model
 
 # This system prompt is ALWAYS included at the bottom of the message.
 UNEDITABLE_SYSTEM_PROMPT = """\nYou can invoke sub-agents by calling tools in this format:
@@ -45,6 +46,10 @@ class GraphConfigPydantic(BaseModel):
             }
         },
     )
+    supervisor_model: str = Field(
+        default="openai:gpt-4.1",
+        metadata={"x_oap_ui_config": {"type": "text", "placeholder": "Enter the model to use for the supervisor."}},
+    )
 
 
 class OAPRemoteGraph(RemoteGraph):
@@ -117,9 +122,30 @@ def make_child_graphs(cfg: GraphConfigPydantic, access_token: Optional[str] = No
     return [create_remote_graph_wrapper(a) for a in cfg.agents]
 
 
-def make_model(cfg: GraphConfigPydantic):
+def get_api_key_for_model(model_name: str, config: RunnableConfig):
+    model_name = model_name.lower()
+    model_to_key = {
+        "openai:": "OPENAI_API_KEY",
+        "anthropic:": "ANTHROPIC_API_KEY", 
+        "google": "GOOGLE_API_KEY"
+    }
+    key_name = next((key for prefix, key in model_to_key.items() 
+                    if model_name.startswith(prefix)), None)
+    if not key_name:
+        return None
+    api_keys = config.get("configurable", {}).get("apiKeys", {})
+    if api_keys and api_keys.get(key_name) and len(api_keys[key_name]) > 0:
+        return api_keys[key_name]
+    # Fallback to environment variable
+    return os.getenv(key_name)
+
+
+def make_model(cfg: GraphConfigPydantic, model_api_key: str):
     """Instantiate the LLM for the supervisor based on the config."""
-    return ChatOpenAI(model="gpt-4o")
+    return init_chat_model(
+        model=cfg.supervisor_model,
+        api_key=model_api_key
+    )
 
 
 def make_prompt(cfg: GraphConfigPydantic):
@@ -136,9 +162,12 @@ def graph(config: RunnableConfig):
     # Pass the token to make_child_graphs, which now handles None values
     child_graphs = make_child_graphs(cfg, supabase_access_token)
 
+    # Get the API key from the RunnableConfig or from the environment variable
+    model_api_key = get_api_key_for_model(cfg.supervisor_model, config) or "No token found"
+
     return create_supervisor(
         child_graphs,
-        model=make_model(cfg),
+        model=make_model(cfg, model_api_key),
         prompt=make_prompt(cfg),
         config_schema=GraphConfigPydantic,
         handoff_tool_prefix="delegate_to_",

pyproject.toml

@@ -14,6 +14,9 @@ dependencies = [
     "supabase>=2.15.1",
     "aiohttp>=3.8.0",
     "langgraph-supervisor>=0.0.27",
+    "langchain>=0.3.26",
+    "langchain-anthropic>=0.3.16",
+    "langchain-google-genai>=2.1.6",
 ]
 
 [tool.setuptools]