agentpro/tools/cbt_tool.py

from enum import Enum
from typing import Dict, Any, List, Optional
import json
from .base import Tool
class CBTExerciseType(Enum):
    ART_OF_WORRY = "Art of Worry"
    WORD_REFRAMING = "Word Reframing"
    STICKY_NOTE_PROJECT = "Sticky Note Project"
    LOST_LUGGAGE = "Lost Luggage"
    JUST_PASSING_THROUGH = "Just Passing Through"
    DOWN_THE_RABBIT_HOLE = "Down the Rabbit Hole"

class CBTExerciseTool(Tool):
    name: str = "CBT Exercise Tool"
    description: str = "a tool that provides specialized guidance for cognitive behavioral therapy exercises. it returns contextual information, prompts, techniques, and progress tracking for specific cbt exercises."
    arg: str = """dict with:
    - "exercise_type": (required) one of: "art of worry", "word reframing", "sticky note project", "lost luggage", 
       "just passing through", "down the rabbit hole"
    - "action": (required) one of: "get_context", "get_prompt", "check_completion", "get_next_technique" 
    - additional parameters based on action:
      - for "get_prompt": include "stage" (integer)
      - for "check_completion": include "progress" (list of strings)
      - for "get_next_technique": include "used_techniques" (list of strings)
    ensure the action input is json parseable."""

    def run(self, prompt: str) -> str:
        # Handle string input
        if isinstance(prompt, str):
            try:
                input_data = json.loads(prompt)
                print("Converted input from str to JSON Dictionary")
            except Exception as e:
                return f"Error: {e}. Input must be valid JSON."
        else:
            input_data = prompt
        
        # Validate input
        if not isinstance(input_data, dict):
            return "Error: Input must be a dictionary."
            
        exercise_name = input_data.get("exercise_type")
        if not exercise_name:
            return "Error: 'exercise_type' is required."
            
        action = input_data.get("action")
        if not action:
            return "Error: 'action' is required."
            
        # Map exercise name to enum
        exercise_map = {
            "art of worry": CBTExerciseType.ART_OF_WORRY,
            "word reframing": CBTExerciseType.WORD_REFRAMING,
            "sticky note project": CBTExerciseType.STICKY_NOTE_PROJECT,
            "lost luggage": CBTExerciseType.LOST_LUGGAGE,
            "just passing through": CBTExerciseType.JUST_PASSING_THROUGH,
            "down the rabbit hole": CBTExerciseType.DOWN_THE_RABBIT_HOLE
        }
        
        # Case-insensitive matching
        exercise_name = exercise_name.lower()
        exercise_type = exercise_map.get(exercise_name)
        if not exercise_type:
            return f"Error: Invalid exercise type '{exercise_name}'. Please choose from: {', '.join(exercise_map.keys())}"
        
        # Process based on action
        try:
            action = action.lower()
            if action == "get_context":
                result = self._get_exercise_context(exercise_type)
                return json.dumps(result)
                
            elif action == "get_prompt":
                stage = input_data.get("stage", 0)
                result = self._get_exercise_prompt(exercise_type, stage)
                return result
                
            elif action == "check_completion":
                progress = input_data.get("progress", [])
                result = self._check_completion(exercise_type, progress)
                return json.dumps(result)
                
            elif action == "get_next_technique":
                used_techniques = input_data.get("used_techniques", [])
                result = self._get_next_technique(exercise_type, used_techniques)
                return result if result else "No additional techniques available."
                
            else:
                return f"Error: Invalid action '{action}'. Please choose from: get_context, get_prompt, check_completion, get_next_technique"
                
        except Exception as e:
            return f"Error using CBT Exercise Tool: {str(e)}"
    
    def _get_exercise_context(self, exercise_type: CBTExerciseType) -> Dict[str, Any]:
        """Returns detailed context for a specific CBT exercise type."""
        return self._exercise_contexts.get(exercise_type, {})
    
    def _get_exercise_prompt(self, exercise_type: CBTExerciseType, stage: int = 0) -> str:
        """Returns a specific prompt for a given exercise type and stage."""
        context = self._get_exercise_context(exercise_type)
        prompts = context.get("prompts", [])
        
        if not prompts or stage >= len(prompts):
            return "What would you like to explore next with this exercise?"
        
        return prompts[stage]
    
    def _check_completion(self, exercise_type: CBTExerciseType, progress: List[str]) -> Dict[str, Any]:
        """Checks progress against completion indicators for an exercise."""
        context = self._get_exercise_context(exercise_type)
        indicators = context.get("completion_indicators", [])
        follow_ups = context.get("follow_up_questions", [])
        
        # Simple matching algorithm - could be enhanced with NLP
        completed = []
        for indicator in indicators:
            for note in progress:
                # Check if any progress note matches this indicator
                if any(keyword in note.lower() for keyword in indicator.lower().split()):
                    completed.append(indicator)
                    break
        
        completion_rate = len(completed) / len(indicators) if indicators else 0
        
        return {
            "completion_rate": completion_rate,
            "completed_indicators": completed,
            "pending_indicators": [i for i in indicators if i not in completed],
            "follow_up_questions": follow_ups,
            "is_complete": completion_rate > 0.75  # Consider complete if 75% of indicators met
        }
    
    def _get_next_technique(self, exercise_type: CBTExerciseType, used_techniques: List[str]) -> Optional[str]:
        """Suggests the next technique to try for a given exercise."""
        context = self._get_exercise_context(exercise_type)
        techniques = context.get("techniques", [])
        
        available = [t for t in techniques if t not in used_techniques]
        return available[0] if available else None
    
    # Exercise contexts database
    _exercise_contexts = {
        CBTExerciseType.ART_OF_WORRY: {
            "description": "Drawing your worries to gain a new perspective on them.",
            "prompts": [
                "What specific worry would you like to explore today?",
                "If you were to draw this worry, what might it look like?",
                "What colors would represent this worry?",
                "How does looking at your worry as an external object change your perspective?",
                "What do you notice about your relationship to this worry now?"
            ],
            "guidance": """
            Guide the user to externalize their worry through visualization or actual drawing.
            Help them describe it in detail - shape, color, texture, size.
            Encourage them to observe the worry from different angles.
            Ask how seeing the worry as separate from themselves changes their relationship to it.
            Explore how the worry's intensity might have changed through this process.
            """,
            "techniques": [
                "Externalization",
                "Visual processing",
                "Cognitive distancing",
                "Perspective taking"
            ],
            "completion_indicators": [
                "User has identified and described a worry",
                "User has visualized/drawn the worry",
                "User has observed it from a distance",
                "User has reflected on how their perspective has changed"
            ],
            "follow_up_questions": [
                "How did this exercise affect your relationship with this worry?",
                "What surprised you about visualizing your worry?",
                "How might you use this technique for other worries in the future?"
            ]
        },
        
        CBTExerciseType.WORD_REFRAMING: {
            "description": "Identifying and changing negative thought patterns through language.",
            "prompts": [
                "What negative thought would you like to work on reframing?",
                "How true is this thought on a scale of 0-100%?",
                "What evidence supports this thought? What evidence contradicts it?",
                "What's a more balanced or helpful way to think about this situation?",
                "How true does the new thought feel on a scale of 0-100%?"
            ],
            "guidance": """
            Help the user identify automatic negative thoughts.
            Guide them to evaluate the evidence for and against these thoughts.
            Assist in developing more balanced alternative thoughts.
            Use Socratic questioning to help them discover alternatives themselves.
            Avoid simply providing reframes - help them develop their own insights.
            """,
            "techniques": [
                "Cognitive restructuring",
                "Evidence examination",
                "Thought records",
                "Balanced thinking"
            ],
            "completion_indicators": [
                "User has identified negative automatic thoughts",
                "User has examined evidence for and against the thought",
                "User has generated more balanced alternatives",
                "User reports reduced belief in the original negative thought"
            ],
            "follow_up_questions": [
                "How does the reframed thought feel compared to the original?",
                "What impact might thinking this new way have on your emotions or behaviors?",
                "When might you need to use this reframing technique in daily life?"
            ]
        },
        
        CBTExerciseType.STICKY_NOTE_PROJECT: {
            "description": "Using sticky notes to track and challenge negative thoughts.",
            "prompts": [
                "What negative thoughts do you frequently experience?",
                "Which one would you like to challenge first?",
                "What evidence contradicts this negative thought?",
                "What would be a more helpful thought to replace it?",
                "Where could you place this positive reminder to see it regularly?"
            ],
            "guidance": """
            Guide users to identify recurring negative thoughts.
            Help them challenge each thought with evidence and logic.
            Support them in creating positive, realistic alternative statements.
            Suggest strategic placement of positive notes in their environment.
            Emphasize the importance of regular visual reinforcement.
            """,
            "techniques": [
                "Visual reminders",
                "Environmental restructuring",
                "Thought challenging",
                "Positive reinforcement"
            ],
            "completion_indicators": [
                "User has identified negative thoughts",
                "User has challenged these thoughts",
                "User has created positive alternatives",
                "User has a plan for placement of reminders"
            ],
            "follow_up_questions": [
                "How might regularly seeing these positive statements affect your thinking?",
                "What other negative thoughts could benefit from this approach?",
                "How will you know if this strategy is working for you?"
            ]
        },
        
       CBTExerciseType.LOST_LUGGAGE: {
        "description": "A metaphor-based exercise to manage ruminative thinking by observing thoughts rather than engaging with them.",
        "prompts": [
            "What are five recurring worrisome thoughts that keep looping in your mind?",
            "How does it feel when you try to stop or suppress these thoughts?",
            "Which of these thoughts feels the heaviest or most exhausting to carry?",
            "What would it be like to simply watch the thought go by, like luggage on a carousel?",
            "Which thought can you practice observing without picking up today?"
        ],
        "guidance": """
        Use the metaphor of a baggage carousel at an airport to help users visualize ruminative thoughts.
        Explain that trying to stop or control these thoughts often intensifies them, like a song on repeat with a broken stop button.
        Encourage users to see their thoughts as bags — constantly circulating — and explore the idea of *not picking them up*.
        Help them identify five ruminative thoughts and guide them to practice non-engagement (watching without grabbing).
        Emphasize that relief comes from letting the thoughts pass, not from solving or analyzing them.
        """,
        "techniques": [
            "Metaphorical thinking",
            "Ruminative thought identification",
            "Mindful observation",
            "Cognitive defusion",
            "Letting go practice"
        ],
        "completion_indicators": [
            "User has listed five recurring ruminative thoughts",
            "User has understood the cost of engaging with every thought",
            "User has visualized thoughts as luggage on a carousel",
            "User has practiced observing at least one thought without engaging"
        ],
        "follow_up_questions": [
            "What did you notice when you allowed yourself to just watch a thought instead of picking it up?",
            "Which thought felt hardest to let go of? Why?",
            "How did your body or emotions respond to the act of observing rather than engaging?",
            "What helps you remember that you don’t have to pick up every thought?",
            "Would you like to revisit this practice tomorrow or try it in real-time when a looping thought appears?"
        ]
             },
        
        CBTExerciseType.JUST_PASSING_THROUGH: {
            "description": "Mindfulness exercise to observe thoughts without attachment.",
            "prompts": [
                "Can you take a few deep breaths and notice what thoughts arise?",
                "Without trying to change them, can you observe these thoughts passing by?",
                "What happens when you don't engage with or follow these thoughts?",
                "Can you picture your thoughts as clouds passing across the sky?",
                "What do you notice about how thoughts naturally come and go?"
            ],
            "guidance": """
            Guide users in basic mindfulness practice.
            Help them observe thoughts without judgment or engagement.
            Introduce metaphors like "thoughts as clouds" or "thoughts as leaves on a stream."
            Emphasize that the goal isn't to stop thoughts but to change relationship to them.
            Encourage noticing the transient nature of thoughts.
            """,
            "techniques": [
                "Mindfulness",
                "Present-moment awareness",
                "Non-attachment",
                "Metacognitive awareness"
            ],
            "completion_indicators": [
                "User has practiced observing thoughts mindfully",
                "User has experienced non-engagement with thoughts",
                "User understands the transient nature of thoughts",
                "User can apply basic mindfulness to thoughts"
            ],
            "follow_up_questions": [
                "What did you notice about how thoughts naturally come and go?",
                "How was this different from how you usually relate to your thoughts?",
                "How might this practice help in stressful situations?"
            ]
        },
        
        CBTExerciseType.DOWN_THE_RABBIT_HOLE: {
        "description": "A sensory grounding and thought-tracking exercise designed to build awareness of how often thoughts drift away from the present moment and into negative loops.",
        "prompts": [
            "Find a comfortable position where you feel supported and relaxed.",
            "Begin by noticing what you can see around you — colors, shapes, light, or shadows.",
            "Shift your attention to what you can hear — nearby or distant sounds.",
            "Bring awareness to any smells or tastes present in this moment.",
            "Notice how your body feels — your back against the chair, your feet on the floor, or the air on your skin.",
            "As you observe, notice when your thoughts start to drift away. Where do they go?",
            "What kinds of thoughts pulled your attention away from the present?",
            "What sensations were easiest or hardest to stay focused on?"
        ],
        "guidance": """
        Help the user find a still, relaxed posture and begin tuning into their immediate physical environment.
        Encourage gentle, non-judgmental awareness of each sensory input: sight, sound, smell, taste, and touch.
        Guide them to notice when their mind drifts and what thoughts arise during that drift.
        Support reflection on how often this happens and what it reveals about their current mental state.
        Emphasize that the goal is not to stay perfectly focused, but to build awareness of their attention patterns.
        """,
        "techniques": [
            "Mindful observation",
            "Sensory grounding",
            "Thought tracking",
            "Cognitive defusion",
            "Present-moment awareness"
        ],
        "completion_indicators": [
            "User engaged in sensory observation across multiple senses",
            "User noticed and reflected on drifting thoughts",
            "User identified common thought patterns that interrupt presence",
            "User explored which senses were more grounding or distracting"
        ],
        "follow_up_questions": [
            "What kinds of thoughts distracted you most often?",
            "Which sense helped you feel the most present?",
            "How did your body feel before and after the exercise?",
            "How could you use this kind of observation in daily life to reset or recenter?"
        ]
    }
    }

if __name__ == "__main__":
    cbt_tool = CBTExerciseTool()
    call = cbt_tool