Add call limit wrapper to launch scripts and colabs.

PiperOrigin-RevId: 676425999
Change-Id: I59940006e942ec2ccbd79a4fffdeb20d90f74fd6

examples/modular/launch.py

@@ -46,6 +46,7 @@
 import datetime
 import importlib
 
+from concordia.language_model import call_limit_wrapper
 from concordia.language_model import utils
 from concordia.utils import measurements as measurements_lib
 import numpy as np
@@ -119,6 +120,7 @@
     embedder=embedder,
     measurements=measurements,
     agent_module=agent_module,
+    override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),
 )
 # Run the simulation
 _, results_log = runnable_simulation()

examples/modular/launch_concordia_challenge_evaluation.py

@@ -72,6 +72,7 @@
 import importlib
 import os
 
+from concordia.language_model import call_limit_wrapper
 from concordia.language_model import utils
 from concordia.utils import concurrency
 from concordia.utils import measurements as measurements_lib
@@ -200,6 +201,7 @@ def _evaluate_all_repetitions_on_one_scenario(
         focal_agent_module=agent_module,
         embedder=embedder,
         measurements=measurements,
+        override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),
     )
     # Run the simulation
     outcome, text_results_log = runnable_simulation()

examples/modular/launch_one_scenario.py

@@ -72,6 +72,7 @@
 import datetime
 import importlib
 
+from concordia.language_model import call_limit_wrapper
 from concordia.language_model import utils
 from concordia.utils import measurements as measurements_lib
 import numpy as np
@@ -116,10 +117,7 @@
     default=1,
     dest='num_repetitions_per_scenario',
 )
-parser.add_argument('--api_key',
-                    action='store',
-                    default=None,
-                    dest='api_key')
+parser.add_argument('--api_key', action='store', default=None, dest='api_key')
 parser.add_argument(
     '--disable_language_model',
     action='store_true',
@@ -188,6 +186,7 @@
       focal_agent_module=agent_module,
       embedder=embedder,
       measurements=measurements,
+      override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),
   )
   # Run the simulation
   outcome, text_results_log = runnable_simulation()
@@ -230,9 +229,7 @@
     background_per_capita_score=np.mean(
         background_per_capita_scores_to_average
     ),
-    ungrouped_per_capita_score=np.mean(
-        ungrouped_per_capita_scores_to_average
-    ),
+    ungrouped_per_capita_score=np.mean(ungrouped_per_capita_scores_to_average),
     simulation_outcomes=tuple(simulation_outcomes),
     focal_is_resident=scenario_config.focal_is_resident,
     api_type=args.api_type,

examples/modular/launch_resident_visitor.py

@@ -47,6 +47,7 @@
 import datetime
 import importlib
 
+from concordia.language_model import call_limit_wrapper
 from concordia.language_model import utils
 from concordia.utils import measurements as measurements_lib
 import numpy as np
@@ -128,6 +129,7 @@
     embedder=embedder,
     measurements=measurements,
     resident_visitor_modules=(resident_agent_module, visitor_agent_module),
+    override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),
 )
 # Run the simulation
 _, results_log = runnable_simulation()

examples/modular/notebook.ipynb

@@ -85,6 +85,7 @@
         "import datetime\n",
         "import importlib\n",
         "\n",
+        "from concordia.language_model import call_limit_wrapper\n",
         "from concordia.language_model import utils\n",
         "from concordia.utils import measurements as measurements_lib\n",
         "import sentence_transformers"
@@ -240,6 +241,7 @@
         "    embedder=embedder,\n",
         "    measurements=measurements,\n",
         "    agent_module=agent_module,\n",
+        "    override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),   \n",
         ")"
       ]
     },

examples/modular/scenario/scenarios.py

@@ -296,6 +296,7 @@ def build_simulation(
     focal_agent_module: types.ModuleType,
     embedder: Callable[[str], np.ndarray],
     measurements: measurements_lib.Measurements,
+    override_agent_model: language_model.LanguageModel | None = None,
     agent_base_module: str = DEFAULT_IMPORT_AGENT_BASE_MODULE,
     support_agent_base_module: str = DEFAULT_IMPORT_SUPPORT_AGENT_MODULE,
     env_base_module: str = DEFAULT_IMPORT_ENV_BASE_MODULE,
@@ -337,6 +338,7 @@ def build_simulation(
       model=model,
       embedder=embedder,
       measurements=measurements,
+      override_agent_model=override_agent_model,
       resident_visitor_modules=(resident_agent_module, visitor_agent_module),
       supporting_agent_module=supporting_agent_module,
       time_and_place_module=scenario_config.time_and_place_module,

examples/tutorials/agent_development.ipynb

@@ -96,6 +96,7 @@
         "\n",
         "from IPython import display\n",
         "\n",
+        "from concordia.language_model import call_limit_wrapper\n",
         "from concordia.language_model import utils\n",
         "from concordia.utils import measurements as measurements_lib\n",
         "import sentence_transformers"
@@ -657,6 +658,7 @@
         "    embedder=embedder,\n",
         "    measurements=measurements,\n",
         "    focal_agent_module=agent_module,\n",
+        "    override_agent_model=call_limit_wrapper.CallLimitLanguageModel(model),\n",
         ")"
       ]
     },