workflow.py

"""
Fraud Model Explainability Assistant - Workflow Implementation

This module defines the FraudExplainabilityWorkflow, a structured workflow pattern
that orchestrates intent classification, tool execution, and response generation
for the fraud analysis assistant.
"""

import os
import json
import logging
import asyncio
from typing import List, Dict, Any, Optional, TypedDict
from dataclasses import dataclass, field

# from strands import Tool # Tool type not directly exported
from strands.models.openai import OpenAIModel

# Import tools
from utils import (
    get_application_summary,
    explain_fraud_score,
    compare_to_population,
    check_fair_lending_flags,
    get_identity_network,
    get_model_performance,
)

# Import Confluence tools if available (handled dynamically)

logger = logging.getLogger(__name__)


class WorkflowState(TypedDict):
    """Represents the state of the workflow execution."""
    input_text: str
    messages: List[Dict[str, str]]
    intent: Optional[str]
    tool_calls: List[Dict[str, Any]]
    tool_outputs: List[Dict[str, Any]]
    final_response: Optional[str]
    error: Optional[str]


class FraudExplainabilityWorkflow:
    """
    Orchestrates the fraud analysis workflow:
    1. Analyze Intent
    2. Route to Tools
    3. Execute Tools
    4. Generate Response
    """

    def __init__(self, model_id: str = "gpt-4o"):
        self.model_id = model_id
        
        # Initialize LLM
        openai_api_key = os.environ.get("OPENAI_API_KEY")
        if not openai_api_key:
            logger.warning("OPENAI_API_KEY not found. Workflow will likely fail.")
        
        self.llm = OpenAIModel(
            client_args={"api_key": openai_api_key},
            model_id=self.model_id,
            params={"temperature": 0.1, "max_tokens": 2048},
        )
        
        # Initialize Tools
        self.tools = self._initialize_tools()
        self.tool_map = {getattr(t, "tool_name", getattr(t, "name", str(t))): t for t in self.tools}

    def _initialize_tools(self) -> List[Any]:
        """Initialize and return the list of available tools."""
        tools = [
            get_application_summary,
            explain_fraud_score,
            compare_to_population,
            check_fair_lending_flags,
            get_identity_network,
            get_model_performance,
        ]
        
        # dynamic import to avoid circular dependency and handle missing deps
        try:
            from app import init_confluence
            from confluence_ingestor.adapters.strands import (
                create_confluence_search_tool,
                create_confluence_loader_tool,
            )
            
            rag = init_confluence()
            if rag:
                tools.append(create_confluence_search_tool(rag=rag, k=5))
                tools.append(create_confluence_loader_tool(max_pages=10))
        except ImportError:
            logger.debug("Confluence integration not available (ImportError).")
        except Exception as e:
            logger.error(f"Failed to add Confluence tools: {e}")
        
        return tools


    async def run(self, input_text: str, context_messages: List[Dict[str, str]] = None) -> str:
        """
        Main entry point for the workflow.
        Executes the steps in order.
        """
        state: WorkflowState = {
            "input_text": input_text,
            "messages": context_messages or [],
            "intent": None,
            "tool_calls": [],
            "tool_outputs": [],
            "final_response": None,
            "error": None
        }

        try:
            logger.info(f"Starting workflow for: {input_text}")
            
            # Step 1: Analyze Intent & Plan Tools
            await self._analyze_intent_and_plan(state)
            
            # Step 2: Execute Tools
            await self._execute_tools(state)
            
            # Step 3: Generate Response
            await self._generate_response(state)
            
            return state["final_response"]

        except Exception as e:
            logger.error(f"Workflow execution failed: {e}", exc_info=True)
            return f"I encountered an error processing your request: {str(e)}"

    async def _call_llm(self, prompt: str) -> str:
        """Helper to call async LLM."""
        messages = [{"role": "user", "content": [{"text": prompt}]}]
        full_text = ""
        async for chunk in self.llm.stream(messages=messages):
            # Extract text from contentBlockDelta
            if "contentBlockDelta" in chunk:
                delta = chunk["contentBlockDelta"].get("delta", {})
                if "text" in delta:
                    full_text += delta["text"]
        return full_text


    async def _analyze_intent_and_plan(self, state: WorkflowState):
        """
        Determine the intent and decide which tools to call.
        """
        prompt = f"""
        You are a routing agent for a Fraud Explainability Assistant.
        Your goal is to analyze the user's request and determine which tools to call.

        User Request: "{state['input_text']}"

        Available Tools:
        - get_application_summary(application_id): Basic info about an application.
        - explain_fraud_score(application_id): Detailed SHAP explanations for score.
        - compare_to_population(application_id, comparison_group): Stats vs approved/denied.
        - check_fair_lending_flags(application_id): Compliance check.
        - get_identity_network(application_id): Linkage analysis.
        - get_model_performance(model_name, portfolio): Model metrics.
        - confluence_search(query): Search documentation/policies.
        - confluence_loader(space_key, page_title): Load full doc pages.

        Return a JSON object with:
        - "intent": Brief description of intent.
        - "tool_calls": List of objects with "tool_name" and "arguments" (dict).
        
        If no tool is needed (e.g., greeting), return empty tool_calls.
        """
        
        # Call LLM for planning
        response_text = (await self._call_llm(prompt)).strip()
        
        # Clean markdown code blocks if present
        if response_text.startswith("```json"):
            response_text = response_text[7:]
        if response_text.endswith("```"):
            response_text = response_text[:-3]
            
        try:
            plan = json.loads(response_text)
            state["intent"] = plan.get("intent", "Unknown")
            state["tool_calls"] = plan.get("tool_calls", [])
            logger.info(f"Intent: {state['intent']}, Tools: {len(state['tool_calls'])}")
        except json.JSONDecodeError:
            logger.error(f"Failed to parse plan JSON: {response_text}")
            state["error"] = "Failed to plan execution."

    async def _execute_tools(self, state: WorkflowState):
        """
        Execute the planned tools and store results.
        """
        for call in state["tool_calls"]:
            tool_name = call["tool_name"]
            args = call.get("arguments", {})
            
            if tool_name in self.tool_map:
                try:
                    tool_instance = self.tool_map[tool_name]
                    logger.info(f"Executing {tool_name} with {args}")
                    
                    # Support both generated tool classes and manual function tools
                    if hasattr(tool_instance, "__call__"):
                         # Check if tool is async
                         if asyncio.iscoroutinefunction(tool_instance):
                             result = await tool_instance(**args)
                         else:
                             result = tool_instance(**args)
                    else:
                         # Fallback if it's a strands Tool object (depends on implementation)
                         # This assumes the tool wrapper handles the call
                         pass

                    state["tool_outputs"].append({
                        "tool_name": tool_name,
                        "result": result
                    })
                except Exception as e:
                    logger.error(f"Tool {tool_name} failed: {e}")
                    state["tool_outputs"].append({
                        "tool_name": tool_name,
                        "error": str(e)
                    })
            else:
                logger.warning(f"Tool {tool_name} not found.")

    async def _generate_response(self, state: WorkflowState):
        """
        Synthesize the final answer using tool outputs.
        """
        context_str = ""
        for output in state["tool_outputs"]:
            if "error" in output:
                context_str += f"\n[Error from {output['tool_name']}]: {output['error']}\n"
            else:
                context_str += f"\n[Result from {output['tool_name']}]:\n{output['result']}\n"
        
        if not context_str and not state["tool_calls"]:
            context_str = "[No tools were called. Answer based on general knowledge or conversational context.]"

        prompt = f"""
        You are a Fraud Model Explainability Assistant.
        
        User Request: "{state['input_text']}"
        
        Context / Tool Outputs:
        {context_str}
        
        Please provide a comprehensive answer to the user's request.
        - Be precise and data-driven.
        - If multiple tools returned data, synthesize them into a coherent narrative.
        - Highlight risk factors and compliance notes if present.
        """
        
        response_text = await self._call_llm(prompt)
        state["final_response"] = response_text