860395cd58
Endpoints:
- GET /api/v1/usage - Overview with totals, per-agent, per-model breakdown
- GET /api/v1/usage/threads - List all thread budgets sorted by usage
- GET /api/v1/usage/threads/{id} - Single thread budget details
- GET /api/v1/usage/agents/{id} - Usage totals for specific agent
- GET /api/v1/usage/models/{model} - Usage totals for specific model
- POST /api/v1/usage/reset - Reset all usage tracking
Models:
- UsageTotals, UsageOverview, UsageResponse
- ThreadBudgetInfo, ThreadBudgetListResponse
- AgentUsageInfo, ModelUsageInfo
Also adds has_budget() method to ThreadBudgetRegistry for checking
if a thread exists without auto-creating it.
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
370 lines
10 KiB
Python
370 lines
10 KiB
Python
"""
|
|
models.py — Pydantic models for AgentServer API.
|
|
|
|
These models define the JSON structure for API responses.
|
|
Uses camelCase for JSON keys (JavaScript convention).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
|
|
def to_camel(string: str) -> str:
|
|
"""Convert snake_case to camelCase."""
|
|
components = string.split("_")
|
|
return components[0] + "".join(x.title() for x in components[1:])
|
|
|
|
|
|
class CamelModel(BaseModel):
|
|
"""Base model with camelCase JSON serialization."""
|
|
|
|
model_config = ConfigDict(
|
|
alias_generator=to_camel,
|
|
populate_by_name=True,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Enums
|
|
# =============================================================================
|
|
|
|
|
|
class AgentState(str, Enum):
|
|
"""Agent processing state."""
|
|
|
|
IDLE = "idle"
|
|
PROCESSING = "processing"
|
|
WAITING = "waiting"
|
|
ERROR = "error"
|
|
PAUSED = "paused"
|
|
|
|
|
|
class ThreadStatus(str, Enum):
|
|
"""Thread lifecycle status."""
|
|
|
|
ACTIVE = "active"
|
|
COMPLETED = "completed"
|
|
ERROR = "error"
|
|
KILLED = "killed"
|
|
|
|
|
|
class OrganismStatus(str, Enum):
|
|
"""Organism running status."""
|
|
|
|
STARTING = "starting"
|
|
RUNNING = "running"
|
|
STOPPING = "stopping"
|
|
STOPPED = "stopped"
|
|
|
|
|
|
# =============================================================================
|
|
# API Response Models
|
|
# =============================================================================
|
|
|
|
|
|
class AgentInfo(CamelModel):
|
|
"""Agent information for API response."""
|
|
|
|
name: str
|
|
description: str
|
|
is_agent: bool = Field(alias="isAgent")
|
|
peers: List[str] = Field(default_factory=list)
|
|
payload_class: str = Field(alias="payloadClass")
|
|
state: AgentState = AgentState.IDLE
|
|
current_thread: Optional[str] = Field(None, alias="currentThread")
|
|
queue_depth: int = Field(0, alias="queueDepth")
|
|
last_activity: Optional[datetime] = Field(None, alias="lastActivity")
|
|
message_count: int = Field(0, alias="messageCount")
|
|
|
|
|
|
class MessageInfo(CamelModel):
|
|
"""Message information for API response."""
|
|
|
|
id: str
|
|
thread_id: str = Field(alias="threadId")
|
|
from_id: str = Field(alias="from")
|
|
to_id: str = Field(alias="to")
|
|
payload_type: str = Field(alias="payloadType")
|
|
payload: Dict[str, Any] = Field(default_factory=dict)
|
|
timestamp: datetime
|
|
slot_index: Optional[int] = Field(None, alias="slotIndex")
|
|
|
|
|
|
class ThreadInfo(CamelModel):
|
|
"""Thread information for API response."""
|
|
|
|
id: str
|
|
status: ThreadStatus = ThreadStatus.ACTIVE
|
|
participants: List[str] = Field(default_factory=list)
|
|
message_count: int = Field(0, alias="messageCount")
|
|
created_at: datetime = Field(alias="createdAt")
|
|
last_activity: Optional[datetime] = Field(None, alias="lastActivity")
|
|
error: Optional[str] = None
|
|
|
|
|
|
class OrganismInfo(CamelModel):
|
|
"""Organism overview for API response."""
|
|
|
|
name: str
|
|
status: OrganismStatus = OrganismStatus.RUNNING
|
|
uptime_seconds: float = Field(0.0, alias="uptimeSeconds")
|
|
agent_count: int = Field(0, alias="agentCount")
|
|
active_threads: int = Field(0, alias="activeThreads")
|
|
total_messages: int = Field(0, alias="totalMessages")
|
|
identity_configured: bool = Field(False, alias="identityConfigured")
|
|
|
|
|
|
class OrganismConfig(CamelModel):
|
|
"""Sanitized organism configuration for API response."""
|
|
|
|
name: str
|
|
port: int = 8765
|
|
thread_scheduling: str = Field("breadth-first", alias="threadScheduling")
|
|
max_concurrent_pipelines: int = Field(50, alias="maxConcurrentPipelines")
|
|
max_concurrent_handlers: int = Field(20, alias="maxConcurrentHandlers")
|
|
listeners: List[str] = Field(default_factory=list)
|
|
# Note: Secrets like API keys are never exposed
|
|
|
|
|
|
# =============================================================================
|
|
# Request Models
|
|
# =============================================================================
|
|
|
|
|
|
class InjectRequest(CamelModel):
|
|
"""Request body for POST /inject."""
|
|
|
|
to: str
|
|
payload: Dict[str, Any]
|
|
thread_id: Optional[str] = Field(None, alias="threadId")
|
|
|
|
|
|
class InjectResponse(CamelModel):
|
|
"""Response body for POST /inject."""
|
|
|
|
thread_id: str = Field(alias="threadId")
|
|
message_id: str = Field(alias="messageId")
|
|
|
|
|
|
class SubscribeRequest(CamelModel):
|
|
"""WebSocket subscription filter."""
|
|
|
|
threads: Optional[List[str]] = None
|
|
agents: Optional[List[str]] = None
|
|
payload_types: Optional[List[str]] = Field(None, alias="payloadTypes")
|
|
events: Optional[List[str]] = None
|
|
|
|
|
|
# =============================================================================
|
|
# WebSocket Event Models
|
|
# =============================================================================
|
|
|
|
|
|
class WSEvent(CamelModel):
|
|
"""Base WebSocket event."""
|
|
|
|
event: str
|
|
|
|
|
|
class WSConnectedEvent(WSEvent):
|
|
"""Sent on WebSocket connection with full state snapshot."""
|
|
|
|
event: str = "connected"
|
|
organism: OrganismInfo
|
|
agents: List[AgentInfo]
|
|
threads: List[ThreadInfo]
|
|
|
|
|
|
class WSAgentStateEvent(WSEvent):
|
|
"""Agent state changed."""
|
|
|
|
event: str = "agent_state"
|
|
agent: str
|
|
state: AgentState
|
|
current_thread: Optional[str] = Field(None, alias="currentThread")
|
|
|
|
|
|
class WSMessageEvent(WSEvent):
|
|
"""New message in the system."""
|
|
|
|
event: str = "message"
|
|
message: MessageInfo
|
|
|
|
|
|
class WSThreadCreatedEvent(WSEvent):
|
|
"""New thread started."""
|
|
|
|
event: str = "thread_created"
|
|
thread: ThreadInfo
|
|
|
|
|
|
class WSThreadUpdatedEvent(WSEvent):
|
|
"""Thread status changed."""
|
|
|
|
event: str = "thread_updated"
|
|
thread_id: str = Field(alias="threadId")
|
|
status: ThreadStatus
|
|
message_count: int = Field(alias="messageCount")
|
|
|
|
|
|
class WSErrorEvent(WSEvent):
|
|
"""Error occurred."""
|
|
|
|
event: str = "error"
|
|
thread_id: Optional[str] = Field(None, alias="threadId")
|
|
agent: Optional[str] = None
|
|
error: str
|
|
timestamp: datetime
|
|
|
|
|
|
# =============================================================================
|
|
# List Response Models
|
|
# =============================================================================
|
|
|
|
|
|
class AgentListResponse(CamelModel):
|
|
"""Response for GET /agents."""
|
|
|
|
agents: List[AgentInfo]
|
|
count: int
|
|
|
|
|
|
class ThreadListResponse(CamelModel):
|
|
"""Response for GET /threads."""
|
|
|
|
threads: List[ThreadInfo]
|
|
count: int
|
|
total: int
|
|
offset: int
|
|
limit: int
|
|
|
|
|
|
class MessageListResponse(CamelModel):
|
|
"""Response for GET /messages or /threads/{id}/messages."""
|
|
|
|
messages: List[MessageInfo]
|
|
count: int
|
|
total: int
|
|
offset: int
|
|
limit: int
|
|
|
|
|
|
class ErrorResponse(CamelModel):
|
|
"""Error response."""
|
|
|
|
error: str
|
|
detail: Optional[str] = None
|
|
|
|
|
|
# =============================================================================
|
|
# Capability/Introspection Models
|
|
# =============================================================================
|
|
|
|
|
|
class CapabilityInfo(CamelModel):
|
|
"""Basic capability info for listing."""
|
|
|
|
name: str
|
|
description: str
|
|
is_agent: bool = Field(alias="isAgent")
|
|
peers: List[str] = Field(default_factory=list)
|
|
root_tag: str = Field(alias="rootTag")
|
|
|
|
|
|
class CapabilityDetail(CamelModel):
|
|
"""Detailed capability info including schema."""
|
|
|
|
name: str
|
|
description: str
|
|
is_agent: bool = Field(alias="isAgent")
|
|
peers: List[str] = Field(default_factory=list)
|
|
root_tag: str = Field(alias="rootTag")
|
|
payload_class: str = Field(alias="payloadClass")
|
|
schema_xsd: Optional[str] = Field(None, alias="schemaXsd")
|
|
example_xml: Optional[str] = Field(None, alias="exampleXml")
|
|
|
|
|
|
class CapabilityListResponse(CamelModel):
|
|
"""Response for GET /capabilities."""
|
|
|
|
capabilities: List[CapabilityInfo]
|
|
count: int
|
|
|
|
|
|
# =============================================================================
|
|
# Usage/Gas Tracking Models
|
|
# =============================================================================
|
|
|
|
|
|
class UsageTotals(CamelModel):
|
|
"""Aggregate usage statistics."""
|
|
|
|
total_tokens: int = Field(0, alias="totalTokens")
|
|
prompt_tokens: int = Field(0, alias="promptTokens")
|
|
completion_tokens: int = Field(0, alias="completionTokens")
|
|
request_count: int = Field(0, alias="requestCount")
|
|
total_cost: float = Field(0.0, alias="totalCost")
|
|
avg_latency_ms: float = Field(0.0, alias="avgLatencyMs")
|
|
|
|
|
|
class ThreadBudgetInfo(CamelModel):
|
|
"""Token budget info for a thread."""
|
|
|
|
thread_id: str = Field(alias="threadId")
|
|
max_tokens: int = Field(alias="maxTokens")
|
|
prompt_tokens: int = Field(alias="promptTokens")
|
|
completion_tokens: int = Field(alias="completionTokens")
|
|
total_tokens: int = Field(alias="totalTokens")
|
|
remaining: int
|
|
percent_used: float = Field(alias="percentUsed")
|
|
is_exhausted: bool = Field(alias="isExhausted")
|
|
|
|
|
|
class AgentUsageInfo(CamelModel):
|
|
"""Usage info for a specific agent."""
|
|
|
|
agent_id: str = Field(alias="agentId")
|
|
total_tokens: int = Field(0, alias="totalTokens")
|
|
prompt_tokens: int = Field(0, alias="promptTokens")
|
|
completion_tokens: int = Field(0, alias="completionTokens")
|
|
request_count: int = Field(0, alias="requestCount")
|
|
total_cost: float = Field(0.0, alias="totalCost")
|
|
|
|
|
|
class ModelUsageInfo(CamelModel):
|
|
"""Usage info for a specific model."""
|
|
|
|
model: str
|
|
total_tokens: int = Field(0, alias="totalTokens")
|
|
prompt_tokens: int = Field(0, alias="promptTokens")
|
|
completion_tokens: int = Field(0, alias="completionTokens")
|
|
request_count: int = Field(0, alias="requestCount")
|
|
total_cost: float = Field(0.0, alias="totalCost")
|
|
|
|
|
|
class UsageOverview(CamelModel):
|
|
"""Complete usage overview (gas gauge)."""
|
|
|
|
totals: UsageTotals
|
|
by_agent: List[AgentUsageInfo] = Field(default_factory=list, alias="byAgent")
|
|
by_model: List[ModelUsageInfo] = Field(default_factory=list, alias="byModel")
|
|
active_threads: int = Field(0, alias="activeThreads")
|
|
|
|
|
|
class UsageResponse(CamelModel):
|
|
"""Response for GET /usage."""
|
|
|
|
usage: UsageOverview
|
|
|
|
|
|
class ThreadBudgetListResponse(CamelModel):
|
|
"""Response for GET /usage/threads."""
|
|
|
|
threads: List[ThreadBudgetInfo]
|
|
count: int
|
|
default_max_tokens: int = Field(alias="defaultMaxTokens")
|