File size: 8,011 Bytes
68b32d7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
"""
Agent Status API Endpoints
Provides status monitoring for AI agents and task execution
"""

import asyncio
from datetime import datetime
import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional
import uuid
from fastapi import BackgroundTasks
from pydantic import BaseModel

from core.base_routes import BaseAPIRouter

logger = logging.getLogger(__name__)

router = BaseAPIRouter(prefix="/api/agent-status", tags=["Agent Status"])

# In-memory storage for agent status (for MVP)
AGENT_STATUS_FILE = Path(__file__).parent.parent / "agent_status.json"

class AgentTask(BaseModel):
    task_id: str
    agent_id: str
    status: str  # pending, running, completed, failed, cancelled
    progress: float = 0.0  # 0.0 to 1.0
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error_message: Optional[str] = None
    result: Optional[Dict[str, Any]] = None
    metadata: Dict[str, Any] = {}

class AgentInfo(BaseModel):
    agent_id: str
    name: str
    type: str
    status: str  # idle, busy, offline
    last_active: Optional[datetime] = None
    current_task: Optional[str] = None
    capabilities: List[str] = []
    health_score: float = 1.0  # 0.0 to 1.0

def load_agent_status() -> Dict[str, Any]:
    """Load agent status from file"""
    if not AGENT_STATUS_FILE.exists():
        return {"agents": {}, "tasks": {}}

    try:
        with open(AGENT_STATUS_FILE, 'r') as f:
            return json.load(f)
    except Exception as e:
        return {"agents": {}, "tasks": {}}

def save_agent_status(data: Dict[str, Any]):
    """Save agent status to file"""
    try:
        with open(AGENT_STATUS_FILE, 'w') as f:
            json.dump(data, f, indent=2, default=str)
    except Exception as e:
        logger.error(f"Error saving agent status: {e}")

@router.get("/agent/status/{task_id}", response_model=AgentTask)
async def get_agent_status(task_id: str):
    """Get status of a specific agent task"""
    data = load_agent_status()

    if task_id not in data.get("tasks", {}):
        # Return a default status for unknown tasks
        return AgentTask(
            task_id=task_id,
            agent_id="unknown",
            status="not_found",
            error_message="Task not found"
        )

    task_data = data["tasks"][task_id]
    return AgentTask(**task_data)

@router.get("/agent/status", response_model=List[AgentTask])
async def get_all_agent_tasks():
    """Get status of all agent tasks"""
    data = load_agent_status()

    tasks = []
    for task_data in data.get("tasks", {}).values():
        tasks.append(AgentTask(**task_data))

    return tasks

@router.get("/agents", response_model=List[AgentInfo])
async def get_all_agents():
    """Get information about all agents"""
    data = load_agent_status()

    agents = []
    for agent_data in data.get("agents", {}).values():
        agents.append(AgentInfo(**agent_data))

    return agents

@router.get("/agents/{agent_id}", response_model=AgentInfo)
async def get_agent_info(agent_id: str):
    """Get information about a specific agent"""
    data = load_agent_status()

    if agent_id not in data.get("agents", {}):
        # Create a default agent if not found
        default_agent = AgentInfo(
            agent_id=agent_id,
            name=f"Agent {agent_id}",
            type="general",
            status="idle",
            last_active=datetime.now(),
            capabilities=["text_processing", "analysis"]
        )

        # Save the default agent
        data.setdefault("agents", {})[agent_id] = default_agent.model_dump()
        save_agent_status(data)

        return default_agent

    agent_data = data["agents"][agent_id]
    return AgentInfo(**agent_data)

@router.post("/agent/{agent_id}/heartbeat")
async def agent_heartbeat(agent_id: str, status: Dict[str, Any]):
    """Update agent heartbeat and status"""
    data = load_agent_status()

    # Update or create agent info
    agent_info = {
        "agent_id": agent_id,
        "name": status.get("name", f"Agent {agent_id}"),
        "type": status.get("type", "general"),
        "status": status.get("status", "idle"),
        "last_active": datetime.now(),
        "current_task": status.get("current_task"),
        "capabilities": status.get("capabilities", []),
        "health_score": status.get("health_score", 1.0)
    }

    data.setdefault("agents", {})[agent_id] = agent_info
    save_agent_status(data)

    return router.success_response(
        data={"timestamp": datetime.now()},
        message="Agent heartbeat updated successfully"
    )

@router.post("/agent/task/{task_id}/update")
async def update_task_status(task_id: str, update: Dict[str, Any]):
    """Update status of a specific task"""
    data = load_agent_status()

    if task_id not in data.get("tasks", {}):
        raise router.not_found_error("Task", task_id)

    # Update task fields
    task_data = data["tasks"][task_id]

    if "status" in update:
        task_data["status"] = update["status"]
        if update["status"] == "running" and not task_data.get("started_at"):
            task_data["started_at"] = datetime.now().isoformat()
        elif update["status"] in ["completed", "failed", "cancelled"]:
            task_data["completed_at"] = datetime.now().isoformat()

    if "progress" in update:
        task_data["progress"] = update["progress"]

    if "error_message" in update:
        task_data["error_message"] = update["error_message"]

    if "result" in update:
        task_data["result"] = update["result"]

    data["tasks"][task_id] = task_data
    save_agent_status(data)

    return router.success_response(message="Task status updated successfully")

@router.post("/agent/task")
async def create_task(task: AgentTask):
    """Create a new agent task"""
    data = load_agent_status()

    # Set timestamps
    if not task.started_at and task.status == "running":
        task.started_at = datetime.now()

    # Convert to dict and save
    task_dict = task.model_dump()
    task_dict["started_at"] = task_dict["started_at"].isoformat() if task_dict["started_at"] else None
    task_dict["completed_at"] = task_dict["completed_at"].isoformat() if task_dict["completed_at"] else None

    data.setdefault("tasks", {})[task.task_id] = task_dict
    save_agent_status(data)

    return router.success_response(
        data={"task_id": task.task_id},
        message="Task created successfully"
    )

@router.delete("/agent/task/{task_id}")
async def delete_task(task_id: str):
    """Delete a task"""
    data = load_agent_status()

    if task_id in data.get("tasks", {}):
        del data["tasks"][task_id]
        save_agent_status(data)
        return router.success_response(message="Task deleted successfully")
    else:
        raise router.not_found_error("Task", task_id)

@router.get("/agent/metrics")
async def get_agent_metrics():
    """Get agent performance metrics"""
    data = load_agent_status()

    total_agents = len(data.get("agents", {}))
    active_agents = len([
        a for a in data.get("agents", {}).values()
        if a.get("status") in ["running", "busy"]
    ])

    total_tasks = len(data.get("tasks", {}))
    completed_tasks = len([
        t for t in data.get("tasks", {}).values()
        if t.get("status") == "completed"
    ])
    failed_tasks = len([
        t for t in data.get("tasks", {}).values()
        if t.get("status") == "failed"
    ])

    return router.success_response(
        data={
            "agents": {
                "total": total_agents,
                "active": active_agents,
                "idle": total_agents - active_agents
            },
            "tasks": {
                "total": total_tasks,
                "completed": completed_tasks,
                "failed": failed_tasks,
                "pending": total_tasks - completed_tasks - failed_tasks
            },
            "success_rate": completed_tasks / max(total_tasks, 1)
        }
    )