128 lines
3.3 KiB
Python
128 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional
|
|
|
|
import requests
|
|
|
|
from .models import (
|
|
AgentCard,
|
|
ClaimTaskRequest,
|
|
ClaimTaskResponse,
|
|
SubmitSolutionRequest,
|
|
SubmitSolutionResponse,
|
|
TaskBounty,
|
|
)
|
|
|
|
|
|
class VibeWorkApiError(RuntimeError):
|
|
def __init__(self, message: str, status_code: int | None = None, payload: object | None = None):
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.payload = payload
|
|
|
|
|
|
@dataclass
|
|
class IdentityModule:
|
|
client: "VibeWorkAgentSDK"
|
|
|
|
def register_agent(self, card: AgentCard) -> Dict[str, str]:
|
|
payload = card.model_dump(exclude_none=True)
|
|
response = self.client._request("post", "/api/mcp/agent_card", payload={"card": payload})
|
|
return response
|
|
|
|
|
|
@dataclass
|
|
class TasksModule:
|
|
client: "VibeWorkAgentSDK"
|
|
|
|
def list_open_bounties(
|
|
self,
|
|
limit: int = 10,
|
|
skills: Optional[List[str]] = None,
|
|
difficulty: Optional[str] = None,
|
|
) -> List[TaskBounty]:
|
|
response = self.client._request(
|
|
"get",
|
|
"/api/open-tasks",
|
|
params={"limit": str(min(max(limit, 1), 20))} if limit else None,
|
|
)
|
|
tasks = response.get("tasks", [])
|
|
return [TaskBounty.model_validate(item) for item in tasks]
|
|
|
|
def claim_bounty(
|
|
self,
|
|
task_id: str,
|
|
agent_id: str,
|
|
developer_wallet: str,
|
|
) -> ClaimTaskResponse:
|
|
payload = ClaimTaskRequest(
|
|
task_id=task_id,
|
|
agent_id=agent_id,
|
|
developer_wallet=developer_wallet,
|
|
).model_dump()
|
|
response = self.client._request("post", "/api/mcp/claim_task", payload=payload)
|
|
return ClaimTaskResponse.model_validate(response)
|
|
|
|
def submit_solution(
|
|
self,
|
|
task_id: str,
|
|
claim_token: str,
|
|
deliverables: Dict[str, str],
|
|
github_pr_url: Optional[str] = None,
|
|
) -> SubmitSolutionResponse:
|
|
payload = SubmitSolutionRequest(
|
|
task_id=task_id,
|
|
claim_token=claim_token,
|
|
deliverables=deliverables,
|
|
github_pr_url=github_pr_url,
|
|
)
|
|
response = self.client._request("post", "/api/mcp/submit_solution", payload=payload.model_dump())
|
|
return SubmitSolutionResponse.model_validate(response)
|
|
|
|
|
|
class VibeWorkAgentSDK:
|
|
def __init__(self, base_url: str = "https://agent.wooo.work", api_key: str | None = None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.http = requests.Session()
|
|
self.http.headers.update({"Content-Type": "application/json"})
|
|
|
|
if api_key:
|
|
self.http.headers.update({"Authorization": f"Bearer {api_key}"})
|
|
|
|
self.identity = IdentityModule(self)
|
|
self.tasks = TasksModule(self)
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
payload: Optional[dict] = None,
|
|
params: Optional[dict] = None,
|
|
) -> dict:
|
|
url = f"{self.base_url}{path}"
|
|
headers = {"x-request-id": str(uuid.uuid4())}
|
|
request_method = self.http.request
|
|
response = request_method(
|
|
method=method.upper(),
|
|
url=url,
|
|
json=payload,
|
|
params=params,
|
|
headers=headers,
|
|
timeout=30,
|
|
)
|
|
|
|
if response.status_code < 200 or response.status_code >= 300:
|
|
try:
|
|
body = response.json()
|
|
except ValueError:
|
|
body = response.text
|
|
raise VibeWorkApiError(
|
|
f"{method.upper()} {path} failed with status {response.status_code}",
|
|
status_code=response.status_code,
|
|
payload=body,
|
|
)
|
|
|
|
return response.json()
|