File size: 11,800 Bytes
3e30d53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 | import pandas as pd
import numpy as np
import plotly.express as px
import json
import logging
import re
from typing import TypedDict, Dict, Any, List
from langgraph.graph import StateGraph, END
from langchain_google_genai import ChatGoogleGenerativeAI
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class AnalysisState(TypedDict):
dataframe: pd.DataFrame
dataset_info: Dict[str, Any]
insights: str
visualizations: List[Dict[str, Any]]
charts: List[Any]
class DataAnalysisAgent:
def __init__(self, llm: ChatGoogleGenerativeAI):
self.llm = llm
self.workflow = self._create_workflow()
def _create_workflow(self):
"""Creates the graph workflow for the data analysis sub-agent."""
workflow = StateGraph(AnalysisState)
workflow.add_node("data_profiler", self._profile_dataset)
# This new node will generate insights AND plan visualizations in one LLM call
workflow.add_node("insight_and_viz_planner", self._generate_insights_and_plan_visualizations)
workflow.add_node("chart_creator", self._create_charts)
workflow.add_edge("data_profiler", "insight_and_viz_planner")
workflow.add_edge("insight_and_viz_planner", "chart_creator")
workflow.add_edge("chart_creator", END)
workflow.set_entry_point("data_profiler")
return workflow.compile()
def _profile_dataset(self, state: AnalysisState):
"""Profiles the dataset to understand its structure for the LLM."""
logger.info("--- 📊 (Sub-Agent) Profiling Data ---")
df_for_profiling = state["dataframe"].copy().reset_index()
profile = {
"shape": df_for_profiling.shape,
"columns": list(df_for_profiling.columns),
"dtypes": {col: str(dtype) for col, dtype in df_for_profiling.dtypes.to_dict().items()},
"numeric_columns": df_for_profiling.select_dtypes(include=[np.number]).columns.tolist(),
"datetime_columns": df_for_profiling.select_dtypes(include=['datetime64']).columns.tolist()
}
logger.info(" Data profile created.")
return {"dataset_info": profile}
def _generate_insights_and_plan_visualizations(self, state: AnalysisState):
"""Generates key insights and plans visualizations in a single LLM call."""
logger.info("--- 🧠 (Sub-Agent) Generating Insights & Visualization Plan ---")
info = state["dataset_info"]
datetime_col = info.get("datetime_columns", [None])[0] or info.get("columns", ["index"])[0]
prompt = f"""
You are an expert financial data scientist. Based on the following data profile from a time-series stock dataset,
generate key insights and plan effective visualizations.
Data Profile: {json.dumps(info, indent=2)}
Instructions:
Your response MUST be ONLY a single valid JSON object. Do not include any other text or markdown.
The JSON object must have two keys: "insights" and "visualizations".
- "insights": A list of 3-5 concise, bullet-point style strings focusing on trends, correlations, and anomalies.
- "visualizations": A list of 3 JSON objects, each planning a chart.
- Plan a line chart for the 'close' price over time using the '{datetime_col}' column.
- Plan a histogram for the 'volume' column.
- Plan one other relevant chart (e.g., scatter plot, bar chart).
Example Response:
{{
"insights": [
"The closing price shows a significant upward trend over the period.",
"Trading volume spiked on dates corresponding to major news events.",
"There is a strong positive correlation between opening and closing prices."
],
"visualizations": [
{{"type": "line", "columns": ["{datetime_col}", "close"], "title": "Closing Price Over Time"}},
{{"type": "histogram", "columns": ["volume"], "title": "Trading Volume Distribution"}},
{{"type": "scatter", "columns": ["open", "close"], "title": "Opening vs. Closing Price"}}
]
}}
"""
response_str = self.llm.invoke(prompt).content
logger.info(f" LLM raw output for insights & viz plan:\n{response_str}")
try:
json_match = re.search(r'\{.*\}', response_str, re.DOTALL)
if not json_match:
raise ValueError("No JSON object found in the LLM response.")
clean_json_str = json_match.group(0)
response_json = json.loads(clean_json_str)
insights_list = response_json.get("insights", [])
insights_str = "\n".join(f"* {insight}" for insight in insights_list)
viz_plan = response_json.get("visualizations", [])
logger.info(" Successfully parsed insights and viz plan.")
return {"insights": insights_str, "visualizations": viz_plan}
except (json.JSONDecodeError, ValueError) as e:
logger.error(f"Failed to parse insights and visualization plan from LLM. Error: {e}")
logger.info(" Using a default visualization plan as a fallback.")
default_plan = [
{"type": "line", "columns": [datetime_col, "close"], "title": "Closing Price Over Time (Default)"},
{"type": "histogram", "columns": ["volume"], "title": "Trading Volume (Default)"}
]
return {"insights": "Analysis generated, but detailed insights could not be parsed.", "visualizations": default_plan}
def _create_charts(self, state: AnalysisState):
"""Creates Plotly charts - HARDCODED for reliability."""
logger.info("--- 🎨 (Sub-Agent) Creating Charts ---")
# 1. Prepare DataFrame
df = state["dataframe"].copy()
if df.index.name in ['timestamp', 'date', 'datetime', 'index']:
df = df.reset_index()
# Normalize column names to lowercase
df.columns = [str(c).lower() for c in df.columns]
charts = []
# Find X-axis column (timestamp)
x_col = None
for candidate in ['timestamp', 'date', 'datetime', 'index']:
if candidate in df.columns:
x_col = candidate
break
if not x_col:
logger.warning(" No timestamp column found. Skipping charts.")
return {"charts": []}
# --- CHART 1: Price History (Line) ---
if 'close' in df.columns:
try:
logger.info(f" Generating Price Chart (x={x_col}, y=close)")
fig = px.line(df, x=x_col, y='close',
title="📈 Price History",
template="plotly_dark",
labels={'close': 'Price ($)', x_col: 'Time'})
fig.update_traces(line_color='#00ff41')
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate price chart: {e}")
# --- CHART 2: Volume (Bar) ---
if 'volume' in df.columns:
try:
logger.info(f" Generating Volume Chart (x={x_col}, y=volume)")
fig = px.bar(df, x=x_col, y='volume',
title="📊 Trading Volume",
template="plotly_dark",
labels={'volume': 'Volume', x_col: 'Time'})
fig.update_traces(marker_color='#ff6b35')
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate volume chart: {e}")
# --- CHART 3: Price vs Volume (Scatter) ---
if 'close' in df.columns and 'volume' in df.columns:
try:
logger.info(" Generating Price vs Volume Scatter Plot")
fig = px.scatter(df, x='volume', y='close',
title="🔍 Price vs Volume Correlation",
template="plotly_dark",
labels={'volume': 'Trading Volume', 'close': 'Price ($)'},
trendline="ols", # Add regression line
opacity=0.6)
fig.update_traces(marker=dict(size=8, color='#4ecdc4'))
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate scatter plot: {e}")
# --- CHART 4: Daily Returns Histogram ---
if 'close' in df.columns and len(df) > 1:
try:
logger.info(" Generating Daily Returns Histogram")
# Calculate returns
df['returns'] = df['close'].pct_change() * 100
df_returns = df.dropna(subset=['returns'])
if not df_returns.empty:
fig = px.histogram(df_returns, x='returns',
nbins=30,
title="📊 Daily Returns Distribution",
template="plotly_dark",
labels={'returns': 'Daily Return (%)'},
color_discrete_sequence=['#9b59b6'])
fig.add_vline(x=0, line_dash="dash", line_color="white",
annotation_text="Zero Return", annotation_position="top")
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate histogram: {e}")
# --- CHART 5: Box Plot (Price Distribution) ---
if 'close' in df.columns:
try:
logger.info(" Generating Box Plot")
fig = px.box(df, y='close',
title="📦 Price Distribution (Box Plot)",
template="plotly_dark",
labels={'close': 'Price ($)'},
color_discrete_sequence=['#a29bfe'])
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate box plot: {e}")
# --- CHART 6: Violin Plot (Volume Distribution) ---
if 'volume' in df.columns:
try:
logger.info(" Generating Violin Plot")
fig = px.violin(df, y='volume',
title="🎻 Volume Distribution (Violin Plot)",
template="plotly_dark",
labels={'volume': 'Trading Volume'},
color_discrete_sequence=['#74b9ff'],
box=True, # Show box plot inside violin
points='all') # Show all data points
charts.append(fig)
except Exception as e:
logger.error(f" Failed to generate violin plot: {e}")
logger.info(f" Successfully created {len(charts)} charts.")
return {"charts": charts}
def run_analysis(self, dataframe: pd.DataFrame):
"""Runs the full analysis workflow on the given DataFrame."""
if dataframe.empty:
logger.warning("Input DataFrame is empty. Skipping analysis.")
return {"insights": "No data available for analysis.", "charts": []}
initial_state = {"dataframe": dataframe}
# The final state will now contain insights and charts after the workflow runs
final_state = self.workflow.invoke(initial_state)
return final_state |