auto_causal/agent.py

"""
LangChain agent for the auto_causal module.

This module configures a LangChain agent with specialized tools for causal inference,
allowing for an interactive approach to analyzing datasets and applying appropriate
causal inference methods.
"""

import logging
from typing import Dict, List, Any, Optional
from langchain.agents.react.agent import create_react_agent
from langchain.agents import AgentExecutor, create_structured_chat_agent, create_tool_calling_agent
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain_core.messages import SystemMessage, HumanMessage
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder
from langchain.tools import tool
# Import the callback handler
from langchain.callbacks.tracers.stdout import ConsoleCallbackHandler
# Import tool rendering utility
from langchain.tools.render import render_text_description
# Import LCEL components
from langchain.agents.format_scratchpad.tools import format_to_tool_messages
from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.language_models import BaseChatModel
from langchain_anthropic.chat_models import convert_to_anthropic_tool
import os
# Import actual tools from the tools directory
from auto_causal.tools.input_parser_tool import input_parser_tool
from auto_causal.tools.dataset_analyzer_tool import dataset_analyzer_tool
from auto_causal.tools.query_interpreter_tool import query_interpreter_tool
from auto_causal.tools.method_selector_tool import method_selector_tool
from auto_causal.tools.method_validator_tool import method_validator_tool
from auto_causal.tools.method_executor_tool import method_executor_tool
from auto_causal.tools.explanation_generator_tool import explanation_generator_tool
from auto_causal.tools.output_formatter_tool import output_formatter_tool
#from auto_causal.prompts import SYSTEM_PROMPT # Assuming SYSTEM_PROMPT is defined here or imported
from langchain_core.output_parsers import StrOutputParser
# Import the centralized factory function
from .config import get_llm_client 
#from .prompts import SYSTEM_PROMPT 
from langchain_core.messages import AIMessage, AIMessageChunk
import re
import json
from typing import Union
from langchain_core.output_parsers import BaseOutputParser
from langchain.schema import AgentAction, AgentFinish
from langchain_anthropic.output_parsers import ToolsOutputParser
from langchain.agents.react.output_parser import ReActOutputParser
from langchain.agents import AgentOutputParser
from langchain.agents.agent import AgentAction, AgentFinish, OutputParserException
import re
from typing import Union, List
from auto_causal.models import *

from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.exceptions import OutputParserException

from langchain.agents.agent import AgentOutputParser
from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS

FINAL_ANSWER_ACTION = "Final Answer:"
MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE = (
    "Invalid Format: Missing 'Action:' after 'Thought:'"
)
MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE = (
    "Invalid Format: Missing 'Action Input:' after 'Action:'"
)
FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE = (
    "Parsing LLM output produced both a final answer and parse-able actions"
)


class ReActMultiInputOutputParser(AgentOutputParser):
    """Parses ReAct-style output that may contain multiple tool calls."""

    def get_format_instructions(self) -> str:
        # You can reuse the original FORMAT_INSTRUCTIONS,
        # but let the model know it may emit multiple actions.
        return FORMAT_INSTRUCTIONS + (
            "\n\nIf you need to call more than one tool, simply repeat:\n"
            "Action: <tool_name>\n"
            "Action Input: <json or text>\n"
            "…for each tool in sequence."
        )

    @property
    def _type(self) -> str:
        return "react-multi-input"

    def parse(self, text: str) -> Union[List[AgentAction], AgentFinish]:
        includes_answer = FINAL_ANSWER_ACTION in text
        print('-------------------')
        print(text)
        print('-------------------')
        # Grab every Action / Action Input block
        pattern = (
            r"Action\s*\d*\s*:[\s]*(.*?)\s*"
            r"Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*?)(?=(?:Action\s*\d*\s*:|$))"
        )
        matches = list(re.finditer(pattern, text, re.DOTALL))

        # If we found tool calls…
        if matches:
            if includes_answer:
                # both a final answer *and* tool calls is ambiguous
                raise OutputParserException(
                    f"{FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE}: {text}"
                )

            actions: List[AgentAction] = []
            for m in matches:
                tool_name = m.group(1).strip()
                tool_input = m.group(2).strip().strip('"')
                print('\n--------------------------')
                print(tool_input)
                print('--------------------------')
                actions.append(AgentAction(tool_name, json.loads(tool_input), text))

            return actions

        # Otherwise, if there's a final answer, finish
        if includes_answer:
            answer = text.split(FINAL_ANSWER_ACTION, 1)[1].strip()
            return AgentFinish({"output": answer}, text)

        # No calls and no final answer → figure out which error to throw
        if not re.search(r"Action\s*\d*\s*Input\s*\d*:", text):
            raise OutputParserException(
                f"Could not parse LLM output: `{text}`",
                observation=MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE,
                llm_output=text,
                send_to_llm=True,
            )

        # Fallback
        raise OutputParserException(f"Could not parse LLM output: `{text}`")

# Set up basic logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def create_agent_prompt(tools: List[tool]) -> ChatPromptTemplate:
    """Create the prompt template for the causal inference agent, emphasizing workflow and data handoff.
       (This is the version required by the LCEL agent structure below)
    """
    # Get the tool descriptions
    tool_description = render_text_description(tools)
    tool_names = ", ".join([t.name for t in tools])

    # Define the system prompt template string
    system_template = """
You are a causal inference expert helping users answer causal questions by following a strict workflow using specialized tools.

Remember you always have to always generate the Thought, Action and Action Input block.
TOOLS:
------
You have access to the following tools:

{tools}

To use a tool, please use the following format:

Thought: Do I need to use a tool? Yes
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action, as a single, valid JSON object string. Check the tool definition for required arguments and structure.
Observation: the result of the action, often containing structured data like 'variables', 'dataset_analysis', 'method_info', etc.

When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:

Thought: Do I need to use a tool? No
Final Answer: [your response here]

DO NOT UNDER ANY CIRCUMSTANCE CALL MORE THAN ONE TOOL IN A  STEP

**IMPORTANT TOOL USAGE:**
1.  **Action Input Format:** The value for 'Action Input' MUST be a single, valid JSON object string. Do NOT include any other text or formatting around the JSON string.
2.  **Argument Gathering:** You MUST gather ALL required arguments for the Action Input JSON from the initial Human input AND the 'Observation' outputs of PREVIOUS steps. Look carefully at the required arguments for the tool you are calling.
3.  **Data Handoff:** The 'Observation' from a previous step often contains structured data needed by the next tool. For example, the 'variables' output from `query_interpreter_tool` contains fields like `treatment_variable`, `outcome_variable`, `covariates`, `time_variable`, `instrument_variable`, `running_variable`, `cutoff_value`, and `is_rct`. When calling `method_selector_tool`, you MUST construct its required `variables` input argument by including **ALL** these relevant fields identified by the `query_interpreter_tool` in the previous Observation. Similarly, pass the full `dataset_analysis`, `dataset_description`, and `original_query` when required by the next tool.

IMPORTANT WORKFLOW:
-------------------
You must follow this exact workflow, selecting the appropriate tool for each step:

1. ALWAYS start with `input_parser_tool` to understand the query
2. THEN use `dataset_analyzer_tool` to analyze the dataset
3. THEN use `query_interpreter_tool` to identify variables (output includes `variables` and `dataset_analysis`)
4. THEN use `method_selector_tool` (input requires `variables` and `dataset_analysis` from previous step)
5. THEN use `method_validator_tool` (input requires `method_info` and `variables` from previous step)
6. THEN use `method_executor_tool` (input requires `method`, `variables`, `dataset_path`)
7. THEN use `explanation_generator_tool` (input requires results, method_info, variables, etc.)
8. FINALLY use `output_formatter_tool` to return the results 

REASONING PROCESS:
------------------
EXPLICITLY REASON about:
1. What step you're currently on (based on previous tool's Observation)
2. Why you're selecting a particular tool (should follow the workflow)
3. How the output of the previous tool (especially structured data like `variables`, `dataset_analysis`, `method_info`) informs the inputs required for the current tool.

IMPORTANT RULES:
1. Do not make more than one tool call in a single step.
2. Do not include ``` in your output at all.
3. Don't use action names like default_api.dataset_analyzer_tool, instead use tool names like dataset_analyzer_tool.
4. Always start, action, and observation with a new line.
5. Don't use '\\' before double quotes
6. Don't include ```json for Action Input. Also ensure that Action Input is a valid json. DO no add any text after Action Iput.
7. You have to always choose one of the tools unless it's the final answer.
Begin!
""" 

    # Create the prompt template
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_template),
        MessagesPlaceholder("chat_history", optional=True), # Use MessagesPlaceholder
        # MessagesPlaceholder("agent_scratchpad"),  

        ("human", "{input}\n Thought:{agent_scratchpad}"),
        # ("ai", "{agent_scratchpad}"),
        # MessagesPlaceholder("agent_scratchpad" ), # Use MessagesPlaceholder
        # "agent_scratchpad"
    ])
    return prompt

def create_causal_agent(llm: BaseChatModel) -> AgentExecutor:
    """
    Create and configure the LangChain agent with causal inference tools.
    (Using explicit LCEL construction, compatible with shared LLM client)
    """
    # Define tools available to the agent
    agent_tools = [
        input_parser_tool,
        dataset_analyzer_tool,
        query_interpreter_tool,
        method_selector_tool,
        method_validator_tool,
        method_executor_tool,
        explanation_generator_tool,
        output_formatter_tool
    ]
    # anthropic_agent_tools = [ convert_to_anthropic_tool(anthropic_tool) for anthropic_tool in agent_tools]
    # Create the prompt using the helper
    prompt = create_agent_prompt(agent_tools)
    # Bind tools to the LLM (using the passed shared instance)
    
    
    # Create memory
    # Consider if memory needs to be passed in or created here
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

    # Manually construct the agent runnable using LCEL
    from langchain_anthropic.output_parsers import ToolsOutputParser
    from langchain.agents.output_parsers.json import JSONAgentOutputParser
    # from langchain.agents.react.output_parser import MultiActionAgentOutputParsers ReActMultiInputOutputParser
    provider = os.getenv("LLM_PROVIDER", "openai")
    if provider == "gemini":
        base_parser=ReActMultiInputOutputParser()
        llm_with_tools = llm.bind_tools(agent_tools)
    else:
        base_parser=ToolsAgentOutputParser()
        llm_with_tools = llm.bind_tools(agent_tools, tool_choice="any")
    agent = create_react_agent(llm_with_tools, agent_tools, prompt, output_parser=base_parser)
    
    
    # Create executor (should now work with the manually constructed agent)
    executor = AgentExecutor(
        agent=agent,
        tools=agent_tools,
        memory=memory, # Pass the memory object
        verbose=True,
        callbacks=[ConsoleCallbackHandler()], # Optional: for console debugging
        handle_parsing_errors=True, # Let AE handle parsing errors
        max_retries = 100
    )
    
    return executor

def run_causal_analysis(query: str, dataset_path: str, 
                        dataset_description: Optional[str] = None, 
                        api_key: Optional[str] = None) -> Dict[str, Any]:
    """
    Run causal analysis on a dataset based on a user query.
    
    Args:
        query: User's causal question
        dataset_path: Path to the dataset
        dataset_description: Optional textual description of the dataset
        api_key: Optional OpenAI API key (DEPRECATED - will be ignored)
        
    Returns:
        Dictionary containing the final formatted analysis results from the agent's last step.
    """
    # Log the start of the analysis
    logger.info("Starting causal analysis run...")
    
    try:
        # --- Instantiate the shared LLM client --- 
        model_name = os.getenv("LLM_MODEL", "gpt-4")
        if model_name in ['o3', 'o4-mini', 'o3-mini']:
            print('-------------------------')
            shared_llm = get_llm_client()
        else:
            shared_llm = get_llm_client(temperature=0) # Or read provider/model from env
        
        # --- Dependency Injection Note (REMAINS RELEVANT) --- 
        # If tools need the LLM, they must be adapted. Example using partial:
        # from functools import partial
        # from .components import input_parser 
        # # Assume input_parser.parse_input needs llm 
        # input_parser_tool_with_llm = tool(partial(input_parser.parse_input, llm=shared_llm)) 
        # Use input_parser_tool_with_llm in the tools list passed to the agent below.
        # Similar adjustments needed for decision_tree._recommend_ps_method if used.
        # --- End Note --- 

        # --- Create agent using the shared LLM --- 
        # agent_executor = create_causal_agent(shared_llm) 
        
        # Construct input, including description if available
        # IMPORTANT: Agent now expects 'input' and potentially 'chat_history'
        # The input needs to contain all initial info the first tool might need.
        input_text = f"My question is: {query}\n"
        input_text += f"The dataset is located at: {dataset_path}\n"
        if dataset_description:
            input_text += f"Dataset Description: {dataset_description}\n"
        input_text += "Please perform the causal analysis following the workflow."
        
        # Log the constructed input text
        logger.info(f"Constructed input for agent: \n{input_text}")

        input_parsing_result = input_parser_tool(input_text)
        dataset_analysis_result = dataset_analyzer_tool.func(dataset_path=input_parsing_result["dataset_path"], dataset_description=input_parsing_result["dataset_description"], original_query=input_parsing_result["original_query"]).analysis_results
        query_info = QueryInfo(
        query_text=input_parsing_result["original_query"],
        potential_treatments=input_parsing_result["extracted_variables"].get("treatment"),
        potential_outcomes=input_parsing_result["extracted_variables"].get("outcome"),
        covariates_hints=input_parsing_result["extracted_variables"].get("covariates_mentioned"),
        instrument_hints=input_parsing_result["extracted_variables"].get("instruments_mentioned")
    )

        query_interpreter_output = query_interpreter_tool.func(query_info=query_info, dataset_analysis=dataset_analysis_result, dataset_description=input_parsing_result["dataset_description"], original_query = input_parsing_result["original_query"]).variables
        method_selector_output = method_selector_tool.func(variables=query_interpreter_output,
            dataset_analysis=dataset_analysis_result,
            dataset_description=input_parsing_result["dataset_description"],
            original_query = input_parsing_result["original_query"],
            excluded_methods=None)
        method_info = MethodInfo(
            **method_selector_output['method_info']
        )
        method_validator_input = MethodValidatorInput(
            method_info=method_info,
            variables=query_interpreter_output,
            dataset_analysis=dataset_analysis_result,
            dataset_description=input_parsing_result["dataset_description"],
            original_query = input_parsing_result["original_query"]
        )
        method_validator_output = method_validator_tool.func(method_validator_input)
        method_executor_input = MethodExecutorInput(
            **method_validator_output
        )
        method_executor_output = method_executor_tool.func(method_executor_input, original_query = input_parsing_result["original_query"])

        explainer_output = explanation_generator_tool.func(            method_info=method_info,
            validation_info=method_validator_output,
            variables=query_interpreter_output,
            results=method_executor_output,
            dataset_analysis=dataset_analysis_result,
            dataset_description=input_parsing_result["dataset_description"],
            original_query = input_parsing_result["original_query"])
        result = explainer_output
        result['results']['results']["method_used"] = method_validator_output['method']
        logger.info(result)
        logger.info("Causal analysis run finished.")
        
        # Ensure result is a dict and extract the 'output' part
        if isinstance(result, dict):
            final_output = result
            if isinstance(final_output, dict):
                return final_output # Return only the dictionary from the final tool
            else:
                logger.error(f"Agent result['output'] was not a dictionary: {type(final_output)}. Returning error dict.")
                return {"error": "Agent did not produce the expected dictionary output in the 'output' key.", "raw_agent_result": result}
        else:
            logger.error(f"Agent returned non-dict type: {type(result)}. Returning error dict.")
            return {"error": "Agent did not return expected dictionary output.", "raw_output": str(result)}

    except ValueError as e:
        logger.error(f"Configuration Error: {e}")
        # Return an error dictionary in case of exception too
        return {"error": f"Error: Configuration issue - {e}"} # Ensure consistent error return type
    except Exception as e:
        logger.error(f"An unexpected error occurred during causal analysis: {e}", exc_info=True)
        # Return an error dictionary in case of exception too
        return {"error": f"An unexpected error occurred: {e}"}