feat: add --server flag and set local AI service timeout

codeflash/api/aiservice.py

@@ -47,7 +47,7 @@ def __init__(self) -> None:
         self.headers = {"Authorization": f"Bearer {get_codeflash_api_key()}", "Connection": "close"}
         self.llm_call_counter = count(1)
         self.is_local = self.base_url == "http://localhost:8000"
-        self.timeout: float | None = None if self.is_local else 90
+        self.timeout: float | None = 300 if self.is_local else 90
 
     def get_next_sequence(self) -> int:
         """Get the next LLM call sequence number."""

codeflash/cli_cmds/cli.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import sys
 from argparse import SUPPRESS, ArgumentParser, Namespace
 from pathlib import Path
@@ -107,6 +108,12 @@ def parse_args() -> Namespace:
         action="store_true",
         help="(Deprecated) Async function optimization is now enabled by default. This flag is ignored.",
     )
+    parser.add_argument(
+        "--server",
+        type=str,
+        choices=["local", "prod"],
+        help="AI service server to use: 'local' for localhost:8000, 'prod' for app.codeflash.ai",
+    )
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]
@@ -121,6 +128,9 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
     )
     from codeflash.code_utils.github_utils import require_github_app_or_exit
 
+    if args.server:
+        os.environ["CODEFLASH_AIS_SERVER"] = args.server
+
     is_init: bool = args.command.startswith("init") if args.command else False
     if args.verbose:
         logging_config.set_level(logging.DEBUG, echo_setting=not is_init)