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 UTC, datetime 

6from enum import StrEnum 

7from typing import Any 

8 

9from pydantic import BaseModel, Field 

10 

11 

12class AgentStatus(StrEnum): 

13 """Agent run status.""" 

14 

15 IDLE = "idle" 

16 RUNNING = "running" 

17 WAITING_TOOL = "waiting_tool" 

18 WAITING_HUMAN = "waiting_human" 

19 COMPLETED = "completed" 

20 FAILED = "failed" 

21 CANCELLED = "cancelled" 

22 TIMEOUT = "timeout" 

23 

24 

25class AgentRunRequest(BaseModel): 

26 """Request to run an agent.""" 

27 

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

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

30 model: str | None = Field(default=None, description="Override the default model") 

31 max_tokens: int | None = Field( 

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

33 ) 

34 temperature: float | None = 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: dict[str, Any] | None = Field( 

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

43 ) 

44 timeout_seconds: int | None = 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 

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

53 agent_name: str 

54 status: AgentStatus 

55 output: str | None = Field(default=None) 

56 tool_calls: list[dict[str, Any]] = Field(default_factory=list) 

57 usage: dict[str, int] | None = Field(default=None) 

58 duration_ms: float = Field(default=0.0) 

59 error: str | None = Field(default=None) 

60 created_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat()) 

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

62 

63 

64class AgentInfo(BaseModel): 

65 """Static agent information.""" 

66 

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: str | None = None 

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

75 

76 

77class AgentListResponse(BaseModel): 

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

79 

80 agents: list[AgentInfo] = Field(default_factory=list) 

81 total: int = 0