129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
"""Tests del grafo compilado: camino feliz con MockProvider y pausa HITL."""
|
|
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from agentforge_core.domain.agent import AgentDefinition, LLMConfig
|
|
from agentforge_core.domain.policy import PolicyDefinition
|
|
from agentforge_core.guardrails.guardrails_ai import GuardrailsAIEngine
|
|
from agentforge_core.llm.mock import MockProvider
|
|
from agentforge_core.runtime.checkpointer import build_checkpointer
|
|
from agentforge_core.runtime.graph import build_graph
|
|
|
|
|
|
def _agent() -> AgentDefinition:
|
|
return AgentDefinition(
|
|
name="incident_analyzer",
|
|
version="v1",
|
|
owner="Juan",
|
|
purpose="Análisis de incidentes",
|
|
state="active",
|
|
guardrails=["default"],
|
|
llm=LLMConfig(provider="mock"),
|
|
system_prompt="Eres un analista. Responde SIEMPRE con JSON.",
|
|
output_schema={"type": "object"},
|
|
risk_threshold_for_hitl=4,
|
|
updated_at=datetime.now(UTC),
|
|
)
|
|
|
|
|
|
def _policy_min() -> PolicyDefinition:
|
|
return PolicyDefinition(
|
|
name="min",
|
|
version="v1",
|
|
description="t",
|
|
input_validators=[],
|
|
output_validators=[],
|
|
)
|
|
|
|
|
|
def _initial_state(trace_id: str, agent: AgentDefinition, user_input: str) -> dict:
|
|
return {
|
|
"trace_id": trace_id,
|
|
"agent_name": agent.name,
|
|
"agent_version": agent.version,
|
|
"user_input": user_input,
|
|
"messages": [],
|
|
"raw_llm_output": None,
|
|
"parsed_output": None,
|
|
"proposed_actions": [],
|
|
"violations": [],
|
|
"decision_path": [],
|
|
"status": "running",
|
|
"error": None,
|
|
"human_decision": None,
|
|
"final_output": None,
|
|
}
|
|
|
|
|
|
async def test_grafo_completa_camino_feliz_sin_hitl(tmp_path: Path) -> None:
|
|
agent = _agent()
|
|
trace_id = str(uuid4())
|
|
config = {"configurable": {"thread_id": trace_id}}
|
|
async with build_checkpointer(tmp_path) as cp:
|
|
graph = build_graph(
|
|
agent_def=agent,
|
|
policy=_policy_min(),
|
|
provider=MockProvider(),
|
|
engine=GuardrailsAIEngine(),
|
|
checkpointer=cp,
|
|
)
|
|
final = await graph.ainvoke(
|
|
_initial_state(trace_id, agent, "degradación MOS pool SBC"), # risk=2 → no HITL
|
|
config=config,
|
|
)
|
|
assert final["status"] == "completed"
|
|
assert final["final_output"]["approved_actions"] # propagó la acción no riesgosa
|
|
|
|
|
|
async def test_grafo_pausa_en_hitl_si_riesgo_alto(tmp_path: Path) -> None:
|
|
agent = _agent()
|
|
trace_id = str(uuid4())
|
|
config = {"configurable": {"thread_id": trace_id}}
|
|
async with build_checkpointer(tmp_path) as cp:
|
|
graph = build_graph(
|
|
agent_def=agent,
|
|
policy=_policy_min(),
|
|
provider=MockProvider(),
|
|
engine=GuardrailsAIEngine(),
|
|
checkpointer=cp,
|
|
)
|
|
await graph.ainvoke(
|
|
_initial_state(trace_id, agent, "caída registros sip"), # act-1 risk=4 → HITL
|
|
config=config,
|
|
)
|
|
# En interrupt, ainvoke devuelve el snapshot; el grafo queda con un paso pendiente.
|
|
state = await graph.aget_state(config)
|
|
assert state.next # interrumpido en approve_gate, esperando resume()
|
|
assert state.values["status"] == "running"
|
|
|
|
|
|
async def test_grafo_bloquea_por_guardrail_de_entrada(tmp_path: Path) -> None:
|
|
agent = _agent()
|
|
policy = PolicyDefinition(
|
|
name="block-pii",
|
|
version="v1",
|
|
description="bloquea emails en la entrada",
|
|
input_validators=[
|
|
{"type": "detect_pii", "config": {"entities": ["EMAIL_ADDRESS"], "severity_on_match": "block"}}
|
|
],
|
|
output_validators=[],
|
|
)
|
|
trace_id = str(uuid4())
|
|
config = {"configurable": {"thread_id": trace_id}}
|
|
async with build_checkpointer(tmp_path) as cp:
|
|
graph = build_graph(
|
|
agent_def=agent,
|
|
policy=policy,
|
|
provider=MockProvider(),
|
|
engine=GuardrailsAIEngine(),
|
|
checkpointer=cp,
|
|
)
|
|
final = await graph.ainvoke(
|
|
_initial_state(trace_id, agent, "manda correo a juan@example.com"),
|
|
config=config,
|
|
)
|
|
assert final["status"] == "blocked_by_guardrail"
|
|
assert final["raw_llm_output"] is None # nunca llegó al LLM
|