Skip to content

Commit c24178a

Browse files
yuxiang-wuyuxiang-wu
authored
Merge pull request #12 from WecoAI/feature/cli-liveness-tracking
feat: Implement CLI Liveness Tracking and Sync with Dev
2 parents 39182a9 + 87ab675 commit c24178a

File tree

5 files changed

+278
-60
lines changed

5 files changed

+278
-60
lines changed

‎pyproject.toml‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,18 @@ build-backend = "setuptools.build_meta"
55

66
[project]
77
name = "weco"
8-
authors = [
9-
{name = "Weco AI Team", email = "contact@weco.ai"},
10-
]
8+
authors = [{ name = "Weco AI Team", email = "contact@weco.ai" }]
119
description = "Documentation for `weco`, a CLI for using Weco AI's code optimizer."
1210
readme = "README.md"
13-
version = "0.2.17"
14-
license = {text = "MIT"}
11+
version = "0.2.18"
12+
license = { text = "MIT" }
1513
requires-python = ">=3.8"
1614
dependencies = ["requests", "rich"]
1715
keywords = ["AI", "Code Optimization", "Code Generation"]
1816
classifiers = [
1917
"Programming Language :: Python :: 3",
2018
"Operating System :: OS Independent",
21-
"License :: OSI Approved :: MIT License"
19+
"License :: OSI Approved :: MIT License",
2220
]
2321

2422
[project.scripts]

‎weco/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import os
2+
import importlib.metadata
23

34
# DO NOT EDIT
4-
__pkg_version__ = "0.2.17"
5+
__pkg_version__ = importlib.metadata.version("weco")
56
__api_version__ = "v1"
67

78
__base_url__ = f"https://api.weco.ai/{__api_version__}"

‎weco/api.py‎

Lines changed: 117 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,20 @@
1-
from typing import Dict, Any
1+
from typing import Dict, Any, Optional
22
import rich
33
import requests
44
from weco import __pkg_version__, __base_url__
55
import sys
6+
from rich.console import Console
67

78

89
def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Console) -> None:
910
"""Extract and display error messages from API responses in a structured format."""
10-
console.print(f"[bold red]{e.response.json()['detail']}[/]")
11-
sys.exit(1)
11+
try:
12+
detail = e.response.json()["detail"]
13+
except (ValueError, KeyError): # Handle cases where response is not JSON or detail key is missing
14+
detail = f"HTTP {e.response.status_code} Error: {e.response.text}"
15+
console.print(f"[bold red]{detail}[/]")
16+
# Avoid exiting here, let the caller decide if the error is fatal
17+
# sys.exit(1)
1218

1319

1420
def start_optimization_session(
@@ -28,25 +34,32 @@ def start_optimization_session(
2834
) -> Dict[str, Any]:
2935
"""Start the optimization session."""
3036
with console.status("[bold green]Starting Optimization..."):
31-
response = requests.post(
32-
f"{__base_url__}/sessions", # Path is relative to base_url
33-
json={
34-
"source_code": source_code,
35-
"additional_instructions": additional_instructions,
36-
"objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
37-
"optimizer": {
38-
"steps": steps,
39-
"code_generator": code_generator_config,
40-
"evaluator": evaluator_config,
41-
"search_policy": search_policy_config,
37+
try:
38+
response = requests.post(
39+
f"{__base_url__}/sessions", # Path is relative to base_url
40+
json={
41+
"source_code": source_code,
42+
"additional_instructions": additional_instructions,
43+
"objective": {"evaluation_command": evaluation_command, "metric_name": metric_name, "maximize": maximize},
44+
"optimizer": {
45+
"steps": steps,
46+
"code_generator": code_generator_config,
47+
"evaluator": evaluator_config,
48+
"search_policy": search_policy_config,
49+
},
50+
"metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
4251
},
43-
"metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
44-
},
45-
headers=auth_headers, # Add headers
46-
timeout=timeout,
47-
)
48-
response.raise_for_status()
49-
return response.json()
52+
headers=auth_headers, # Add headers
53+
timeout=timeout,
54+
)
55+
response.raise_for_status()
56+
return response.json()
57+
except requests.exceptions.HTTPError as e:
58+
handle_api_error(e, console)
59+
sys.exit(1) # Exit if starting session fails
60+
except requests.exceptions.RequestException as e:
61+
console.print(f"[bold red]Network Error starting session: {e}[/]")
62+
sys.exit(1)
5063

5164

5265
def evaluate_feedback_then_suggest_next_solution(
@@ -58,29 +71,92 @@ def evaluate_feedback_then_suggest_next_solution(
5871
timeout: int = 800,
5972
) -> Dict[str, Any]:
6073
"""Evaluate the feedback and suggest the next solution."""
61-
response = requests.post(
62-
f"{__base_url__}/sessions/{session_id}/suggest", # Path is relative to base_url
63-
json={
64-
"execution_output": execution_output,
65-
"additional_instructions": additional_instructions,
66-
"metadata": {**api_keys},
67-
},
68-
headers=auth_headers, # Add headers
69-
timeout=timeout,
70-
)
71-
response.raise_for_status()
72-
return response.json()
74+
try:
75+
response = requests.post(
76+
f"{__base_url__}/sessions/{session_id}/suggest", # Path is relative to base_url
77+
json={
78+
"execution_output": execution_output,
79+
"additional_instructions": additional_instructions,
80+
"metadata": {**api_keys},
81+
},
82+
headers=auth_headers, # Add headers
83+
timeout=timeout,
84+
)
85+
response.raise_for_status()
86+
return response.json()
87+
except requests.exceptions.HTTPError as e:
88+
# Allow caller to handle suggest errors, maybe retry or terminate
89+
handle_api_error(e, Console()) # Use default console if none passed
90+
raise # Re-raise the exception
91+
except requests.exceptions.RequestException as e:
92+
print(f"Network Error during suggest: {e}") # Use print as console might not be available
93+
raise # Re-raise the exception
7394

7495

7596
def get_optimization_session_status(
7697
session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
7798
) -> Dict[str, Any]:
7899
"""Get the current status of the optimization session."""
79-
response = requests.get(
80-
f"{__base_url__}/sessions/{session_id}", # Path is relative to base_url
81-
params={"include_history": include_history},
82-
headers=auth_headers,
83-
timeout=timeout,
84-
)
85-
response.raise_for_status()
86-
return response.json()
100+
try:
101+
response = requests.get(
102+
f"{__base_url__}/sessions/{session_id}", # Path is relative to base_url
103+
params={"include_history": include_history},
104+
headers=auth_headers,
105+
timeout=timeout,
106+
)
107+
response.raise_for_status()
108+
return response.json()
109+
except requests.exceptions.HTTPError as e:
110+
handle_api_error(e, Console()) # Use default console
111+
raise # Re-raise
112+
except requests.exceptions.RequestException as e:
113+
print(f"Network Error getting status: {e}")
114+
raise # Re-raise
115+
116+
117+
def send_heartbeat(
118+
session_id: str,
119+
auth_headers: dict = {},
120+
timeout: int = 10, # Shorter timeout for non-critical heartbeat
121+
) -> bool:
122+
"""Send a heartbeat signal to the backend."""
123+
try:
124+
response = requests.put(f"{__base_url__}/sessions/{session_id}/heartbeat", headers=auth_headers, timeout=timeout)
125+
response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
126+
return True
127+
except requests.exceptions.HTTPError as e:
128+
# Log non-critical errors like 409 Conflict (session not running)
129+
if e.response.status_code == 409:
130+
print(f"Heartbeat ignored: Session {session_id} is not running.", file=sys.stderr)
131+
else:
132+
print(f"Heartbeat failed for session {session_id}: HTTP {e.response.status_code}", file=sys.stderr)
133+
# Don't exit, just report failure
134+
return False
135+
except requests.exceptions.RequestException as e:
136+
# Network errors are also non-fatal for heartbeats
137+
print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
138+
return False
139+
140+
141+
def report_termination(
142+
session_id: str,
143+
status_update: str,
144+
reason: str,
145+
details: Optional[str] = None,
146+
auth_headers: dict = {},
147+
timeout: int = 30, # Reasonably longer timeout for important termination message
148+
) -> bool:
149+
"""Report the termination reason to the backend."""
150+
try:
151+
response = requests.post(
152+
f"{__base_url__}/sessions/{session_id}/terminate",
153+
json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
154+
headers=auth_headers,
155+
timeout=timeout,
156+
)
157+
response.raise_for_status()
158+
return True
159+
except requests.exceptions.RequestException as e:
160+
# Log failure, but don't prevent CLI exit
161+
print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
162+
return False

0 commit comments

Comments
 (0)