api/services/planning.py

"""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,
            )