File size: 3,305 Bytes
1f2d50a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
"""Planning service for KGraph-MCP API."""
import logging
from agents.planner import SimplePlannerAgent
from ..models.responses import (
PlannedStepResponse,
PlanResponse,
PromptInfo,
ToolSuggestionResponse,
)
logger = logging.getLogger(__name__)
class PlanningService:
"""Service for handling planning-related operations."""
def __init__(self, planner_agent: SimplePlannerAgent):
"""Initialize the planning service."""
self.planner_agent = planner_agent
def suggest_tools(self, query: str, top_k: int = 3) -> list[ToolSuggestionResponse]:
"""Suggest tools based on user query."""
try:
# Get tool suggestions from the planner agent
planned_steps = self.planner_agent.generate_plan(query, top_k=top_k)
# Convert to response models
suggestions = []
for step in planned_steps:
suggestion = ToolSuggestionResponse(
tool_id=step.tool.tool_id,
name=step.tool.name,
description=step.tool.description,
tags=step.tool.tags or [],
invocation_command_stub=step.tool.invocation_command_stub,
)
suggestions.append(suggestion)
return suggestions
except Exception as e:
logger.error(f"Error suggesting tools: {e}")
return []
def generate_plan(self, query: str, top_k: int = 3) -> PlanResponse:
"""Generate a comprehensive plan with tool+prompt combinations."""
try:
# Get planned steps from the planner agent
planned_steps = self.planner_agent.generate_plan(query, top_k=top_k)
# Convert to response models
plan_steps = []
for step in planned_steps:
tool_response = ToolSuggestionResponse(
tool_id=step.tool.tool_id,
name=step.tool.name,
description=step.tool.description,
tags=step.tool.tags or [],
invocation_command_stub=step.tool.invocation_command_stub,
)
prompt_info = PromptInfo(
prompt_id=step.prompt.prompt_id,
name=step.prompt.name,
description=step.prompt.description,
template_string=step.prompt.template_string,
difficulty_level=step.prompt.difficulty_level,
input_variables=step.prompt.input_variables,
)
planned_step_response = PlannedStepResponse(
tool=tool_response,
prompt=prompt_info,
relevance_score=step.relevance_score,
summary=step.summary,
)
plan_steps.append(planned_step_response)
return PlanResponse(
query=query,
planned_steps=plan_steps,
total_steps=len(plan_steps),
)
except Exception as e:
logger.error(f"Error generating plan: {e}")
return PlanResponse(
query=query,
planned_steps=[],
total_steps=0,
)
|