Skip to content

Commit f095b25

Browse files
polebugpolebug
committed
feat(router): implement agent stop api
1 parent bc49cf0 commit f095b25

File tree

9 files changed

+170
-75
lines changed

9 files changed

+170
-75
lines changed

‎openagent/router/routes/agent.py‎

Lines changed: 67 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
PublicAgentResponse,
2525
ResponseModel,
2626
)
27-
from openagent.tools import BaseTool, ToolConfig, ToolParameters
27+
from openagent.tools import BaseTool, ToolConfig
2828

2929
auth_handler = Auth()
3030

@@ -175,7 +175,7 @@ def list_agents(
175175
)
176176
def get_agent(
177177
agent_id: int,
178-
wallet_address: str = Depends(auth_handler.auth_wrapper),
178+
wallet_address: str | None = Depends(auth_handler.optional_auth_wrapper),
179179
db: Session = Depends(get_db),
180180
) -> ResponseModel[AgentResponse] | APIExceptionResponse:
181181
try:
@@ -186,14 +186,16 @@ def get_agent(
186186
error=f"Agent with ID {agent_id} not found",
187187
)
188188

189-
if agent.wallet_address.lower() != wallet_address.lower():
190-
return APIExceptionResponse(
191-
status_code=status.HTTP_403_FORBIDDEN,
192-
error="Not authorized to query this agent",
193-
)
189+
# return full agent info if authenticated and wallet addresses match
190+
if wallet_address and agent.wallet_address.lower() == wallet_address.lower():
191+
response_data = AgentResponse.model_validate(agent)
192+
else:
193+
# return public info for unauthenticated users or non-owners
194+
response_data = PublicAgentResponse.model_validate(agent)
195+
194196
return ResponseModel(
195197
code=status.HTTP_200_OK,
196-
data=AgentResponse.model_validate(agent),
198+
data=AgentResponse.model_validate(response_data),
197199
message="Agent retrieved successfully",
198200
)
199201
except Exception as error:
@@ -355,6 +357,59 @@ def run_agent(
355357
)
356358

357359

360+
@router.post(
361+
"/{agent_id}/stop",
362+
response_model=ResponseModel[AgentResponse],
363+
summary="Stop an agent",
364+
description="Stop an agent by setting its status to inactive",
365+
responses={
366+
200: {"description": "Successfully stopped agent"},
367+
403: {"description": "Not authorized to stop this agent"},
368+
404: {"description": "Agent not found"},
369+
500: {"description": "Internal server error"},
370+
},
371+
)
372+
def stop_agent(
373+
agent_id: int,
374+
wallet_address: str = Depends(auth_handler.auth_wrapper),
375+
db: Session = Depends(get_db),
376+
) -> ResponseModel[AgentResponse] | APIExceptionResponse:
377+
try:
378+
# get agent
379+
agent = db.query(Agent).filter(Agent.id == agent_id).first()
380+
if not agent:
381+
return APIExceptionResponse(
382+
status_code=status.HTTP_404_NOT_FOUND,
383+
error=f"Agent with ID {agent_id} not found",
384+
)
385+
386+
# check if the user is authorized to stop this agent
387+
if agent.wallet_address.lower() != wallet_address.lower():
388+
return APIExceptionResponse(
389+
status_code=status.HTTP_403_FORBIDDEN,
390+
error="Not authorized to stop this agent",
391+
)
392+
393+
# update the agent status to inactive
394+
agent.status = AgentStatus.INACTIVE
395+
db.commit()
396+
db.refresh(agent)
397+
398+
# TODO: stop any running agent processes
399+
400+
return ResponseModel(
401+
code=status.HTTP_200_OK,
402+
data=AgentResponse.model_validate(agent),
403+
message="Agent stopped successfully",
404+
)
405+
except Exception as error:
406+
db.rollback()
407+
return APIExceptionResponse(
408+
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
409+
error=error,
410+
)
411+
412+
358413
@router.post(
359414
"/{agent_id}/execute/{tool_name}",
360415
response_model=ResponseModel[dict[str, Any]],
@@ -473,9 +528,11 @@ def build_model(model: Model) -> AI_Model:
473528
raise ValueError(f"Unsupported model: {model}")
474529

475530

476-
def initialize_tool_executor(agent: Agent, tool: Tool, model: Model, tool_config: ToolConfig) -> BaseTool:
531+
def initialize_tool_executor(
532+
agent: Agent, tool: Tool, model: Model, tool_config: ToolConfig
533+
) -> BaseTool:
477534
model_instance = build_model(model)
478-
535+
479536
return get_tool_executor(agent, tool, model_instance, tool_config)
480537

481538

‎openagent/router/routes/model.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from fastapi import APIRouter, Depends, status
1+
from fastapi import APIRouter, Depends, status, Query
22
from sqlalchemy.orm import Session
33

44
from openagent.database import get_db
@@ -24,11 +24,21 @@
2424
},
2525
)
2626
def list_models(
27-
page: int = 0, limit: int = 10, db: Session = Depends(get_db)
27+
page: int = 0,
28+
limit: int = 10,
29+
ids: list[int] | None = Query(default=None),
30+
db: Session = Depends(get_db),
2831
) -> ResponseModel[ModelListResponse] | APIExceptionResponse:
2932
try:
30-
total = db.query(Model).count()
31-
models = db.query(Model).offset(page * limit).limit(limit).all()
33+
query = db.query(Model)
34+
35+
# Add filter for model ids if provided
36+
if ids:
37+
query = query.filter(Model.id.in_(ids))
38+
39+
total = query.count()
40+
models = query.offset(page * limit).limit(limit).all()
41+
3242
return ResponseModel(
3343
code=status.HTTP_200_OK,
3444
data=ModelListResponse(

‎openagent/router/routes/models/auth.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
1010
from web3 import Web3
1111

12-
security = HTTPBearer()
12+
security = HTTPBearer(auto_error=False)
1313

1414

1515
class Auth:
@@ -71,7 +71,25 @@ def verify_wallet_signature(
7171
)
7272

7373
def auth_wrapper(
74-
self, auth: HTTPAuthorizationCredentials = Security(security)
74+
self, auth: HTTPAuthorizationCredentials | None = Security(security)
7575
) -> str:
76+
if not auth:
77+
raise HTTPException(
78+
status_code=status.HTTP_401_UNAUTHORIZED,
79+
detail="No authorization token provided",
80+
headers={"WWW-Authenticate": "Bearer"},
81+
)
82+
7683
payload = self.decode_token(auth.credentials)
7784
return payload["wallet_address"]
85+
86+
def optional_auth_wrapper(
87+
self, auth: HTTPAuthorizationCredentials | None = Security(security)
88+
) -> str | None:
89+
if not auth:
90+
return None
91+
try:
92+
payload = self.decode_token(auth.credentials)
93+
return payload["wallet_address"]
94+
except HTTPException:
95+
return None

‎openagent/router/routes/models/response.py‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,16 @@ class AgentResponse(BaseModel):
3838
created_at: datetime
3939
updated_at: datetime
4040

41+
4142
class PublicToolConfigResponse(BaseModel):
4243
model_config = ConfigDict(from_attributes=True)
4344

4445
name: str
4546
description: str | None = None
4647
tool_id: int
4748
model_id: int
48-
49+
50+
4951
class PublicAgentResponse(BaseModel):
5052
model_config = ConfigDict(from_attributes=True)
5153

@@ -70,18 +72,19 @@ class PublicAgentResponse(BaseModel):
7072

7173
@classmethod
7274
def from_orm(cls, obj):
73-
if hasattr(obj, 'tool_configs') and obj.tool_configs:
75+
if hasattr(obj, "tool_configs") and obj.tool_configs:
7476
obj.tool_configs = [
7577
PublicToolConfigResponse(
76-
name=tc['name'],
77-
description=tc.get('description'),
78-
tool_id=tc['tool_id'],
79-
model_id=tc['model_id']
78+
name=tc["name"],
79+
description=tc.get("description"),
80+
tool_id=tc["tool_id"],
81+
model_id=tc["model_id"],
8082
)
8183
for tc in obj.tool_configs
8284
]
8385
return super().from_orm(obj)
8486

87+
8588
class AgentListResponse(BaseModel):
8689
model_config = ConfigDict(from_attributes=True)
8790

@@ -91,7 +94,7 @@ class AgentListResponse(BaseModel):
9194

9295
class ModelResponse(BaseModel):
9396
model_config = ConfigDict(from_attributes=True)
94-
97+
9598
id: int
9699
name: str
97100
description: str | None = None
@@ -106,7 +109,7 @@ class ModelListResponse(BaseModel):
106109

107110
class ToolResponse(BaseModel):
108111
model_config = ConfigDict(from_attributes=True)
109-
112+
110113
id: int
111114
name: str
112115
description: str | None = None

‎openagent/router/routes/tool.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from fastapi import APIRouter, Depends, status
1+
from fastapi import APIRouter, Depends, status, Query
22
from sqlalchemy.orm import Session
33

44
from openagent.database import get_db
@@ -24,11 +24,21 @@
2424
},
2525
)
2626
def list_tools(
27-
page: int = 0, limit: int = 10, db: Session = Depends(get_db)
27+
page: int = 0,
28+
limit: int = 10,
29+
ids: list[int] | None = Query(default=None),
30+
db: Session = Depends(get_db),
2831
) -> ResponseModel[ToolListResponse] | APIExceptionResponse:
2932
try:
30-
total = db.query(Tool).count()
31-
tools = db.query(Tool).offset(page * limit).limit(limit).all()
33+
query = db.query(Tool)
34+
35+
# Add filter for tool ids if provided
36+
if ids:
37+
query = query.filter(Tool.id.in_(ids))
38+
39+
total = query.count()
40+
tools = query.offset(page * limit).limit(limit).all()
41+
3242
return ResponseModel(
3343
code=status.HTTP_200_OK,
3444
data=ToolListResponse(

‎openagent/tools/__init__.py‎

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -5,41 +5,40 @@
55
from phi.tools import Toolkit
66
from pydantic import BaseModel, ConfigDict, model_validator
77

8-
8+
99
class TriggerType(Enum):
1010
SCHEDULED = "scheduled"
1111
AUTO = "auto"
1212
Manual = "manual"
13-
13+
1414
def __str__(self) -> str:
1515
return self.value
1616

17+
1718
class ToolParameters(BaseModel):
1819
model_config = ConfigDict(
1920
from_attributes=True,
2021
arbitrary_types_allowed=True,
21-
json_encoders={TriggerType: lambda v: v.value}
22+
json_encoders={TriggerType: lambda v: v.value},
2223
)
23-
24+
2425
trigger_type: TriggerType
2526
schedule: str | None = None # cron, such as "0 */2 * * *"
2627
config: dict | None = None
2728

2829
def validate_schedule(self):
2930
if self.trigger_type == TriggerType.SCHEDULED and not self.schedule:
3031
raise ValueError("Schedule must be set when trigger_type is SCHEDULED")
31-
32+
3233
def model_dump(self, *args, **kwargs) -> dict:
3334
# Add custom serialization for TriggerType
3435
data = super().model_dump(*args, **kwargs)
35-
data['trigger_type'] = self.trigger_type.value
36+
data["trigger_type"] = self.trigger_type.value
3637
return data
3738

39+
3840
class ToolConfig(BaseModel):
39-
model_config = ConfigDict(
40-
from_attributes=True,
41-
arbitrary_types_allowed=True
42-
)
41+
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)
4342

4443
name: str
4544
description: str | None = None
@@ -51,33 +50,34 @@ def validate_parameters(self):
5150
if self.parameters:
5251
self.parameters.validate_schedule()
5352

54-
5553
def model_dump(self, *args, **kwargs) -> dict:
5654
data = {
5755
"name": self.name,
5856
"description": self.description,
5957
"tool_id": self.tool_id,
6058
"model_id": self.model_id,
6159
}
62-
60+
6361
if self.parameters:
6462
data["parameters"] = self.parameters.model_dump()
65-
63+
6664
return data
6765

66+
6867
class TwitterToolParameters(ToolParameters):
69-
model_config = ConfigDict(
70-
from_attributes=True,
71-
arbitrary_types_allowed=True
72-
)
68+
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)
7369

74-
@model_validator(mode='before')
70+
@model_validator(mode="before")
7571
@classmethod
7672
def validate_twitter_config(cls, data: Dict) -> Dict:
77-
if isinstance(data, dict) and "config" in data and isinstance(data["config"], dict):
73+
if (
74+
isinstance(data, dict)
75+
and "config" in data
76+
and isinstance(data["config"], dict)
77+
):
7878
data["config"] = {
7979
"access_token": data["config"].get("access_token"),
80-
"access_token_secret": data["config"].get("access_token_secret")
80+
"access_token_secret": data["config"].get("access_token_secret"),
8181
}
8282
return data
8383

@@ -113,10 +113,11 @@ def validate_params(self, params: dict[str, Any]) -> tuple[bool, str]:
113113
"""
114114
pass
115115

116+
116117
__all__ = [
117118
"BaseTool",
118119
"ToolConfig",
119120
"ToolParameters",
120121
"TwitterToolParameters",
121122
"TriggerType",
122-
]
123+
]

0 commit comments

Comments
 (0)