src/tools/reasoning.py

from textwrap import dedent
from typing import Any, Dict, List, Optional

from smolagents.tools import Tool  # SmolAgents base class


# ---------------------------------------------------------------------
# Helper enum – kept as str literals so we avoid any Agno dependency.
# ---------------------------------------------------------------------
class NextAction:
    CONTINUE = "continue"
    VALIDATE = "validate"
    FINAL_ANSWER = "final_answer"


# ---------------------------------------------------------------------
# THINK TOOL -----------------------------------------------------------
# ---------------------------------------------------------------------
class ThinkTool(Tool):
    name = "think"
    description = (
        "Internal scratch‑pad. Use this to reason step‑by‑step before "
        "calling other tools or replying to the user."
    )
    inputs = {
        "title": {"type": "string", "description": "Concise title"},
        "thought": {"type": "string", "description": "Detailed reasoning"},
        "action": {
            "type": "string",
            "description": "Intended next action",
            "nullable": True,
        },
        "confidence": {
            "type": "number",
            "description": "Confidence 0–1",
            "nullable": True,
        },
        "run_id": {
            "type": "string",
            "description": "Execution identifier",
            "nullable": True,
        },
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self._history: Dict[str, List[Dict[str, Any]]] = {}

    def forward(  # noqa: N802  (SmolAgents allows camelCase here)
        self,
        title: str,
        thought: str,
        action: Optional[str] = None,
        confidence: float = 0.8,
        run_id: str = "default",
    ) -> str:
        """Store and pretty‑print reasoning history."""
        step = {
            "title": title,
            "reasoning": thought,
            "action": action,
            "confidence": confidence,
        }
        self._history.setdefault(run_id, []).append(step)

        # Pretty print full chain so the LLM can “see” prior steps
        formatted = ""
        for idx, s in enumerate(self._history[run_id], 1):
            formatted += (
                dedent(
                    f"""\
                Step {idx}:
                Title: {s["title"]}
                Reasoning: {s["reasoning"]}
                Action: {s["action"]}
                Confidence: {s["confidence"]}
                """
                )
                + "\n"
            )
        return formatted.strip()


# ---------------------------------------------------------------------
# ANALYZE TOOL ---------------------------------------------------------
# ---------------------------------------------------------------------
class AnalyzeTool(Tool):
    name = "analyze"
    description = (
        "Evaluate the result of previous actions and decide whether to "
        "continue, validate, or provide a final answer. "
    )
    inputs = {
        "title": {"type": "string", "description": "Concise title"},
        "result": {"type": "string", "description": "Outcome being analysed"},
        "analysis": {"type": "string", "description": "Your analysis"},
        "next_action": {
            "type": "string",
            "description": "'continue' | 'validate' | 'final_answer'",
            "nullable": True,
        },
        "confidence": {
            "type": "number",
            "description": "Confidence 0–1",
            "nullable": True,
        },
        "run_id": {
            "type": "string",
            "description": "Execution identifier",
            "nullable": True,
        },
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self._history: Dict[str, List[Dict[str, Any]]] = {}

    def forward(
        self,
        title: str,
        result: str,
        analysis: str,
        next_action: str = NextAction.CONTINUE,
        confidence: float = 0.8,
        run_id: str = "default",
    ) -> str:
        if next_action not in {
            NextAction.CONTINUE,
            NextAction.VALIDATE,
            NextAction.FINAL_ANSWER,
        }:
            raise ValueError(
                f"next_action must be one of "
                f"{NextAction.CONTINUE}, {NextAction.VALIDATE}, "
                f"{NextAction.FINAL_ANSWER}"
            )

        step = {
            "title": title,
            "result": result,
            "reasoning": analysis,
            "next_action": next_action,
            "confidence": confidence,
        }
        self._history.setdefault(run_id, []).append(step)

        formatted = ""
        for idx, s in enumerate(self._history[run_id], 1):
            formatted += (
                dedent(
                    f"""\
                Step {idx}:
                Title: {s["title"]}
                Result: {s.get("result")}
                Reasoning: {s["reasoning"]}
                Next action: {s.get("next_action")}
                Confidence: {s["confidence"]}
                """
                )
                + "\n"
            )
        return formatted.strip()


# ---------------------------------------------------------------------
# TOOLKIT WRAPPER ------------------------------------------------------
# ---------------------------------------------------------------------
class ReasoningToolkit:
    """
    Convenience wrapper so you can write:

        from reasoning_tools import ReasoningToolkit
        toolkit = ReasoningToolkit()
        agent = CodeAgent(tools=toolkit.tools, model=...)
    """

    DEFAULT_INSTRUCTIONS = dedent(
        """\
        You have access to two internal tools – **think** and **analyze** –
        for chain‑of‑thought reasoning. **Always** call `think` before
        external tool calls or final answers, then call `analyze` to
        decide whether to continue, validate, or finish."""
    )

    def __init__(self, think: bool = True, analyze: bool = True):
        self.tools: List[Tool] = []
        if think:
            self.tools.append(ThinkTool())
        if analyze:
            self.tools.append(AnalyzeTool())

    def with_instructions(self, extra: str | None = None) -> str:
        return self.DEFAULT_INSTRUCTIONS + ("\n" + extra if extra else "")