Implement complete agent platform with four major layers:
P5 - Engine Layer (LangGraph):
- AgentState with JSON serialization validation
- 5 async nodes (ask, plan, act, verify, report) with error boundaries
- Template resolution system for {{user_input}} and {{steps.xxx}}
- StateGraph with conditional edges and interrupt_after for approval flow
- 17/17 tests passing, 90% coverage
P6 - LLM Layer:
- ClaudeClient with Decimal-based cost calculation (ROUND_HALF_UP)
- Custom exceptions (ClaudeAPIError, ClaudeTimeoutError, ClaudeRateLimitError)
- PromptBuilder with cache_control support
- ModelRouter with sensitivity-based routing (local vs Claude)
- 18/18 tests passing, 90% coverage
P7 - Audit Layer (Three-tier):
- Sanitizer: mask phone/ID card/API keys, contains_sensitive, mask_dict
- Encryptor: AES-256-GCM with keyring support, 12-byte IV
- AuditLogger: L1 (DB), L2 (OTel), L3 (Encrypted payloads)
- 24/24 tests passing, 86% coverage
P8 - API Layer (FastAPI):
- Lifespan management with 8-step startup sequence
- Middleware: TracingMiddleware, RateLimitMiddleware (Redis), CORS
- Auth: JWT decode with python-jose, RBAC with permission matrix
- Routes: health, chat (SSE), approve, trace
- 13/13 tests passing, 91% coverage
All gate checks passing:
- pytest with 85%+ coverage
- mypy type checking
- ruff style checking
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
124 lines
6.4 KiB
Python
124 lines
6.4 KiB
Python
"""
|
|
AgentState — LangGraph state model (fully JSON-serializable).
|
|
|
|
Follows strict LangGraph patterns:
|
|
- BaseModel with frozen=False (nodes mutate in-place)
|
|
- No large objects stored directly (only references)
|
|
- Checkpoint-compatible (PostgresSaver requires JSON serialization)
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
class AgentState(BaseModel):
|
|
"""
|
|
LangGraph state for the agent execution flow.
|
|
|
|
All fields must be JSON-serializable for checkpoint persistence.
|
|
Large objects (documents, raw data) are stored externally with
|
|
only references (UUIDs, paths) kept in state.
|
|
"""
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Identity
|
|
# ═══════════════════════════════════════════════════════════════
|
|
trace_id: str
|
|
session_id: str
|
|
user_id: str
|
|
tenant_id: str
|
|
user_role: str
|
|
data_scope: list[str] = Field(default_factory=list)
|
|
profile_type: str
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Execution Context
|
|
# ═══════════════════════════════════════════════════════════════
|
|
user_input: str
|
|
extracted_params: dict[str, Any] = Field(default_factory=dict)
|
|
execution_plan: list[dict[str, Any]] = Field(default_factory=list)
|
|
current_step_index: int = 0
|
|
step_outputs: dict[str, Any] = Field(default_factory=dict)
|
|
# step_outputs stores only summaries and data_ref, NOT raw large objects
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Approval
|
|
# ═══════════════════════════════════════════════════════════════
|
|
pending_approval: dict[str, Any] | None = None # ApprovalPayload dict
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Quality Assessment
|
|
# ═══════════════════════════════════════════════════════════════
|
|
evidence_score: Decimal | None = None
|
|
evidence_level: str | None = None # sufficient/limited/insufficient
|
|
citations: list[dict[str, Any]] = Field(default_factory=list)
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Output
|
|
# ═══════════════════════════════════════════════════════════════
|
|
final_output: str | None = None
|
|
output_type: str | None = None # document/text/insufficient_data/clarification_needed
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Error
|
|
# ═══════════════════════════════════════════════════════════════
|
|
error: dict[str, Any] | None = None # {"code", "message", "node", "trace_id"}
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Internal Tracking (not exposed to user)
|
|
# ═══════════════════════════════════════════════════════════════
|
|
node_durations: dict[str, float] = Field(default_factory=dict)
|
|
total_cost_cny: Decimal = Field(default=Decimal("0"))
|
|
|
|
# ═══════════════════════════════════════════════════════════════
|
|
# Helper Methods
|
|
# ═══════════════════════════════════════════════════════════════
|
|
|
|
def is_clarification_needed(self) -> bool:
|
|
"""Check if user clarification is required."""
|
|
return self.output_type == "clarification_needed"
|
|
|
|
def has_error(self) -> bool:
|
|
"""Check if an error occurred during execution."""
|
|
return self.error is not None
|
|
|
|
def add_step_output(
|
|
self, key: str, summary: str, data_ref: str | None = None
|
|
) -> None:
|
|
"""
|
|
Store step result summary (not raw large objects).
|
|
|
|
Args:
|
|
key: Step identifier (output_key from execution plan)
|
|
summary: Human-readable summary of the step result
|
|
data_ref: Optional reference to external storage (UUID, path, etc.)
|
|
"""
|
|
self.step_outputs[key] = {
|
|
"summary": summary,
|
|
"data_ref": data_ref,
|
|
"key": key,
|
|
}
|
|
|
|
@classmethod
|
|
def validate_serializable(cls, state: "AgentState") -> bool:
|
|
"""
|
|
Validate that state can be JSON-serialized (LangGraph checkpoint requirement).
|
|
|
|
Args:
|
|
state: AgentState instance to validate
|
|
|
|
Returns:
|
|
True if serializable, False otherwise
|
|
"""
|
|
import json
|
|
|
|
try:
|
|
# Use Pydantic's model_dump with mode='json' to handle Decimal and other types
|
|
dumped = state.model_dump(mode='json')
|
|
json.dumps(dumped)
|
|
return True
|
|
except (TypeError, ValueError):
|
|
return False
|