import time from langgraph.types import Command from langgraph.graph import END from llm import llm_model from pydantic import BaseModel,Field from typing import List,Literal, Optional from langchain_core.messages import SystemMessage, HumanMessage, AIMessage from Prompts import supervisor_prompt,PreprocessingPlanner_prompt,cleaner_prompt,Reporter_prompt,VISUALIZATION_PROMPT from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain.agents import AgentExecutor, create_tool_calling_agent from Toolkit.Tools import python_repl_ast,eda_fact_sheet from Guardrails.Preprocessing import StructuredPlanOutput from Guardrails.cleaner import CleaningSummary from Guardrails.report import BusinessReport from Guardrails.visualizer import VisualizationReport from langchain.output_parsers import PydanticOutputParser,OutputFixingParser class Router(BaseModel): next: Literal["PreprocessingPlanner_node","Cleaner_node","Reporter_node","visualizer_node",END]= Field(description="The next node to route to. Must be one of the available nodes.") reasoning: str = Field(description="A short reasoning for the decision made.") class AgentStateModel(BaseModel): messages: Optional[List] = None Instructions: Optional[str] = None Analysis: Optional[List[dict]] = None clean: Optional[List[dict]] = None Report: Optional[List[dict]] = None Visualizations: Optional[List[dict]] = None Path: Optional[str] = None next: Optional[str] = None current_reasoning: Optional[str] = None class DataAnalystAgent: def __init__(self): self.llm_model = llm_model def supervisor_node(self,state:AgentStateModel) -> Command[Literal["PreprocessingPlanner_node","Cleaner_node","Reporter_node","visualizer_node", END]]: """ The central router of the workflow. It evaluates the current state and the last message to decide the next action. This node is designed to be highly token-efficient by creating a lean summary of the state instead of passing the full, verbose state objects to the LLM. """ print("**************************below is my state right after entering****************************") print(state) print("************************** SUPERVISOR: EVALUATING STATE ****************************") state_summary = ( f"Current Workflow Status:\n" f"- Analysis Plan Generated: {'Yes' if state.Analysis else 'No'}\n" f"- Cleaning Plan Generated: {'Yes' if state.clean else 'No'}\n" f"- Report Generated: {'Yes' if state.Report else 'No'}\n" f"- Visualizations Generated: {'Yes' if state.Visualizations else 'No'}\n" ) messages_for_llm = [ SystemMessage(content=supervisor_prompt), HumanMessage(content=state_summary), ] if state.messages: last_message = state.messages[-1] # Add a prefix to clearly label the last message for the LLM last_message_content = f"Last Event:\nThe last node to run was '{last_message.name}'. It reported the following:\n---\n{last_message.content}\n---" messages_for_llm.append(HumanMessage(content=last_message_content)) print(f"--- Attaching last event from '{last_message.name}' ---") else: # Handle the very first run where there are no messages messages_for_llm.append(HumanMessage(content="Last Event: None. This is the first step of the workflow.")) messages_for_this_attempt = list(messages_for_llm) print("***********************Invoking LLM for routing decision************************") parser = PydanticOutputParser(pydantic_object=Router) fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=self.llm_model) # Build the retry chain chain = self.llm_model | fixing_parser # Add retries max_attempts = 3 attempt = 0 error_msg = None response = None while attempt < max_attempts: attempt += 1 print(f"--- Attempt {attempt} ---") # Compose messages for this attempt messages_for_this_attempt = list(messages_for_llm) if error_msg: # Inject previous error info to let LLM know what failed messages_for_this_attempt.append(HumanMessage(content=f"Previous attempt failed due to: {error_msg}. Please follow the schema strictly: {Router.model_json_schema()}")) try: response = chain.invoke(messages_for_this_attempt) break except Exception as e: error_msg = str(e) print(f"--- Error on attempt {attempt}: {error_msg} ---") # If last attempt, will exit loop and propagate error if response is None: # All retries failed, fallback error fallback_msg = f"All {max_attempts} attempts failed. Last error: {error_msg}" print(f"--- Supervisor node failed ---\n{fallback_msg}") return Command( goto="END", update={ "next": "END", "current_reasoning": fallback_msg } ) goto = response.next print("********************************this is my goto*************************") print(goto) print("********************************") print(response.reasoning) if goto == "END": goto = END print("**************************below is my state****************************") print(state) return Command(goto=goto, update={'next': goto, 'current_reasoning': response.reasoning} ) def PreprocessingPlanner_node(self, state: AgentStateModel) -> Command[Literal['supervisor']]: print("*****************called PreprocessingPlanner node************") Instructions = state.Instructions parser = PydanticOutputParser(pydantic_object=StructuredPlanOutput) fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=self.llm_model) task_prompt = ( f"Find the instructions given by the user here : {Instructions} and follow this {PreprocessingPlanner_prompt} to the letter.modify in this path:{state.Path}" ) print(f"--- Sending this direct task to the agent ---\n{task_prompt}\n---------------------------------------------") system_prompt = ChatPromptTemplate.from_messages([ ("system", "You are a DataFrame analyzer. Your primary tool is `eda_fact_sheet`. " "First, call the tool to get data insights. Then, based on the tool's output, " "provide a final answer formatted as a JSON object containing the preprocessing plan and summaries."), ("human", "{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ]) Analyzer_agent = create_tool_calling_agent( llm=self.llm_model, tools=[eda_fact_sheet], prompt=system_prompt ) agent_executor = AgentExecutor( agent=Analyzer_agent, tools=[eda_fact_sheet], verbose=True, handle_parsing_errors=True, return_intermediate_steps=True ) # 5. Wrap execution in a retry loop max_attempts = 3 attempt = 0 while attempt < max_attempts: attempt += 1 try: result = agent_executor.invoke({ "input": task_prompt, }) # Try parsing the final output final_output_string = result.get("output", "") parsed_output: StructuredPlanOutput = fixing_parser.parse(final_output_string) # Successfully parsed → extract plan and summary plan_dict = {"plan": [step.model_dump() for step in parsed_output.plan]} summary_str = f"{parsed_output.summary}\n{parsed_output.details}" # Update state and return return Command( update={ "messages": [ AIMessage(content=summary_str, name="PreprocessingPlanner_node") ], "Analysis": [{"final_answer": plan_dict}] }, goto="supervisor", ) except Exception as e: error_msg = ( f"Attempt {attempt} failed due to error: {str(e)}. " f"Please strictly follow the schema: {StructuredPlanOutput.model_json_schema()}" ) print(f"--- Runtime/Parsing error encountered ---\n{error_msg}") # Inject error into prompt for next retry # Use an f-string to properly embed all variables task_prompt = f"The previous attempt failed with this error: {error_msg}. Please correct your tool usage and try again. Here is the original task:\n---\n{task_prompt}" # If all attempts fail, fallback to supervisor with error message return Command( update={ "messages": [ AIMessage(content="Error: The analysis agent failed to produce a valid preprocessing plan after multiple attempts.", name="Analyzer_node_Error") ], "Analysis": [{"error": "Parsing failed after retries"}] }, goto="supervisor", ) def Cleaner_node(self, state: AgentStateModel) -> Command[Literal['supervisor']]: print("*****************called cleaner node************") Path = state.Path preprocessing_plan = state.Analysis[0]['final_answer']['plan'] batched_plan = [] current_batch = [] for column_action in preprocessing_plan: # Add the current item to the batch current_batch.append(column_action) # If the batch is now full, add it to our final list and reset it if len(current_batch) == 4: batched_plan.append(current_batch) current_batch = [] # After the loop, check if there are any leftover items in the last batch if current_batch: batched_plan.append(current_batch) parser = PydanticOutputParser(pydantic_object=CleaningSummary) fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=self.llm_model) system_prompt = ChatPromptTemplate.from_messages( [( "system", "Follow the instructions here : {cleaner_prompt} and in the input to the letter and make the necessary changes to the dataframe." ), ("human", "{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ] ) Cleaner_agent = create_tool_calling_agent( llm=self.llm_model, tools=[python_repl_ast], prompt=system_prompt ) agent_executor = AgentExecutor( agent=Cleaner_agent, tools=[python_repl_ast], verbose=True, # handle_parsing_errors=True, # return_intermediate_steps=True ) # 3. Loop Through Batches and Invoke the Agent for Each all_batch_results = [] final_clean_outputs = [] for i, batch in enumerate(batched_plan, start=1): clean_plan_str = str(batch) task_prompt = ( f"Apply the following cleaning plan (batch {i} of {len(batched_plan)}) to the dataset at path: {Path}\n" f"Plan details:\n{clean_plan_str}" ) print(f"--- Sending task for Batch {i} to the agent ---\n{task_prompt}\n---------------------------------------------") max_attempts = 3 attempt = 0 batch_success = False while attempt < max_attempts: attempt += 1 try: result = agent_executor.invoke({"input": task_prompt,"cleaner_prompt":cleaner_prompt,"agent_scratchpad": [] }) final_output_string = result.get("output", "") # Try parsing the output parsed_output: CleaningSummary = fixing_parser.parse(final_output_string) # On success, store results and break the retry loop all_batch_results.append(parsed_output) final_clean_outputs.append({"final_answer": final_output_string}) batch_success = True print(f"--- Batch {i} successful on attempt {attempt} ---") time.sleep(10) break except Exception as e: error_msg = ( f"Attempt {attempt} for batch {i} failed with an error: {str(e)}. " "You MUST provide a final answer that is a valid JSON object. Please review the plan and strictly follow the schema." ) print(f"--- Runtime/Parsing error for Batch {i} ---\n{error_msg}") # Inject error context for the next retry attempt on this specific batch task_prompt = f"Your previous attempt failed with this error: {error_msg}\n\nPlease re-execute the original plan:\n{task_prompt}" # If a batch fails after all retries, exit the entire node with an error if not batch_success: error_message = f"Error: The cleaner agent failed to process batch {i} after {max_attempts} attempts." return Command( update={ "messages": [AIMessage(content=error_message, name="Cleaner_node_Error")], "clean": [{"error": f"Processing failed at batch {i}"}] }, goto="supervisor", ) final_summary = "All cleaning batches completed successfully.\n\n" for idx, summary in enumerate(all_batch_results, start=1): final_summary += f"--- Batch {idx} Summary ---\nSummary: {summary.summary}\nDetails: {summary.details}\n\n" return Command( update={ "messages": [AIMessage(content=final_summary.strip(), name="cleaner_node")], "clean": final_clean_outputs, }, goto="supervisor", ) def Reporter_node(self, state: AgentStateModel) -> Command[Literal['supervisor']]: print("*****************called Reporter node************") Instructions = state.Instructions df_path = state.Path # --- STEP 1: Perform Reconnaissance MANUALLY (and only once) --- print("--- Reporter: Performing initial data reconnaissance with eda_fact_sheet ---") try: recon_result_str = str(eda_fact_sheet.run(path=df_path)) except Exception as e: return Command(update={"messages": [AIMessage(content=f"Error during initial data recon: {e}", name="Reporter_node_Error")]}, goto="supervisor") print(f"--- Reporter: Condensing {len(recon_result_str)} characters of context... ---") condensation_prompt = ChatPromptTemplate.from_messages([ ("system", "You are a data analysis assistant. Summarize the following verbose JSON data profile into a concise, human-readable format for another AI agent to use. Focus on column names, data types, and key stats."), ("human", "Please summarize this data profile:\n\n{profile}") ]) summarizer_chain = condensation_prompt | self.llm_model condensed_summary = summarizer_chain.invoke({"profile": recon_result_str}).content print(f"--- Reporter: Condensed Summary Created ---") # 1. Instantiate the parser for our new structured output parser = PydanticOutputParser(pydantic_object=BusinessReport) fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=self.llm_model) task_prompt = (f"User Instructions: {Instructions}\n\nHere is a condensed summary of the dataset: \n---\n{condensed_summary}\n---\n\nNow, follow your main instructions: {Reporter_prompt}") print(f"--- Reporter: Sending condensed task to the main analysis agent ---") print(f"--- Sending this direct task to the agent ---\n{task_prompt}\n---------------------------------------------") system_prompt = ChatPromptTemplate.from_messages([ ("system", "You are a Business Intelligence consultant. You have access to `eda_fact_sheet(df_path)` for initial recon and `python_repl_ast(query)` for deep-dive analysis. " "The CSV is at this path: {df_path}. " "Your mission is to autonomously analyze the data and produce a strategic business report. "), ("human", "{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ]) Reporter_agent = create_tool_calling_agent( llm=self.llm_model, tools=[python_repl_ast], prompt=system_prompt ) agent_executor = AgentExecutor( agent=Reporter_agent, tools=[python_repl_ast], verbose=True, # handle_parsing_errors=True, # return_intermediate_steps=True ) # 2. Wrap execution in a retry loop for robust parsing max_attempts = 3 attempt = 0 while attempt < max_attempts: attempt += 1 try: result = agent_executor.invoke({ "input": task_prompt, "df_path": df_path, }) final_output_string = result.get("output", "") parsed_output: BusinessReport = fixing_parser.parse(final_output_string) summary_str = f"{parsed_output.subject}\n\n{parsed_output.executive_summary}" return Command( update={ "messages": [AIMessage(content=summary_str, name="Reporter_node")], "Report": [{"final_answer": parsed_output.model_dump()}] }, goto="supervisor", ) except Exception as e: error_msg = ( f"Attempt {attempt} failed due to a parsing error: {str(e)}. " f"You MUST provide a final answer that is a valid JSON object. Please strictly follow this schema: " f"{BusinessReport.model_json_schema()}" ) print(f"--- Runtime/Parsing error encountered ---\n{error_msg}") task_prompt =f"here is the error from the previous execution: {error_msg} so process the workflow accordingly" + task_prompt return Command( update={ "messages": [ AIMessage(content="Error: The reporter agent failed to produce a valid business report after multiple attempts.", name="Reporter_node_Error") ], "Report": [{"error": "Parsing failed after all retries"}] }, goto="supervisor", ) def visualizer_node(self, state: AgentStateModel) -> Command[Literal['supervisor']]: """ This node directs an agent to perform EDA and generate 10 business-focused visualizations. It enforces a structured JSON output for the final report and includes retry logic for parsing. """ print("***************** called Visualizer node ************") df_path = state.Path print("--- Visualizer: Performing initial data reconnaissance... ---") try: recon_result_str = str(eda_fact_sheet.run(path=df_path)) except Exception as e: return Command(update={"messages": [AIMessage(content=f"Error during initial data recon: {e}", name="Visualizer_node_Error")]}, goto="supervisor") print(f"--- Visualizer: Condensing {len(recon_result_str)} characters of context... ---") condensation_prompt = ChatPromptTemplate.from_messages([ ("system", "You are a data analysis assistant. Summarize the following verbose JSON data profile into a concise, human-readable format for another AI agent to use for creating visualizations."), ("human", "Please summarize this data profile:\n\n{profile}") ]) summarizer_chain = condensation_prompt | self.llm_model condensed_summary = summarizer_chain.invoke({"profile": recon_result_str}).content print(f"--- Visualizer: Condensed Summary Created ---") parser = PydanticOutputParser(pydantic_object=VisualizationReport) fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=self.llm_model) task_prompt = (f"Based on this data summary:\n---\n{condensed_summary}\n---\n\nNow, follow your main instructions to create visualizations: {VISUALIZATION_PROMPT}") print(f"--- Visualizer: Sending condensed task to the plotting agent ---") print(f"--- Sending this direct task to the agent ---\n{task_prompt}\n---------------------------------------------") system_prompt = ChatPromptTemplate.from_messages([ ("system", "You are a Data Visualization Specialist. You have access to `eda_fact_sheet(df_path)` for initial recon and `python_repl_ast(query)` for generating and saving plots. " "The CSV is at this path: {df_path}. " "Your mission is to autonomously analyze the data and produce a series of 10 visualizations as instructed. "), ("human", "{input}"), MessagesPlaceholder(variable_name="agent_scratchpad") ]) visualizer_agent = create_tool_calling_agent( llm=self.llm_model, tools=[python_repl_ast], prompt=system_prompt ) agent_executor = AgentExecutor( agent=visualizer_agent, tools=[python_repl_ast], verbose=True, # handle_parsing_errors=True, # return_intermediate_steps=True ) max_attempts = 3 attempt = 0 while attempt < max_attempts: attempt += 1 try: result = agent_executor.invoke({ "input": task_prompt, "df_path": df_path, }) final_output_string = result.get("output", "") parsed_output: VisualizationReport = fixing_parser.parse(final_output_string) summary_str = f"Successfully generated a report with {len(parsed_output.visualizations)} visualizations." return Command( update={ "messages": [AIMessage(content=summary_str, name="Visualizer_node")], "Visualizations": [{"final_answer": parsed_output.model_dump()}] }, goto="supervisor", ) except Exception as e: error_msg = ( f"Attempt {attempt} failed due to a parsing error: {str(e)}. " f"You MUST provide a final answer that is a valid JSON object. Please strictly follow this schema: " f"{VisualizationReport.model_json_schema()}" ) print(f"--- Runtime/Parsing error encountered ---\n{error_msg}") task_prompt = f"here is the error from the last execution: {error_msg} so process the workflow accordingly" + "\n\n" + task_prompt return Command( update={ "messages": [ AIMessage(content="Error: The visualizer agent failed to produce a valid report after multiple attempts.", name="Visualizer_node_Error") ], "Visualizations": [{"error": "Parsing failed after all retries"}] }, goto="supervisor", ) "Data is about sales,provide the data overview along with the preprocessing steps needed to perform EDA" """ Path = r"D:\Code Assistant\superstore sales.csv" D:\Code Assistant\superstore sales.csv """