improve how errors are displayed to users

Author: DhruvSrikanth

weco/__init__.py

@@ -1,7 +1,7 @@
 import os
 
 # DO NOT EDIT
-__pkg_version__ = "0.2.15"
+__pkg_version__ = "0.2.16"
 __api_version__ = "v1"
 
 __base_url__ = f"https://api.weco.ai/{__api_version__}"

weco/api.py

@@ -28,33 +28,28 @@ def start_optimization_session(
 ) -> Dict[str, Any]:
     """Start the optimization session."""
     with console.status("[bold green]Starting Optimization..."):
-        try:
-            # __base_url__ already contains /v1
-            response = requests.post(
-                f"{__base_url__}/sessions",  # Path is relative to base_url
-                json={
-                    "source_code": source_code,
-                    "additional_instructions": additional_instructions,
-                    "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
-                    "optimizer": {
-                        "steps": steps,
-                        "code_generator": code_generator_config,
-                        "evaluator": evaluator_config,
-                        "search_policy": search_policy_config,
-                    },
-                    "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
+        response = requests.post(
+            f"{__base_url__}/sessions",  # Path is relative to base_url
+            json={
+                "source_code": source_code,
+                "additional_instructions": additional_instructions,
+                "objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
+                "optimizer": {
+                    "steps": steps,
+                    "code_generator": code_generator_config,
+                    "evaluator": evaluator_config,
+                    "search_policy": search_policy_config,
                 },
-                headers=auth_headers,  # Add headers
-                timeout=timeout,
-            )
-            response.raise_for_status()
-            return response.json()
-        except requests.exceptions.HTTPError as e:
-            handle_api_error(e=e, console=console)
+                "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
+            },
+            headers=auth_headers,  # Add headers
+            timeout=timeout,
+        )
+        response.raise_for_status()
+        return response.json()
 
 
 def evaluate_feedback_then_suggest_next_solution(
-    console: rich.console.Console,
     session_id: str,
     execution_output: str,
     additional_instructions: str = None,
@@ -63,41 +58,29 @@ def evaluate_feedback_then_suggest_next_solution(
     timeout: int = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
-    try:
-        # __base_url__ already contains /v1
-        response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
-            json={
-                "execution_output": execution_output,
-                "additional_instructions": additional_instructions,
-                "metadata": {**api_keys},
-            },
-            headers=auth_headers,  # Add headers
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e=e, console=console)
+    response = requests.post(
+        f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
+        json={
+            "execution_output": execution_output,
+            "additional_instructions": additional_instructions,
+            "metadata": {**api_keys},
+        },
+        headers=auth_headers,  # Add headers
+        timeout=timeout,
+    )
+    response.raise_for_status()
+    return response.json()
 
 
 def get_optimization_session_status(
-    console: rich.console.Console,
-    session_id: str,
-    include_history: bool = False,
-    auth_headers: dict = {},
-    timeout: int = 800,  # Add auth_headers
+    session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
 ) -> Dict[str, Any]:
     """Get the current status of the optimization session."""
-    try:
-        # __base_url__ already contains /v1
-        response = requests.get(
-            f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
-            params={"include_history": include_history},
-            headers=auth_headers,  # Add headers
-            timeout=timeout,
-        )
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as e:
-        handle_api_error(e=e, console=console)
+    response = requests.get(
+        f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
+        params={"include_history": include_history},
+        headers=auth_headers,
+        timeout=timeout,
+    )
+    response.raise_for_status()
+    return response.json()

weco/cli.py

@@ -258,7 +258,7 @@ def main() -> None:
             code_generator_config = {"model": args.model}
             evaluator_config = {
                 "model": args.model,
-                "include_analysis": False,  # NOTE: False for now
+                "include_analysis": True,
             }
             search_policy_config = {
                 "num_drafts": max(1, math.ceil(0.15 * steps)),
@@ -388,7 +388,6 @@ def main() -> None:
 
                     # Send feedback and get next suggestion
                     eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
-                        console=console,
                         session_id=session_id,
                         execution_output=term_out,
                         additional_instructions=current_additional_instructions,  # Pass current instructions
@@ -408,11 +407,7 @@ def main() -> None:
                     # Get the optimization session status for
                     # the best solution, its score, and the history to plot the tree
                     status_response = get_optimization_session_status(
-                        console=console,
-                        session_id=session_id,
-                        include_history=True,
-                        timeout=timeout,
-                        auth_headers=auth_headers,
+                        session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                     )
 
                     # Update the step of the progress bar
@@ -493,7 +488,6 @@ def main() -> None:
 
                 # Ensure we pass evaluation results for the last step's generated solution
                 eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
-                    console=console,
                     session_id=session_id,
                     execution_output=term_out,
                     additional_instructions=current_additional_instructions,
@@ -510,7 +504,7 @@ def main() -> None:
                 # Get the optimization session status for
                 # the best solution, its score, and the history to plot the tree
                 status_response = get_optimization_session_status(
-                    console=console, session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                 )
                 # Build the metric tree
                 tree_panel.build_metric_tree(nodes=status_response["history"])
@@ -575,7 +569,11 @@ def main() -> None:
             console.print(end_optimization_layout)
 
         except Exception as e:
-            console.print(Panel(f"[bold red]Error: {str(e)}", title="[bold red]Error", border_style="red"))
+            try:
+                error_message = e.response.json()["detail"]
+            except Exception:
+                error_message = str(e)
+            console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Error", border_style="red"))
             # Print traceback for debugging
-            console.print_exception(show_locals=True)
+            # console.print_exception(show_locals=False)
             sys.exit(1)

weco/panels.py

@@ -121,6 +121,7 @@ def __init__(
         self.metric = metric
         self.is_buggy = is_buggy
         self.evaluated = True
+        self.name = ""
 
 
 class MetricTree:
@@ -181,16 +182,17 @@ def build_metric_tree(self, nodes: List[dict]):
         nodes.sort(key=lambda x: x["step"])
 
         # Finally build the new tree
-        for node in nodes:
-            self.metric_tree.add_node(
-                Node(
-                    id=node["solution_id"],
-                    parent_id=node["parent_id"],
-                    code=node["code"],
-                    metric=node["metric_value"],
-                    is_buggy=node["is_buggy"],
-                )
+        for i, node in enumerate(nodes):
+            node = Node(
+                id=node["solution_id"],
+                parent_id=node["parent_id"],
+                code=node["code"],
+                metric=node["metric_value"],
+                is_buggy=node["is_buggy"],
             )
+            if i == 0:
+                node.name = "baseline"
+            self.metric_tree.add_node(node)
 
     def set_unevaluated_node(self, node_id: str):
         """Set the unevaluated node."""
@@ -232,12 +234,15 @@ def append_rec(node: Node, tree: Tree):
                     style = None
                     text = f"{node.metric:.3f}"
 
+                # add the node name info
+                text = f"{node.name} {text}".strip()
+
             s = f"[{f'{style} ' if style is not None else ''}{color}]● {text}"
             subtree = tree.add(s)
             for child in node.children:
                 append_rec(child, subtree)
 
-        tree = Tree("🌳")
+        tree = Tree("", hide_root=True)
         for n in self.metric_tree.get_draft_nodes():
             append_rec(n, tree)