121 lines
4.6 KiB
Python
121 lines
4.6 KiB
Python
"""
|
|
Phase 2 #2C: Loop Primitive + Verifier
|
|
Adds loop edges and verifier nodes to the orchestrator.
|
|
"""
|
|
import asyncio
|
|
from typing import Dict, List, Any, Callable
|
|
|
|
|
|
class LoopConfig:
|
|
"""Configuration for a loop in the agent DAG."""
|
|
def __init__(self, max_iterations: int = 5, condition: str = "not_passed",
|
|
threshold: float = 0.8):
|
|
self.max_iterations = max_iterations
|
|
self.condition = condition
|
|
self.threshold = threshold
|
|
self.current_iteration = 0
|
|
|
|
def should_continue(self, verifier_result: dict) -> bool:
|
|
self.current_iteration += 1
|
|
if self.current_iteration > self.max_iterations:
|
|
return False
|
|
if self.condition == "not_passed":
|
|
return not verifier_result.get("passed", False)
|
|
elif self.condition == "score_below":
|
|
return verifier_result.get("score", 0) < self.threshold
|
|
elif self.condition == "always":
|
|
return self.current_iteration <= self.max_iterations
|
|
return False
|
|
|
|
|
|
class VerifierNode:
|
|
"""Node that checks whether a goal has been achieved."""
|
|
|
|
def __init__(self, goal: str, check_type: str = "quality"):
|
|
self.goal = goal
|
|
self.check_type = check_type
|
|
|
|
async def verify(self, output: str, context: dict) -> dict:
|
|
"""Run verification via LLM to check if goal is met."""
|
|
try:
|
|
import httpx, json
|
|
prompt = (
|
|
f"Goal: {self.goal}\n"
|
|
f"Output to verify:\n{output[:2000]}\n\n"
|
|
f"Check if the output meets the goal. Return JSON:\n"
|
|
f'{{"passed": true/false, "score": 0.0-1.0, "reason": "brief explanation"}}'
|
|
)
|
|
async with httpx.AsyncClient(timeout=60) as client:
|
|
resp = await client.post(
|
|
"http://qihu-aiflow-litellm:4000/v1/chat/completions",
|
|
headers={
|
|
"Authorization": "Bearer sk-litellm-master-key-aiflow-2026",
|
|
"Content-Type": "application/json"
|
|
},
|
|
json={
|
|
"model": "deepseek-v4-flash",
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": 0.1,
|
|
"max_tokens": 200,
|
|
}
|
|
)
|
|
if resp.status_code == 200:
|
|
text = resp.json()["choices"][0]["message"]["content"].strip()
|
|
if text.startswith("```"):
|
|
text = text.split("\n", 1)[1][:-3]
|
|
return json.loads(text)
|
|
except Exception:
|
|
pass
|
|
return {"passed": True, "score": 1.0, "reason": "verification skipped (offline fallback)"}
|
|
|
|
|
|
class LoopExecutor:
|
|
"""Executes a node graph with loop support."""
|
|
|
|
def __init__(self, nodes: List[Dict], edges: List[Dict]):
|
|
self.nodes = {n["id"]: n for n in nodes}
|
|
self.edges = edges
|
|
self.loop_configs: Dict[str, LoopConfig] = {}
|
|
self.verifiers: Dict[str, VerifierNode] = {}
|
|
|
|
def add_loop(self, from_node: str, to_node: str, config: LoopConfig):
|
|
self.loop_configs[from_node] = config
|
|
|
|
def add_verifier(self, node_id: str, verifier: VerifierNode):
|
|
self.verifiers[node_id] = verifier
|
|
|
|
async def execute(self, node_executor: Callable, initial_input: dict) -> dict:
|
|
"""Execute DAG with loop support."""
|
|
results = {}
|
|
to_process = ["input-1"] # Start from input node
|
|
|
|
while to_process:
|
|
node_id = to_process.pop(0)
|
|
node = self.nodes.get(node_id)
|
|
if not node:
|
|
continue
|
|
|
|
# Check loop limit
|
|
if node_id in self.loop_configs:
|
|
lc = self.loop_configs[node_id]
|
|
if lc.current_iteration >= lc.max_iterations:
|
|
continue
|
|
|
|
output = await node_executor(node, initial_input, results)
|
|
results[node_id] = output
|
|
|
|
# Find outgoing edges
|
|
next_nodes = [e["target"] for e in self.edges if e["source"] == node_id]
|
|
for nn in next_nodes:
|
|
if nn in self.loop_configs:
|
|
# Check if loop should continue
|
|
lc = self.loop_configs[nn]
|
|
if lc.should_continue(output):
|
|
target = [e["target"] for e in self.edges if e["source"] == nn]
|
|
if target:
|
|
to_process.append(target[0]) # Loop back
|
|
else:
|
|
to_process.append(nn)
|
|
|
|
return results
|