Merge pull request #26 from google/ericpts/model_flag

Make the model configurable via flag.

Author: ericpts

agent.py

@@ -38,9 +38,7 @@ def __init__(
         self,
         browser_computer: Computer,
         query: str,
-        model_name: Literal[
-            "computer-use-exp-6-11"
-        ] = "computer-use-exp-6-11",
+        model_name: str
     ):
         self._browser_computer = browser_computer
         self._query = query

main.py

@@ -59,6 +59,11 @@ def main() -> int:
         default=False,
         help="If possible, highlight the location of the mouse.",
     )
+    parser.add_argument(
+        "--model",
+        default='computer-use-exp-6-11',
+        help="Set which main model to use.",
+    )
     args = parser.parse_args()
 
     if args.env == "cloud-run":
@@ -83,6 +88,7 @@ def main() -> int:
         agent = BrowserAgent(
             browser_computer=browser_computer,
             query=args.query,
+            model_name=args.model,
         )
         agent.agent_loop()
     return 0