diff --git a/backend/agent/router.py b/backend/agent/router.py new file mode 100644 index 0000000..abe008b --- /dev/null +++ b/backend/agent/router.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional + +from sqlalchemy.orm import Session + +from ..app_settings import get_workflow_router_model_preference +from ..ollama_client import chat_typed, list_model_catalog +from .models import WorkflowDefinition + + +ROUTER_FORMAT = { + "type": "object", + "properties": { + "workflow_id": {"type": "string"}, + "confidence": {"type": "number"}, + "reason": {"type": "string"}, + "inputs": {"type": "object"}, + }, + "required": ["workflow_id", "confidence", "reason", "inputs"], + "additionalProperties": False, +} + + +def workflow_manifests(db: Session) -> List[Dict[str, Any]]: + workflows = db.query(WorkflowDefinition).filter(WorkflowDefinition.enabled.is_(True)).order_by(WorkflowDefinition.name.asc()).all() + manifests = [] + for item in workflows: + manifests.append({ + "workflow_id": item.id, + "slug": item.slug, + "name": item.name, + "routing_description": item.routing_description, + "examples": json.loads(item.routing_examples_json or "[]"), + "estimated_cost_class": item.estimated_cost_class, + "required_capabilities": json.loads(item.required_capabilities_json or "[]"), + }) + return manifests + + +async def select_workflow( + db: Session, + *, + message: str, + recent_messages: List[Dict[str, Any]], + attachments: List[Dict[str, Any]], + library_slug: Optional[str], + router_model: Optional[str], + chat_model: Optional[str], + confidence_threshold: float = 0.55, +) -> Dict[str, Any]: + fallback = db.query(WorkflowDefinition).filter(WorkflowDefinition.slug == "input-output").first() + if fallback is None: + raise RuntimeError("The built-in input-output workflow is missing.") + fallback_result = { + "workflow_id": fallback.id, + "workflow_slug": fallback.slug, + "confidence": 0.0, + "reason": "Router fallback to direct chat.", + "inputs": {}, + "fallback": True, + } + manifests = workflow_manifests(db) + if attachments: + vision = next((item for item in manifests if item["slug"] == "vision-answer"), None) + if vision: + return {**fallback_result, "workflow_id": vision["workflow_id"], "workflow_slug": vision["slug"], "confidence": 1.0, "reason": "The request includes attachments.", "fallback": False} + + preferred = str(router_model or get_workflow_router_model_preference() or "").strip() + model = preferred or str(chat_model or "").strip() + if not model: + return fallback_result + if preferred: + try: + catalog = await list_model_catalog() + if preferred not in (catalog.get("chat_models") or []): + model = str(chat_model or "").strip() + except Exception: + model = str(chat_model or "").strip() or preferred + if not model: + return fallback_result + + compact_history = [ + {"role": item.get("role"), "content": str(item.get("content") or "")[:700]} + for item in recent_messages[-4:] + ] + prompt = ( + "Select exactly one enabled workflow for the current request. Do not invent workflows. " + "Prefer the lowest-cost workflow that can complete the task. Return JSON only.\n\n" + f"Current request: {message}\n" + f"Recent messages: {json.dumps(compact_history, ensure_ascii=False)}\n" + f"Attachment count: {len(attachments)}\n" + f"Selected knowledge database: {library_slug or 'none'}\n" + f"Enabled workflows: {json.dumps(manifests, ensure_ascii=False)}" + ) + try: + result = await chat_typed( + model, + [{"role": "user", "content": prompt}], + format=ROUTER_FORMAT, + options={"temperature": 0, "num_ctx": 4096}, + ) + parsed = json.loads(result.content) + workflow_id = str(parsed.get("workflow_id") or "") + selected = next((item for item in manifests if item["workflow_id"] == workflow_id or item["slug"] == workflow_id), None) + confidence = float(parsed.get("confidence") or 0) + if selected is None or confidence < confidence_threshold: + return fallback_result + if "rag" in selected["required_capabilities"] and not library_slug: + return fallback_result + return { + "workflow_id": selected["workflow_id"], + "workflow_slug": selected["slug"], + "confidence": max(0.0, min(confidence, 1.0)), + "reason": str(parsed.get("reason") or ""), + "inputs": parsed.get("inputs") if isinstance(parsed.get("inputs"), dict) else {}, + "fallback": False, + "model": model, + } + except Exception as exc: + return {**fallback_result, "reason": f"Router failed: {type(exc).__name__}: {exc}"}