weco 0.2.23__tar.gz → 0.2.25__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {weco-0.2.23 → weco-0.2.25}/.gitignore +2 -0
  2. {weco-0.2.23 → weco-0.2.25}/PKG-INFO +31 -18
  3. {weco-0.2.23 → weco-0.2.25}/README.md +28 -17
  4. weco-0.2.25/contributing.md +23 -0
  5. {weco-0.2.23 → weco-0.2.25}/pyproject.toml +3 -1
  6. {weco-0.2.23 → weco-0.2.25}/weco/api.py +60 -27
  7. {weco-0.2.23 → weco-0.2.25}/weco/auth.py +7 -5
  8. {weco-0.2.23 → weco-0.2.25}/weco/chatbot.py +34 -23
  9. {weco-0.2.23 → weco-0.2.25}/weco/cli.py +10 -1
  10. weco-0.2.25/weco/constants.py +7 -0
  11. {weco-0.2.23 → weco-0.2.25}/weco/optimizer.py +26 -21
  12. {weco-0.2.23 → weco-0.2.25}/weco/panels.py +105 -59
  13. {weco-0.2.23 → weco-0.2.25}/weco/utils.py +47 -12
  14. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/PKG-INFO +31 -18
  15. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/SOURCES.txt +2 -1
  16. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/requires.txt +2 -0
  17. weco-0.2.23/.repomixignore +0 -4
  18. {weco-0.2.23 → weco-0.2.25}/.github/workflows/lint.yml +0 -0
  19. {weco-0.2.23 → weco-0.2.25}/.github/workflows/release.yml +0 -0
  20. {weco-0.2.23 → weco-0.2.25}/LICENSE +0 -0
  21. {weco-0.2.23 → weco-0.2.25}/assets/example-optimization.gif +0 -0
  22. {weco-0.2.23 → weco-0.2.25}/assets/weco.svg +0 -0
  23. {weco-0.2.23 → weco-0.2.25}/examples/cuda/README.md +0 -0
  24. {weco-0.2.23 → weco-0.2.25}/examples/cuda/evaluate.py +0 -0
  25. {weco-0.2.23 → weco-0.2.25}/examples/cuda/guide.md +0 -0
  26. {weco-0.2.23 → weco-0.2.25}/examples/cuda/optimize.py +0 -0
  27. {weco-0.2.23 → weco-0.2.25}/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb +0 -0
  28. {weco-0.2.23 → weco-0.2.25}/examples/hello-kernel-world/evaluate.py +0 -0
  29. {weco-0.2.23 → weco-0.2.25}/examples/hello-kernel-world/optimize.py +0 -0
  30. {weco-0.2.23 → weco-0.2.25}/examples/prompt/README.md +0 -0
  31. {weco-0.2.23 → weco-0.2.25}/examples/prompt/eval.py +0 -0
  32. {weco-0.2.23 → weco-0.2.25}/examples/prompt/optimize.py +0 -0
  33. {weco-0.2.23 → weco-0.2.25}/examples/prompt/prompt_guide.md +0 -0
  34. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/README.md +0 -0
  35. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/competition_description.md +0 -0
  36. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/data/sample_submission.csv +0 -0
  37. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/data/test.csv +0 -0
  38. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/data/train.csv +0 -0
  39. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/evaluate.py +0 -0
  40. {weco-0.2.23 → weco-0.2.25}/examples/spaceship-titanic/train.py +0 -0
  41. {weco-0.2.23 → weco-0.2.25}/examples/triton/README.md +0 -0
  42. {weco-0.2.23 → weco-0.2.25}/examples/triton/evaluate.py +0 -0
  43. {weco-0.2.23 → weco-0.2.25}/examples/triton/optimize.py +0 -0
  44. {weco-0.2.23 → weco-0.2.25}/setup.cfg +0 -0
  45. {weco-0.2.23 → weco-0.2.25}/weco/__init__.py +0 -0
  46. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/dependency_links.txt +0 -0
  47. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/entry_points.txt +0 -0
  48. {weco-0.2.23 → weco-0.2.25}/weco.egg-info/top_level.txt +0 -0
@@ -69,8 +69,10 @@ etc/
69
69
  # AI generated files
70
70
  digest.txt
71
71
  .runs/
72
+ CLAUDE.md
72
73
 
73
74
  *.pyc
74
75
 
75
76
  # Repomix
77
+ .repomixignore
76
78
  repomix-output.*
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: weco
3
- Version: 0.2.23
3
+ Version: 0.2.25
4
4
  Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
5
5
  Author-email: Weco AI Team <contact@weco.ai>
6
6
  License: MIT
@@ -16,6 +16,8 @@ Requires-Dist: requests
16
16
  Requires-Dist: rich
17
17
  Requires-Dist: packaging
18
18
  Requires-Dist: gitingest
19
+ Requires-Dist: fastapi
20
+ Requires-Dist: slowapi
19
21
  Provides-Extra: dev
20
22
  Requires-Dist: ruff; extra == "dev"
21
23
  Requires-Dist: build; extra == "dev"
@@ -158,6 +160,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
158
160
  | `-M, --model` | Model identifier for the LLM to use (e.g., `o4-mini`, `claude-sonnet-4-0`). | `o4-mini` when `OPENAI_API_KEY` is set; `claude-sonnet-4-0` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro` when `GEMINI_API_KEY` is set. | `-M o4-mini` |
159
161
  | `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM. | `None` | `-i instructions.md` or `-i "Optimize the model for faster inference"`|
160
162
  | `-l, --log-dir` | Path to the directory to log intermediate steps and final optimization result. | `.runs/` | `-l ./logs/` |
163
+ | `--eval-timeout` | Timeout in seconds for each step in evaluation. | No timeout (unlimited) | `--eval-timeout 3600` |
161
164
 
162
165
  ---
163
166
 
@@ -248,28 +251,38 @@ Final speedup value = 1.5
248
251
 
249
252
  Weco will parse this output to extract the numerical value (1.5 in this case) associated with the metric name ('speedup').
250
253
 
251
- ## Contributing
254
+ ## Supported Models
252
255
 
253
- We welcome your contributions! To get started:
256
+ Weco supports the following LLM models:
254
257
 
255
- 1. **Fork & Clone the Repository:**
256
- ```bash
257
- git clone https://github.com/WecoAI/weco-cli.git
258
- cd weco-cli
259
- ```
258
+ ### OpenAI Models
259
+ - `o3`
260
+ - `o3-mini`
261
+ - `o4-mini`
262
+ - `o1-pro`
263
+ - `o1`
264
+ - `gpt-4.1`
265
+ - `gpt-4.1-mini`
266
+ - `gpt-4.1-nano`
267
+ - `gpt-4o`
268
+ - `gpt-4o-mini`
260
269
 
261
- 2. **Install Dependencies:**
262
- ```bash
263
- pip install -e ".[dev]"
264
- ```
270
+ ### Anthropic Models
271
+ - `claude-opus-4-0`
272
+ - `claude-sonnet-4-0`
273
+ - `claude-3-7-sonnet-latest`
265
274
 
266
- 3. **Create a Feature Branch:**
267
- ```bash
268
- git checkout -b feature/your-feature-name
269
- ```
275
+ ### Gemini Models
276
+ - `gemini-2.5-pro`
277
+ - `gemini-2.5-flash`
278
+ - `gemini-2.5-flash-lite`
279
+
280
+ You can specify any of these models using the `-M` or `--model` flag. Ensure you have the corresponding API key set as an environment variable for the model provider you wish to use.
270
281
 
271
- 4. **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
282
+ ---
283
+
284
+ ## Contributing
272
285
 
273
- 5. **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.
286
+ We welcome contributions! Please see [contributing.md](contributing.md) for detailed guidelines on how to contribute to this project.
274
287
 
275
288
  ---
@@ -134,6 +134,7 @@ For more advanced examples, including [Triton](/examples/triton/README.md), [CUD
134
134
  | `-M, --model` | Model identifier for the LLM to use (e.g., `o4-mini`, `claude-sonnet-4-0`). | `o4-mini` when `OPENAI_API_KEY` is set; `claude-sonnet-4-0` when `ANTHROPIC_API_KEY` is set; `gemini-2.5-pro` when `GEMINI_API_KEY` is set. | `-M o4-mini` |
135
135
  | `-i, --additional-instructions`| Natural language description of specific instructions **or** path to a file containing detailed instructions to guide the LLM. | `None` | `-i instructions.md` or `-i "Optimize the model for faster inference"`|
136
136
  | `-l, --log-dir` | Path to the directory to log intermediate steps and final optimization result. | `.runs/` | `-l ./logs/` |
137
+ | `--eval-timeout` | Timeout in seconds for each step in evaluation. | No timeout (unlimited) | `--eval-timeout 3600` |
137
138
 
138
139
  ---
139
140
 
@@ -224,28 +225,38 @@ Final speedup value = 1.5
224
225
 
225
226
  Weco will parse this output to extract the numerical value (1.5 in this case) associated with the metric name ('speedup').
226
227
 
227
- ## Contributing
228
+ ## Supported Models
228
229
 
229
- We welcome your contributions! To get started:
230
+ Weco supports the following LLM models:
230
231
 
231
- 1. **Fork & Clone the Repository:**
232
- ```bash
233
- git clone https://github.com/WecoAI/weco-cli.git
234
- cd weco-cli
235
- ```
232
+ ### OpenAI Models
233
+ - `o3`
234
+ - `o3-mini`
235
+ - `o4-mini`
236
+ - `o1-pro`
237
+ - `o1`
238
+ - `gpt-4.1`
239
+ - `gpt-4.1-mini`
240
+ - `gpt-4.1-nano`
241
+ - `gpt-4o`
242
+ - `gpt-4o-mini`
236
243
 
237
- 2. **Install Dependencies:**
238
- ```bash
239
- pip install -e ".[dev]"
240
- ```
244
+ ### Anthropic Models
245
+ - `claude-opus-4-0`
246
+ - `claude-sonnet-4-0`
247
+ - `claude-3-7-sonnet-latest`
241
248
 
242
- 3. **Create a Feature Branch:**
243
- ```bash
244
- git checkout -b feature/your-feature-name
245
- ```
249
+ ### Gemini Models
250
+ - `gemini-2.5-pro`
251
+ - `gemini-2.5-flash`
252
+ - `gemini-2.5-flash-lite`
253
+
254
+ You can specify any of these models using the `-M` or `--model` flag. Ensure you have the corresponding API key set as an environment variable for the model provider you wish to use.
246
255
 
247
- 4. **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
256
+ ---
257
+
258
+ ## Contributing
248
259
 
249
- 5. **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.
260
+ We welcome contributions! Please see [contributing.md](contributing.md) for detailed guidelines on how to contribute to this project.
250
261
 
251
262
  ---
@@ -0,0 +1,23 @@
1
+ # Contributing
2
+
3
+ We welcome your contributions! To get started:
4
+
5
+ 1. **Fork & Clone the Repository:**
6
+ ```bash
7
+ git clone https://github.com/WecoAI/weco-cli.git
8
+ cd weco-cli
9
+ ```
10
+
11
+ 2. **Install Dependencies:**
12
+ ```bash
13
+ pip install -e ".[dev]"
14
+ ```
15
+
16
+ 3. **Create a Feature Branch:**
17
+ ```bash
18
+ git checkout -b feature/your-feature-name
19
+ ```
20
+
21
+ 4. **Make Changes:** Ensure your code adheres to our style guidelines and includes relevant tests.
22
+
23
+ 5. **Commit, Push & Open a PR**: Commit your changes, and open a pull request with a clear description of your enhancements.
@@ -8,7 +8,7 @@ name = "weco"
8
8
  authors = [{ name = "Weco AI Team", email = "contact@weco.ai" }]
9
9
  description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
10
10
  readme = "README.md"
11
- version = "0.2.23"
11
+ version = "0.2.25"
12
12
  license = { text = "MIT" }
13
13
  requires-python = ">=3.8"
14
14
  dependencies = [
@@ -16,6 +16,8 @@ dependencies = [
16
16
  "rich",
17
17
  "packaging",
18
18
  "gitingest",
19
+ "fastapi",
20
+ "slowapi",
19
21
  ]
20
22
  keywords = ["AI", "Code Optimization", "Code Generation"]
21
23
  classifiers = [
@@ -4,6 +4,7 @@ import requests
4
4
  from rich.console import Console
5
5
 
6
6
  from weco import __pkg_version__, __base_url__
7
+ from .constants import DEFAULT_API_TIMEOUT
7
8
 
8
9
 
9
10
  def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None:
@@ -30,7 +31,7 @@ def start_optimization_run(
30
31
  additional_instructions: str = None,
31
32
  api_keys: Dict[str, Any] = {},
32
33
  auth_headers: dict = {},
33
- timeout: Union[int, Tuple[int, int]] = 800,
34
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
34
35
  ) -> Dict[str, Any]:
35
36
  """Start the optimization run."""
36
37
  with console.status("[bold green]Starting Optimization..."):
@@ -53,22 +54,29 @@ def start_optimization_run(
53
54
  timeout=timeout,
54
55
  )
55
56
  response.raise_for_status()
56
- return response.json()
57
+ result = response.json()
58
+ # Handle None values for code and plan fields
59
+ if result.get("plan") is None:
60
+ result["plan"] = ""
61
+ if result.get("code") is None:
62
+ result["code"] = ""
63
+ return result
57
64
  except requests.exceptions.HTTPError as e:
58
65
  handle_api_error(e, console)
59
- sys.exit(1)
66
+ raise
60
67
  except Exception as e:
61
68
  console.print(f"[bold red]Error starting run: {e}[/]")
62
- sys.exit(1)
69
+ raise
63
70
 
64
71
 
65
72
  def evaluate_feedback_then_suggest_next_solution(
73
+ console: Console,
66
74
  run_id: str,
67
75
  execution_output: str,
68
76
  additional_instructions: str = None,
69
77
  api_keys: Dict[str, Any] = {},
70
78
  auth_headers: dict = {},
71
- timeout: Union[int, Tuple[int, int]] = 800,
79
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
72
80
  ) -> Dict[str, Any]:
73
81
  """Evaluate the feedback and suggest the next solution."""
74
82
  try:
@@ -83,10 +91,17 @@ def evaluate_feedback_then_suggest_next_solution(
83
91
  timeout=timeout,
84
92
  )
85
93
  response.raise_for_status()
86
- return response.json()
94
+ result = response.json()
95
+ # Handle None values for code and plan fields
96
+ if result.get("plan") is None:
97
+ result["plan"] = ""
98
+ if result.get("code") is None:
99
+ result["code"] = ""
100
+
101
+ return result
87
102
  except requests.exceptions.HTTPError as e:
88
103
  # Allow caller to handle suggest errors, maybe retry or terminate
89
- handle_api_error(e, Console()) # Use default console if none passed
104
+ handle_api_error(e, console) # Use default console if none passed
90
105
  raise # Re-raise the exception
91
106
  except Exception as e:
92
107
  print(f"Error: {e}") # Use print as console might not be available
@@ -94,7 +109,11 @@ def evaluate_feedback_then_suggest_next_solution(
94
109
 
95
110
 
96
111
  def get_optimization_run_status(
97
- run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 800
112
+ console: Console,
113
+ run_id: str,
114
+ include_history: bool = False,
115
+ auth_headers: dict = {},
116
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
98
117
  ) -> Dict[str, Any]:
99
118
  """Get the current status of the optimization run."""
100
119
  try:
@@ -102,16 +121,30 @@ def get_optimization_run_status(
102
121
  f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
103
122
  )
104
123
  response.raise_for_status()
105
- return response.json()
124
+ result = response.json()
125
+ # Handle None values for code and plan fields in best_result and nodes
126
+ if result.get("best_result"):
127
+ if result["best_result"].get("code") is None:
128
+ result["best_result"]["code"] = ""
129
+ if result["best_result"].get("plan") is None:
130
+ result["best_result"]["plan"] = ""
131
+ # Handle None values for code and plan fields in nodes array
132
+ if result.get("nodes"):
133
+ for i, node in enumerate(result["nodes"]):
134
+ if node.get("plan") is None:
135
+ result["nodes"][i]["plan"] = ""
136
+ if node.get("code") is None:
137
+ result["nodes"][i]["code"] = ""
138
+ return result
106
139
  except requests.exceptions.HTTPError as e:
107
- handle_api_error(e, Console()) # Use default console
140
+ handle_api_error(e, console) # Use default console
108
141
  raise # Re-raise
109
142
  except Exception as e:
110
143
  print(f"Error getting run status: {e}")
111
144
  raise # Re-raise
112
145
 
113
146
 
114
- def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = 10) -> bool:
147
+ def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (10, 10)) -> bool:
115
148
  """Send a heartbeat signal to the backend."""
116
149
  try:
117
150
  response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
@@ -119,9 +152,9 @@ def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tup
119
152
  return True
120
153
  except requests.exceptions.HTTPError as e:
121
154
  if e.response.status_code == 409:
122
- print(f"Heartbeat ignored: Run {run_id} is not running.", file=sys.stderr)
155
+ print(f"Polling ignore: Run {run_id} is not running.", file=sys.stderr)
123
156
  else:
124
- print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
157
+ print(f"Polling failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
125
158
  return False
126
159
  except Exception as e:
127
160
  print(f"Error sending heartbeat for run {run_id}: {e}", file=sys.stderr)
@@ -134,7 +167,7 @@ def report_termination(
134
167
  reason: str,
135
168
  details: Optional[str] = None,
136
169
  auth_headers: dict = {},
137
- timeout: Union[int, Tuple[int, int]] = 30,
170
+ timeout: Union[int, Tuple[int, int]] = (10, 30),
138
171
  ) -> bool:
139
172
  """Report the termination reason to the backend."""
140
173
  try:
@@ -172,22 +205,22 @@ def _determine_model_and_api_key() -> tuple[str, dict[str, str]]:
172
205
  api_key_dict = {"GEMINI_API_KEY": llm_api_keys["GEMINI_API_KEY"]}
173
206
  else:
174
207
  # This should never happen if determine_default_model works correctly
175
- raise ValueError(f"Unknown model returned: {model}")
208
+ raise ValueError(f"Unknown default model choice: {model}")
176
209
 
177
210
  return model, api_key_dict
178
211
 
179
212
 
180
213
  def get_optimization_suggestions_from_codebase(
214
+ console: Console,
181
215
  gitingest_summary: str,
182
216
  gitingest_tree: str,
183
217
  gitingest_content_str: str,
184
- console: Console,
185
218
  auth_headers: dict = {},
186
- timeout: Union[int, Tuple[int, int]] = 800,
219
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
187
220
  ) -> Optional[List[Dict[str, Any]]]:
188
221
  """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
222
+ model, api_key_dict = _determine_model_and_api_key()
189
223
  try:
190
- model, api_key_dict = _determine_model_and_api_key()
191
224
  response = requests.post(
192
225
  f"{__base_url__}/onboard/analyze-codebase",
193
226
  json={
@@ -213,16 +246,16 @@ def get_optimization_suggestions_from_codebase(
213
246
 
214
247
 
215
248
  def generate_evaluation_script_and_metrics(
249
+ console: Console,
216
250
  target_file: str,
217
251
  description: str,
218
252
  gitingest_content_str: str,
219
- console: Console,
220
253
  auth_headers: dict = {},
221
- timeout: Union[int, Tuple[int, int]] = 800,
254
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
222
255
  ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
223
256
  """Generate evaluation script and determine metrics using the model-agnostic backend API."""
257
+ model, api_key_dict = _determine_model_and_api_key()
224
258
  try:
225
- model, api_key_dict = _determine_model_and_api_key()
226
259
  response = requests.post(
227
260
  f"{__base_url__}/onboard/generate-script",
228
261
  json={
@@ -247,18 +280,18 @@ def generate_evaluation_script_and_metrics(
247
280
 
248
281
 
249
282
  def analyze_evaluation_environment(
283
+ console: Console,
250
284
  target_file: str,
251
285
  description: str,
252
286
  gitingest_summary: str,
253
287
  gitingest_tree: str,
254
288
  gitingest_content_str: str,
255
- console: Console,
256
289
  auth_headers: dict = {},
257
- timeout: Union[int, Tuple[int, int]] = 800,
290
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
258
291
  ) -> Optional[Dict[str, Any]]:
259
292
  """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
293
+ model, api_key_dict = _determine_model_and_api_key()
260
294
  try:
261
- model, api_key_dict = _determine_model_and_api_key()
262
295
  response = requests.post(
263
296
  f"{__base_url__}/onboard/analyze-environment",
264
297
  json={
@@ -285,16 +318,16 @@ def analyze_evaluation_environment(
285
318
 
286
319
 
287
320
  def analyze_script_execution_requirements(
321
+ console: Console,
288
322
  script_content: str,
289
323
  script_path: str,
290
324
  target_file: str,
291
- console: Console,
292
325
  auth_headers: dict = {},
293
- timeout: Union[int, Tuple[int, int]] = 800,
326
+ timeout: Union[int, Tuple[int, int]] = DEFAULT_API_TIMEOUT,
294
327
  ) -> Optional[str]:
295
328
  """Analyze script to determine proper execution command using the model-agnostic backend API."""
329
+ model, api_key_dict = _determine_model_and_api_key()
296
330
  try:
297
- model, api_key_dict = _determine_model_and_api_key()
298
331
  response = requests.post(
299
332
  f"{__base_url__}/onboard/analyze-script",
300
333
  json={
@@ -35,7 +35,7 @@ def save_api_key(api_key: str):
35
35
  # Set file permissions to read/write for owner only (600)
36
36
  os.chmod(CREDENTIALS_FILE, stat.S_IRUSR | stat.S_IWUSR)
37
37
  except OSError as e:
38
- print(f"Error: Could not write credentials file or set permissions on {CREDENTIALS_FILE}: {e}")
38
+ print(f"Error: Unable to save credentials file or set permissions on {CREDENTIALS_FILE}: {e}")
39
39
 
40
40
 
41
41
  def load_weco_api_key() -> str | None:
@@ -53,7 +53,7 @@ def load_weco_api_key() -> str | None:
53
53
  credentials = json.load(f)
54
54
  return credentials.get("api_key")
55
55
  except (IOError, json.JSONDecodeError, OSError) as e:
56
- print(f"Warning: Could not read or parse credentials file at {CREDENTIALS_FILE}: {e}")
56
+ print(f"Warning: Unable to read credentials file at {CREDENTIALS_FILE}: {e}")
57
57
  return None
58
58
 
59
59
 
@@ -64,7 +64,7 @@ def clear_api_key():
64
64
  os.remove(CREDENTIALS_FILE)
65
65
  print("Logged out successfully.")
66
66
  except OSError as e:
67
- print(f"Error: Could not remove credentials file at {CREDENTIALS_FILE}: {e}")
67
+ print(f"Error: Unable to remove credentials file at {CREDENTIALS_FILE}: {e}")
68
68
  else:
69
69
  print("Already logged out.")
70
70
 
@@ -129,7 +129,9 @@ def perform_login(console: Console):
129
129
  continue # Continue polling
130
130
  else:
131
131
  # Unexpected 202 response format
132
- console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
132
+ console.print(
133
+ f"\n[bold red]Error:[/] Received unexpected response from authentication server: {token_data}"
134
+ )
133
135
  return False
134
136
  # Check for standard OAuth2 errors (often 400 Bad Request)
135
137
  elif token_response.status_code == 400:
@@ -146,7 +148,7 @@ def perform_login(console: Console):
146
148
  console.print("\n[bold red]Error:[/] Authorization denied by user.")
147
149
  return False
148
150
  else: # invalid_grant, etc.
149
- error_desc = token_data.get("error_description", "Unknown error during polling.")
151
+ error_desc = token_data.get("error_description", "Unknown authentication error occurred.")
150
152
  console.print(f"\n[bold red]Error:[/] {error_desc} ({error_code})")
151
153
  return False
152
154
 
@@ -50,7 +50,7 @@ class UserInteractionHelper:
50
50
 
51
51
  if attempts >= max_retries:
52
52
  self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
53
- raise Exception("Maximum retry attempts exceeded")
53
+ raise Exception("Maximum retry attempts exceeded. Please try again.")
54
54
 
55
55
  # Show available options without the full prompt
56
56
  if choices:
@@ -66,7 +66,7 @@ class UserInteractionHelper:
66
66
  continue
67
67
 
68
68
  # This should never be reached due to the exception above, but just in case
69
- raise Exception("Unexpected error in choice selection")
69
+ raise Exception("Unexpected error while selecting a choice")
70
70
 
71
71
  def get_choice_numeric(self, prompt: str, max_number: int, default: int = None, max_retries: int = 5) -> int:
72
72
  """Get numeric choice with validation and error handling."""
@@ -87,7 +87,7 @@ class UserInteractionHelper:
87
87
 
88
88
  if attempts >= max_retries:
89
89
  self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
90
- raise Exception("Maximum retry attempts exceeded")
90
+ raise Exception("Maximum retry attempts exceeded. Please try again.")
91
91
 
92
92
  # Show valid range
93
93
  self.console.print(f"Please enter a number between [bold]1[/] and [bold]{max_number}[/]")
@@ -115,7 +115,7 @@ class UserInteractionHelper:
115
115
 
116
116
  if attempts >= max_retries:
117
117
  self.console.print(f"[red]Maximum retry attempts ({max_retries}) reached. Exiting.[/]")
118
- raise Exception("Maximum retry attempts exceeded")
118
+ raise Exception("Maximum retry attempts exceeded. Please try again.")
119
119
 
120
120
  self.console.print("Valid options: [bold]y[/] / [bold]n[/]")
121
121
  if default:
@@ -123,7 +123,7 @@ class UserInteractionHelper:
123
123
 
124
124
  continue
125
125
 
126
- raise Exception("Unexpected error in yes/no selection")
126
+ raise Exception("Unexpected error while selecting an option")
127
127
 
128
128
  def display_optimization_options_table(self, options: List[Dict[str, str]]) -> None:
129
129
  """Display optimization options in a formatted table."""
@@ -215,7 +215,10 @@ class Chatbot:
215
215
 
216
216
  with self.console.status("[bold green]Generating optimization suggestions...[/]"):
217
217
  result = get_optimization_suggestions_from_codebase(
218
- self.gitingest_summary, self.gitingest_tree, self.gitingest_content_str, self.console
218
+ console=self.console,
219
+ gitingest_summary=self.gitingest_summary,
220
+ gitingest_tree=self.gitingest_tree,
221
+ gitingest_content_str=self.gitingest_content_str,
219
222
  )
220
223
 
221
224
  if result and isinstance(result, list):
@@ -224,7 +227,7 @@ class Chatbot:
224
227
  options = None
225
228
 
226
229
  if not options or not isinstance(options, list):
227
- self.console.print("[red]Failed to get valid optimization options.[/]")
230
+ self.console.print("[red]Unable to retrieve valid optimization options from the backend.[/]")
228
231
  return None
229
232
 
230
233
  if not options:
@@ -324,17 +327,19 @@ class Chatbot:
324
327
  elif action == "g" or action == "r":
325
328
  with self.console.status("[bold green]Generating evaluation script and determining metrics...[/]"):
326
329
  result = generate_evaluation_script_and_metrics(
327
- selected_option["target_file"],
328
- selected_option["description"],
329
- self.gitingest_content_str,
330
- self.console,
330
+ console=self.console,
331
+ target_file=selected_option["target_file"],
332
+ description=selected_option["description"],
333
+ gitingest_content_str=self.gitingest_content_str,
331
334
  )
332
335
  if result and result[0]:
333
336
  eval_script_content, metric_name, goal, reasoning = result
334
337
  if reasoning:
335
338
  self.console.print(f"[dim]Reasoning: {reasoning}[/]")
336
339
  else:
337
- self.console.print("[red]Failed to generate an evaluation script.[/]")
340
+ self.console.print(
341
+ "[red]Unable to generate an evaluation script. Please try providing a custom script path instead.[/]"
342
+ )
338
343
  eval_script_content = None
339
344
  metric_name = None
340
345
  goal = None
@@ -371,7 +376,10 @@ class Chatbot:
371
376
  # Analyze the script to determine the proper execution command
372
377
  with self.console.status("[bold green]Analyzing script execution requirements...[/]"):
373
378
  eval_command = analyze_script_execution_requirements(
374
- eval_script_content, eval_script_path_str, selected_option["target_file"], self.console
379
+ console=self.console,
380
+ script_content=eval_script_content,
381
+ script_path=eval_script_path_str,
382
+ target_file=selected_option["target_file"],
375
383
  )
376
384
 
377
385
  return {
@@ -386,16 +394,16 @@ class Chatbot:
386
394
  """Get or create evaluation script configuration using intelligent conversation-guided approach."""
387
395
  with self.console.status("[bold green]Analyzing evaluation environment...[/]"):
388
396
  analysis = analyze_evaluation_environment(
389
- selected_option["target_file"],
390
- selected_option["description"],
391
- self.gitingest_summary,
392
- self.gitingest_tree,
393
- self.gitingest_content_str,
394
- self.console,
397
+ console=self.console,
398
+ target_file=selected_option["target_file"],
399
+ description=selected_option["description"],
400
+ gitingest_summary=self.gitingest_summary,
401
+ gitingest_tree=self.gitingest_tree,
402
+ gitingest_content_str=self.gitingest_content_str,
395
403
  )
396
404
 
397
405
  if not analysis:
398
- self.console.print("[yellow]Failed to analyze evaluation environment. Falling back to generation.[/]")
406
+ self.console.print("[yellow]Unable to analyze evaluation environment. Falling back to script generation.[/]")
399
407
  return self.handle_script_generation_workflow(selected_option)
400
408
 
401
409
  self.evaluation_analysis = analysis
@@ -529,7 +537,10 @@ class Chatbot:
529
537
  if not eval_command or eval_command == f"python {script_path}":
530
538
  with self.console.status("[bold green]Analyzing script execution requirements...[/]"):
531
539
  eval_command = analyze_script_execution_requirements(
532
- script_content, script_path, selected_option["target_file"], self.console
540
+ console=self.console,
541
+ script_content=script_content,
542
+ script_path=script_path,
543
+ target_file=selected_option["target_file"],
533
544
  )
534
545
 
535
546
  self.current_step = "confirmation"
@@ -711,7 +722,7 @@ class Chatbot:
711
722
  """Setup evaluation environment for the selected optimization."""
712
723
  eval_config = self.get_evaluation_configuration(selected_option)
713
724
  if not eval_config:
714
- self.console.print("[red]Evaluation script setup failed.[/]")
725
+ self.console.print("[red]Evaluation script setup failed. Please check your script configuration and try again.[/]")
715
726
  return None
716
727
 
717
728
  eval_config = self.confirm_and_finalize_evaluation_config(eval_config)
@@ -791,7 +802,7 @@ def run_onboarding_chatbot(
791
802
  chatbot = Chatbot(project_path, console, run_parser, model)
792
803
  chatbot.start()
793
804
  except Exception as e:
794
- console.print(f"[bold red]An unexpected error occurred in the chatbot: {e}[/]")
805
+ console.print(f"[bold red]An unexpected error occurred: {e}[/]")
795
806
  import traceback
796
807
 
797
808
  traceback.print_exc()
@@ -61,6 +61,12 @@ def configure_run_parser(run_parser: argparse.ArgumentParser) -> None:
61
61
  type=str,
62
62
  help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
63
63
  )
64
+ run_parser.add_argument(
65
+ "--eval-timeout",
66
+ type=int,
67
+ default=None,
68
+ help="Timeout in seconds for each evaluation. No timeout by default. Example: --eval-timeout 3600",
69
+ )
64
70
 
65
71
 
66
72
  def execute_run_command(args: argparse.Namespace) -> None:
@@ -77,6 +83,7 @@ def execute_run_command(args: argparse.Namespace) -> None:
77
83
  log_dir=args.log_dir,
78
84
  additional_instructions=args.additional_instructions,
79
85
  console=console,
86
+ eval_timeout=args.eval_timeout,
80
87
  )
81
88
  exit_code = 0 if success else 1
82
89
  sys.exit(exit_code)
@@ -177,7 +184,9 @@ def main() -> None:
177
184
 
178
185
  project_path = pathlib.Path(filtered_args[0]) if filtered_args else pathlib.Path.cwd()
179
186
  if not project_path.is_dir():
180
- console.print(f"[bold red]Error:[/] Path '{project_path}' is not a valid directory.")
187
+ console.print(
188
+ f"[bold red]Error:[/] The path '{project_path}' is not a valid directory. Please provide a valid directory path."
189
+ )
181
190
  sys.exit(1)
182
191
 
183
192
  # Pass the run_parser and model to the chatbot