81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
from services.runtime_variables import DynamicVariableStore
|
||
|
|
from services.workflow.models import LLMRouteResult, RouteStatus
|
||
|
|
from services.workflow.routing import WorkflowEdgeEvaluator
|
||
|
|
from services.workflow_engine import WorkflowEngine
|
||
|
|
|
||
|
|
|
||
|
|
def routing_graph() -> dict:
|
||
|
|
return {
|
||
|
|
"specVersion": 3,
|
||
|
|
"settings": {},
|
||
|
|
"nodes": [
|
||
|
|
{"id": "start", "type": "start", "data": {}},
|
||
|
|
{"id": "agent", "type": "agent", "data": {"name": "Agent"}},
|
||
|
|
{"id": "matched", "type": "end", "data": {}},
|
||
|
|
{"id": "fallback", "type": "end", "data": {}},
|
||
|
|
],
|
||
|
|
"edges": [
|
||
|
|
{
|
||
|
|
"id": "condition",
|
||
|
|
"source": "agent",
|
||
|
|
"target": "matched",
|
||
|
|
"data": {
|
||
|
|
"mode": "llm",
|
||
|
|
"priority": 10,
|
||
|
|
"condition": "用户要求结束",
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"id": "default",
|
||
|
|
"source": "agent",
|
||
|
|
"target": "fallback",
|
||
|
|
"data": {"mode": "always", "priority": 10},
|
||
|
|
},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class WorkflowEdgeEvaluatorTest(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def test_router_error_does_not_take_default_path(self):
|
||
|
|
class ErrorRouter:
|
||
|
|
async def select_edge(self, **_kwargs):
|
||
|
|
return LLMRouteResult(
|
||
|
|
status=RouteStatus.ERROR,
|
||
|
|
error="provider unavailable",
|
||
|
|
)
|
||
|
|
|
||
|
|
evaluator = WorkflowEdgeEvaluator(
|
||
|
|
WorkflowEngine(routing_graph()),
|
||
|
|
DynamicVariableStore({}),
|
||
|
|
lambda _node_id: ErrorRouter(),
|
||
|
|
)
|
||
|
|
|
||
|
|
decision = await evaluator.evaluate("agent")
|
||
|
|
|
||
|
|
self.assertEqual(decision.status, RouteStatus.ERROR)
|
||
|
|
self.assertIsNone(decision.edge)
|
||
|
|
|
||
|
|
async def test_explicit_no_match_takes_default_path(self):
|
||
|
|
class NoMatchRouter:
|
||
|
|
async def select_edge(self, **_kwargs):
|
||
|
|
return LLMRouteResult(status=RouteStatus.NO_MATCH)
|
||
|
|
|
||
|
|
evaluator = WorkflowEdgeEvaluator(
|
||
|
|
WorkflowEngine(routing_graph()),
|
||
|
|
DynamicVariableStore({}),
|
||
|
|
lambda _node_id: NoMatchRouter(),
|
||
|
|
)
|
||
|
|
|
||
|
|
decision = await evaluator.evaluate("agent")
|
||
|
|
|
||
|
|
self.assertEqual(decision.status, RouteStatus.MATCHED)
|
||
|
|
self.assertEqual(decision.edge["id"], "default")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|