agentic_workflow/planner.py

import json
import chromadb
from agentic_workflow.config import LLM_MODEL
from langchain_ollama import ChatOllama
from dataclasses import dataclass, field

llm = ChatOllama(model=LLM_MODEL, temperature=0)

@dataclass
class PlanStep:
    name: str       # e.g. "retrieval", "pii_check"
    enabled: bool   # whether this step should run
    reason: str     # why it was enabled or skipped

@dataclass
class Plan:
    steps: list[PlanStep] = field(default_factory=list)

    def is_enabled(self, step_name: str) -> bool:
        """Quick lookup — used by the pipeline to check if a step should run."""
        return any(s.name == step_name and s.enabled for s in self.steps)

    def summary(self) -> None:
        """Prints a readable breakdown of the plan."""
        print("\n── Execution Plan ──────────────────────────")
        for s in self.steps:
            status = "RUN " if s.enabled else "SKIP"
            print(f"  [{status}] {s.name:<20} → {s.reason}")
        print("────────────────────────────────────────────\n")

def create_plan(intent_result: dict) -> Plan:
    """
    Reads the intent analysis output and decides which steps
    the pipeline should execute, with a reason for each decision.

    Args:
        intent_result: dict returned by analyse_intent(), e.g.
            {
                "intent": "factual",
                "needs_context": True,
                "needs_history": False
            }

    Returns:
        Plan object with a PlanStep for each possible action.
    """
    intent        = intent_result.get("intent", "factual")
    needs_context = intent_result.get("needs_context", True)
    needs_history = intent_result.get("needs_history", False)

    steps = []

    # PII check
    # Always runs — sensitive queries must be gated regardless of intent
    steps.append(PlanStep(
        name="pii_check",
        enabled=True,
        reason="Always runs to protect against personal data leakage"
    ))

    # Vector retrieval
    # Only runs for factual queries that need external knowledge
    steps.append(PlanStep(
        name="retrieval",
        enabled=needs_context,
        reason=(
            "Query requires knowledge base context"
            if needs_context
            else f"Skipped — '{intent}' intent does not need retrieval"
        )
    ))

    # Conversation history
    # Only runs when the query references a prior turn
    steps.append(PlanStep(
        name="history",
        enabled=needs_history,
        reason=(
            "Query appears to reference a prior turn"
            if needs_history
            else "Skipped — query is self-contained"
        )
    ))

    # LLM generation
    # Always runs — we always need a response
    steps.append(PlanStep(
        name="llm_generation",
        enabled=True,
        reason="Always runs to generate the final response"
    ))

    # Quality evaluation
    # Always runs — every response should be checked before returning
    steps.append(PlanStep(
        name="quality_eval",
        enabled=True,
        reason="Always runs to verify response quality before returning"
    ))

    return Plan(steps=steps)