Update app.py

app.py

@@ -1,70 +1,217 @@
-# app.py
-import os
-import pandas as pd
 import gradio as gr
+import pandas as pd
+import numpy as np
+import os
+import statsmodels.api as sm
+from io import StringIO
 
-from tools.sql_tool import SQLTool
-from tools.ts_preprocess import build_timeseries
+# --- LangChain Imports ---
+from langchain_groq import ChatGroq
+from langchain.agents import AgentExecutor, create_tool_calling_agent
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import tool
+from langchain_core.messages import SystemMessage
 
-DUCKDB_PATH = os.getenv("DUCKDB_PATH", "alm.duckdb")
+# --- ASSUMPTION ---
+# Assuming you have a file named 'sql_tools.py' in the same directory
+# with your pre-built and decorated @tool functions.
+try:
+    from sql_tools import run_duckdb_query, get_table_schema
+except ImportError:
+    print("WARNING: Could not import from 'sql_tools.py'.")
+    print("Using placeholder functions. Please create 'sql_tools.py'.")
+    
+    # Create placeholder tools if the file is missing, so the app can start
+    @tool
+    def run_duckdb_query(query: str) -> str:
+        """
+        [PLACEHOLDER] Runs a read-only SQL query.
+        Please create sql_tools.py to implement this.
+        """
+        if "schema" in query.lower() or "describe" in query.lower():
+            return "report_date DATE, portfolio_id VARCHAR, sector VARCHAR, market_value_usd DOUBLE"
+        return "Error: 'sql_tools.py' not found. This is a placeholder."
 
-sql_tool = SQLTool(DUCKDB_PATH)
-RESOLVED_PATH = sql_tool.get_full_table_path()  # e.g., my_db.main.masterdataset_v or main.masterdataset_v
+    @tool
+    def get_table_schema(table_name: str = "positions") -> str:
+        """
+        [PLACEHOLDER] Returns the schema for the 'positions' table.
+        Please create sql_tools.py to implement this.
+        """
+        return "report_date DATE, portfolio_id VARCHAR, sector VARCHAR, market_value_usd DOUBLE"
 
-INTRO = f"""
-### ALM LLM — Demo
 
-Connected to **DuckDB** at `{DUCKDB_PATH}`  
-Using table **{RESOLVED_PATH}** (auto-resolved).
-"""
+# --- Agent Tools ---
+# These tools perform analysis on data *after* it has been fetched.
 
-def run_nl(nl_query: str):
-    if not nl_query or not nl_query.strip():
-        return pd.DataFrame(), "", "Please enter a query.", pd.DataFrame(), pd.DataFrame()
+@tool
+def calculate_summary_statistics_from_data(data_string: str, column: str) -> str:
+    """
+    Calculates summary statistics (mean, median, std, min, max) for a specific
+    'column' from a 'data_string'.
+    'data_string' should be the string output from the `run_duckdb_query` tool.
+    """
     try:
-        df, sql, why = sql_tool.query_from_nl(nl_query)
+        # Convert the string data back into a DataFrame
+        data_df = pd.read_csv(StringIO(data_string.strip()), delim_whitespace=True, header=0)
+        
+        # HACK: The string output might have an extra index column, let's find the real columns
+        if column not in data_df.columns:
+            # Try reading again, assuming first column is an unnamed index
+            data_df = pd.read_csv(StringIO(data_string.strip()), delim_whitespace=True, header=0, index_col=0)
+            if column not in data_df.columns:
+                 return f"Error: Column '{column}' not found in data."
+
+        stats = {
+            "column": column,
+            "mean": data_df[column].mean(),
+            "median": data_df[column].median(),
+            "std_dev": data_df[column].std(),
+            "min": data_df[column].min(),
+            "max": data_df[column].max(),
+            "count": data_df[column].count()
+        }
+        return str(stats)
     except Exception as e:
-        return pd.DataFrame(), "", f"Error: {e}", pd.DataFrame(), pd.DataFrame()
-    try:
-        cf, gap = build_timeseries(df)
-    except Exception:
-        cf, gap = pd.DataFrame(), pd.DataFrame()
-    return df, sql.strip(), why, cf, gap
-
-def run_sql(sql_text: str):
-    if not sql_text or not sql_text.strip():
-        return pd.DataFrame(), "Please paste a SQL statement.", pd.DataFrame(), pd.DataFrame()
+        return f"Error in calculate_summary_statistics: {e}. Data input was: '{data_string[:200]}...'"
+
+@tool
+def perform_arima_forecast_from_data(data_string: str, time_column: str, value_column: str, forecast_periods: int) -> str:
+    """
+    Performs an ARIMA(1,1,1) forecast on a 'data_string'.
+    'data_string': The string output from `run_duckdb_query`.
+    'time_column': The name of the date/time column in the data.
+    'value_column': The name of the numerical column to forecast.
+    'forecast_periods': The number of periods (e.g., days) to forecast.
+    
+    The data MUST be ordered by the time_column before being passed to this tool.
+    """
     try:
-        df = sql_tool.run_sql(sql_text)
+        # Convert the string data back into a DataFrame
+        data_df = pd.read_csv(StringIO(data_string.strip()), delim_whitespace=True, header=0)
+
+        # HACK: The string output might have an extra index column
+        if time_column not in data_df.columns:
+            data_df = pd.read_csv(StringIO(data_string.strip()), delim_whitespace=True, header=0, index_col=0)
+            if time_column not in data_df.columns:
+                return f"Error: Time column '{time_column}' not found in data."
+        
+        if value_column not in data_df.columns:
+            return f"Error: Value column '{value_column}' not found in data."
+
+        if data_df.empty:
+            return "Error: Query returned no data."
+
+        # Prepare data for statsmodels
+        data_df[time_column] = pd.to_datetime(data_df[time_column])
+        data_df = data_df.set_index(time_column)
+        data_df = data_df.asfreq('D') # Ensure daily frequency, fill gaps if any
+        data_df[value_column] = data_df[value_column].fillna(method='ffill')
+        
+        model = sm.tsa.ARIMA(data_df[value_column], order=(1, 1, 1))
+        results = model.fit()
+        forecast = results.forecast(steps=forecast_periods)
+        
+        forecast_df = pd.DataFrame({
+            'date': forecast.index.strftime('%Y-%m-%d'),
+            'forecasted_value': forecast.values
+        })
+        
+        return f"Forecast successful. Last historical value was {data_df[value_column].iloc[-1]:.2f}.\nForecast:\n{forecast_df.to_string()}"
+    
     except Exception as e:
-        return pd.DataFrame(), f"Error: {e}", pd.DataFrame(), pd.DataFrame()
-    try:
-        cf, gap = build_timeseries(df)
-    except Exception:
-        cf, gap = pd.DataFrame(), pd.DataFrame()
-    return df, "OK", cf, gap
-
-with gr.Blocks(title="ALM LLM") as demo:
-    gr.Markdown(INTRO)
-
-    with gr.Tab("Ask in Natural Language"):
-        nl = gr.Textbox(label="Ask a question", placeholder="e.g., show me the top 10 fds by portfolio value", lines=2)
-        btn = gr.Button("Run")
-        sql_out = gr.Textbox(label="Generated SQL", interactive=False)
-        why_out = gr.Textbox(label="Reasoning", interactive=False)
-        df_out  = gr.Dataframe(label="Query Result", interactive=True)
-        cf_out  = gr.Dataframe(label="Projected Cash-Flows (if applicable)", interactive=True)
-        gap_out = gr.Dataframe(label="Liquidity Gap (monthly)", interactive=True)
-        btn.click(fn=run_nl, inputs=[nl], outputs=[df_out, sql_out, why_out, cf_out, gap_out])
-
-    with gr.Tab("Run Raw SQL"):
-        sql_in = gr.Code(label="SQL", language="sql", value=f"SELECT * FROM {RESOLVED_PATH} LIMIT 20;")
-        btn2 = gr.Button("Execute")
-        df2   = gr.Dataframe(label="Result", interactive=True)
-        status = gr.Textbox(label="Status", interactive=False)
-        cf2    = gr.Dataframe(label="Projected Cash-Flows (if applicable)", interactive=True)
-        gap2   = gr.Dataframe(label="Liquidity Gap (monthly)", interactive=True)
-        btn2.click(fn=run_sql, inputs=[sql_in], outputs=[df2, status, cf2, gap2])
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+        return f"Error in perform_arima_forecast: {e}. Data input was: '{data_string[:200]}...'"
+
+# --- Main Agent and UI Setup ---
+
+# Check for the GROQ_API_KEY in Hugging Face Space Secrets
+if "GROQ_API_KEY" not in os.environ:
+    print("GROQ_API_KEY not found in secrets!")
+    def missing_key_error(message, history):
+        return "Error: `GROQ_API_KEY` is not set in this Space's Secrets. Please add it to use the app."
+    
+    gr.ChatInterface(
+        missing_key_error,
+        title="Agentic Portfolio Analyst",
+        description="Error: GROQ_API_KEY secret is missing."
+    ).launch()
+
+else:
+    print("GROQ_API_KEY found. Initializing agent...")
+    llm = ChatGroq(model_name="llama-3.3-70b-versatile")
+    
+    # 2. Collect all our tools (imported and local)
+    tools = [
+        run_duckdb_query,
+        get_table_schema,
+        calculate_summary_statistics_from_data,
+        perform_arima_forecast_from_data
+    ]
+
+    # 3. Create the Agent Prompt
+    system_prompt = """
+    You are an expert portfolio analyst. You have access to SQL tools and analysis tools.
+    
+    Your logic MUST follow these steps:
+    1.  Use `get_table_schema` to understand the data.
+    2.  Use `run_duckdb_query` to fetch the raw data you need.
+    3.  If analysis (statistics or forecasting) is needed, take the string output
+        from `run_duckdb_query` and pass it *directly* to either
+        `calculate_summary_statistics_from_data` or `perform_arima_forecast_from_data`.
+        
+    Example for forecasting:
+    1.  Call `run_duckdb_query("SELECT report_date, SUM(market_value_usd) AS total_value FROM positions WHERE sector = 'Tech' GROUP BY report_date ORDER BY report_date")`.
+    2.  Get the result string: " report_date total_value \n 2024-01-01 100000.0 \n 2024-01-02 100500.0 \n ..."
+    3.  Call `perform_arima_forecast_from_data(data_string=" report_date total_value \n 2024-01-01 100000.0 \n ...", time_column="report_date", value_column="total_value", forecast_periods=30)`.
+    
+    Answer the user's request based on the final tool output.
+    """
+    
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            SystemMessage(content=system_prompt),
+            ("placeholder", "{chat_history}"),
+            ("human", "{input}"),
+            ("placeholder", "{agent_scratchpad}"),
+        ]
+    )
+
+    # 4. Create the Agent
+    agent = create_tool_calling_agent(llm, tools, prompt)
+
+    # 5. Create the Agent Executor
+    agent_executor = AgentExecutor(
+        agent=agent,
+        tools=tools,
+        verbose=True
+    )
+
+    # 6. Define the function for Gradio
+    def run_agent(message, history):
+        chat_history = []
+        for human_msg, ai_msg in history:
+            chat_history.append(("human", human_msg))
+            chat_history.append(("ai", ai_msg))
+        
+        try:
+            response = agent_executor.invoke({
+                "input": message,
+                "chat_history": chat_history
+            })
+            return response["output"]
+        except Exception as e:
+            return f"An error occurred: {e}"
+
+    # 7. Launch the Gradio App
+    gr.ChatInterface(
+        run_agent,
+        title="Agentic Portfolio Analyst",
+        description="Ask me questions about your portfolio. (This app uses imported SQL tools).",
+        examples=[
+            "What is the schema of the positions table?",
+            "What's the total market value by sector on the last available date?",
+            "Give me summary statistics for the 'Tech' sector's market value from portfolio P-123. Use the 'market_value_usd' column for stats.",
+            "What is the 30-day forecast for the total market value of portfolio P-123? Use 'total_value' for the forecast value column."
+        ]
+    ).launch()
+