Coverage for agentos/models/agent.py: 100%

1"""AgentOS Agent Models — request/response types for agent lifecycle.""" 

2 

3from __future__ import annotations 

4 

5from datetime import datetime, timezone 

6from enum import Enum 

7from typing import Any, Dict, List, Optional 

8 

9from pydantic import BaseModel, Field 

10 

11 

12class AgentStatus(str, Enum): 

13 """Agent run status.""" 

14 IDLE = "idle" 

15 RUNNING = "running" 

16 WAITING_TOOL = "waiting_tool" 

17 WAITING_HUMAN = "waiting_human" 

18 COMPLETED = "completed" 

19 FAILED = "failed" 

20 CANCELLED = "cancelled" 

21 TIMEOUT = "timeout" 

22 

23 

24class AgentRunRequest(BaseModel): 

25 """Request to run an agent.""" 

26 agent_name: str = Field(description="Agent identifier") 

27 input: str = Field(description="User input/message to the agent") 

28 model: Optional[str] = Field( 

29 default=None, description="Override the default model" 

30 ) 

31 max_tokens: Optional[int] = Field( 

32 default=None, ge=1, le=128000, description="Max tokens for the response" 

33 ) 

34 temperature: Optional[float] = Field( 

35 default=None, ge=0.0, le=2.0, description="Sampling temperature" 

36 ) 

37 stream: bool = Field(default=False, description="Enable SSE streaming") 

38 metadata: Dict[str, Any] = Field( 

39 default_factory=dict, description="Arbitrary metadata for tracing" 

40 ) 

41 context: Optional[Dict[str, Any]] = Field( 

42 default=None, description="Additional context injected into agent" 

43 ) 

44 timeout_seconds: Optional[int] = Field( 

45 default=None, ge=1, le=3600, description="Max execution time in seconds" 

46 ) 

47 

48 

49class AgentRunResponse(BaseModel): 

50 """Response from an agent run.""" 

51 run_id: str = Field(description="Unique run identifier") 

52 agent_name: str 

53 status: AgentStatus 

54 output: Optional[str] = Field(default=None) 

55 tool_calls: List[Dict[str, Any]] = Field(default_factory=list) 

56 usage: Optional[Dict[str, int]] = Field(default=None) 

57 duration_ms: float = Field(default=0.0) 

58 error: Optional[str] = Field(default=None) 

59 created_at: str = Field( 

60 default_factory=lambda: datetime.now(timezone.utc).isoformat() 

61 ) 

62 metadata: Dict[str, Any] = Field(default_factory=dict) 

63 

64 

65class AgentInfo(BaseModel): 

66 """Static agent information.""" 

67 name: str 

68 description: str = "" 

69 model: str = "" 

70 version: str = "1.0.0" 

71 tools: List[str] = Field(default_factory=list) 

72 tags: List[str] = Field(default_factory=list) 

73 created_at: Optional[str] = None 

74 metadata: Dict[str, Any] = Field(default_factory=dict) 

75 

76 

77class AgentListResponse(BaseModel): 

78 """List of registered agents.""" 

79 agents: List[AgentInfo] = Field(default_factory=list) 

80 total: int = 0