src/agents/tools.py

"""
tools.py
--------
Custom Tools for the SPE GCS 2026 ML Challenge Agents.

1. StatefulPythonExecutionTool: Safely executes generated Pandas code, keeping state.
2. IADC_SearchTool: Queries the local IADC ChromaDB for drilling concepts.
3. VolveHistory_SearchTool: Queries the Volve DDR ChromaDB for historical events.
"""
import os
import io
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.express as px
import plotly.graph_objects as go
from pathlib import Path
from crewai.tools import BaseTool
from pydantic import Field
from langchain_chroma import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from dotenv import load_dotenv

load_dotenv()

BASE_DIR = Path(__file__).resolve().parents[2]
DATA_DIR = BASE_DIR / "data" / "processed"
DDR_DIR = str(DATA_DIR / "ddr")
WITSML_DIR = str(DATA_DIR / "witsml")
OUTPUTS_DIR = BASE_DIR / "outputs" / "figures"
OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)

# ── 1. Stateful Python Execution Tool ─────────────────────────────────────────

class DataInventoryTool(BaseTool):
    name: str = "data_inventory_inspector"
    description: str = "Use this tool to see what data (CSVs, WITSML, Files) are available across DDR_DIR and WITSML_DIR. Returns a summary of wells and datasets."

    def _run(self, query: str = "") -> str:
        summary = ["### Project Data Inventory"]
        
        # DDR Directory
        ddr_path = os.environ.get('DDR_DIR', DDR_DIR)
        if os.path.exists(ddr_path):
            files = os.listdir(ddr_path)
            summary.append(f"\n**DDR Directory ({ddr_path}):**")
            summary.append(f"- Total Files: {len(files)}")
            csvs = [f for f in files if f.endswith('.csv')]
            wells = set([f.split('_activities')[0].split('_daily')[0] for f in csvs if '_' in f])
            summary.append(f"- Detected Wells: {', '.join(sorted(list(wells))[:10])}...")
            if '_ddr_extraction_summary.csv' in files:
                summary.append("- [Key File]: `_ddr_extraction_summary.csv` (High-level well metadata)")
            if '_ddr_all_activities.csv' in files:
                summary.append("- [Key File]: `_ddr_all_activities.csv` (Granular time-log across all wells)")
        
        # WITSML Directory
        witsml_path = os.environ.get('WITSML_DIR', WITSML_DIR)
        if os.path.exists(witsml_path):
            wells_witsml = [d for d in os.listdir(witsml_path) if os.path.isdir(os.path.join(witsml_path, d))]
            summary.append(f"\n**WITSML Directory ({witsml_path}):**")
            summary.append(f"- Well Folders: {', '.join(wells_witsml)}")
            
        # PDF Reports
        pdf_path = "data/raw/Reports"
        if os.path.exists(pdf_path):
            pdfs = [f for f in os.listdir(pdf_path) if f.endswith('.pdf')]
            summary.append(f"\n**PDF Knowledge Source:**")
            summary.append(f"- Reports: {', '.join(pdfs)}")

        return "\n".join(summary)

def save_plotly_html(fig, filename_without_ext):
    """Helper to be passed to the agent REPL so it can easily save html."""
    path = os.path.join(str(OUTPUTS_DIR), f"{filename_without_ext}.html")
    fig.write_html(path)
    print(f"Interactive Plotly chart saved to: {path}")

# Global REPL state so variables persist between tool calls in the same run
_repl_globals = {
    "pd": pd,
    "plt": plt,
    "np": np,
    "px": px,
    "go": go,
    "os": os,
    "DDR_DIR": DDR_DIR,
    "WITSML_DIR": WITSML_DIR,
    "OUTPUTS_DIR": str(OUTPUTS_DIR),
    "DataInventory": DataInventoryTool(),
    "save_plotly_html": save_plotly_html
}

class StatefulPythonExecutionTool(BaseTool):
    name: str = "Python REPL Data Analyst"
    description: str = (
        "Execute Python code (especially Pandas, and Plotly) to analyze data. "
        "Variables defined here PERSIST between calls. "
        "You have access to Plotly via `px` (plotly.express) and `go` (plotly.graph_objects). "
        "IMPORTANT FOR VISUALIZATIONS: Use Plotly instead of Matplotlib whenever possible. "
        "After creating a Plotly figure `fig`, save it using the provided helper: `save_plotly_html(fig, 'my_chart_name')`. "
        "Always use `print()` or `print(df.to_markdown())` to output the results so you can read them. "
        "Truncate massive outputs; do not print DataFrames with >50 rows."
    )
    
    def _run(self, code: str) -> str:
        # Strip markdown code block formatting if present
        code = code.strip()
        if code.startswith("```python"):
            code = code[9:]
        elif code.startswith("```"):
            code = code[3:]
        if code.endswith("```"):
            code = code[:-3]
        code = code.strip()

        old_stdout = sys.stdout
        redirected_output = sys.stdout = io.StringIO()
        
        try:
            exec(code, getattr(sys.modules[__name__], '_repl_globals'))
        except Exception as e:
            sys.stdout = old_stdout
            return f"Error executing code:\n{e}"
        finally:
            sys.stdout = old_stdout
            
        output = redirected_output.getvalue()
        
        # Hard limits on output size to protect the LLM context window
        if not output.strip():
            return "Code executed successfully, but nothing was printed. Please `print()` the result to see it."
        
        if len(output) > 8000:
            return output[:8000] + "\n\n... [OUTPUT TRUNCATED: Result exceeded 8000 characters. Please refine your code to print smaller summaries.]"
            
        return output

# ── 2. Vector Search Tools ─────────────────────────────────────────────────────

# Lazy singletons for the two vector databases
_iadc_db = None
_volve_db = None
_embeddings = None
EMBEDDING_MODEL = "models/gemini-embedding-2-preview"

def get_embeddings():
    global _embeddings
    if _embeddings is None:
        api_key = os.environ.get("GOOGLE_API_KEY")
        _embeddings = GoogleGenerativeAIEmbeddings(
            model=EMBEDDING_MODEL,
            google_api_key=api_key
        )
    return _embeddings

def get_iadc_db():
    global _iadc_db
    if _iadc_db is None:
        db_path = BASE_DIR / "data" / "viking_context" / "chroma_fallback"
        _iadc_db = Chroma(persist_directory=str(db_path), embedding_function=get_embeddings())
    return _iadc_db

def get_volve_db():
    global _volve_db
    if _volve_db is None:
        db_path = BASE_DIR / "data" / "viking_context" / "chroma_fallback"
        _volve_db = Chroma(persist_directory=str(db_path), embedding_function=get_embeddings())
    return _volve_db

class IADC_SearchTool(BaseTool):
    name: str = "Drilling Knowledge (IADC) DB Search"
    description: str = (
        "Search the IADC drilling glossary and general Wikipedia technical articles. "
        "Use this for DEFINITIONS and THEORY (e.g. 'What is a BHA?', 'What causes stuck pipe?'). "
        "Do NOT use this for specific Volve well events."
    )
    
    def _run(self, query: str) -> str:
        try:
            db = get_iadc_db()
            # OpenViking namespace filter
            viking_filter = {"viking_namespace": "resources/iadc/"} 
            results = db.similarity_search(query, k=3, filter=viking_filter)
            if not results:
                return "No relevant IADC information found in OpenViking context."
            output = []
            for i, doc in enumerate(results):
                source = doc.metadata.get('source', 'Unknown')
                # Clean up path to just file name
                if isinstance(source, str) and '/' in source:
                    source = source.split('/')[-1]
                output.append(f"[Source: {source}]: {doc.page_content}")
            return "\n\n".join(output)
        except Exception as e:
            return f"Error searching IADC DB: {e}"

class VolveHistory_SearchTool(BaseTool):
    name: str = "Volve Campaign History DB Search"
    description: str = (
        "Search the historical Daily Drilling Reports (DDR) from the Volve campaign. "
        "Use this for HISTORICAL EVENTS and EQUIPMENTS (e.g. 'What BHA components failed on well 15/9-F-1 C?', 'Find instances of stuck pipe', 'Motor performance'). "
        "Do NOT use this for general definitions."
    )
    
    def _run(self, query: str) -> str:
        try:
            # 1. Semantic Search (OpenViking L2 Overview via Gemini 2)
            db = get_volve_db()
            viking_filter = {"viking_namespace": "resources/volve/"} 
            results = db.similarity_search(query, k=10, filter=viking_filter)
            
            output = []
            seen_content = set()
            
            # Identify high-value keywords for fallback (OpenViking L0 Hybrid Glob logic)
            keywords = ["whipstock", "milling", "stuck", "fishing", "loss", "kick", "cement", "casing", "liner", "window", "weather", "heave", "bha", "assembly", "motor", "mwd", "lwd", "bit", "failure", "twist off"]
            query_keywords = [k for k in keywords if k in query.lower()]

            # 2. Keyword Fallback: If no results or if specific keywords were missed
            found_keywords = False
            for doc in results:
                for k in query_keywords:
                    if k.upper() in doc.page_content.upper():
                        found_keywords = True
                        break
            
            # If we didn't find specific matches, try a literal scan of the narratives CSV
            if not found_keywords and query_keywords:
                csv_path = BASE_DIR / "data" / "processed" / "serialized_text" / "ddr_narratives.csv"
                if csv_path.exists():
                    import pandas as pd
                    df = pd.read_csv(csv_path)
                    # Simple keyword filter
                    mask = df['text'].str.lower().str.contains('|'.join(query_keywords), na=False)
                    kw_results = df[mask].tail(10) # Get latest 10 matches
                    if not kw_results.empty:
                        for idx, row in kw_results.iterrows():
                            content = row['text']
                            if content not in seen_content:
                                output.append(f"[Volve-KeywordMatch]:\n{content}")
                                seen_content.add(content)

            # Add semantic results (avoiding duplicates)
            for i, doc in enumerate(results):
                if doc.page_content not in seen_content:
                    source = doc.metadata.get('source', 'Unknown source')
                    if isinstance(source, str) and '/' in source:
                        source = source.split('/')[-1]
                    output.append(f"[Source: {source}]:\n{doc.page_content}")
                    seen_content.add(doc.page_content)

            if not output:
                return "No historical Volve events found matching this query."
            
            result_str = "\n\n---\n\n".join(output)
            if len(result_str) > 12000:
                return result_str[:12000] + "\n...[TRUNCATED]"
            return result_str
            
        except Exception as e:
            return f"Error searching Volve History DB: {e}"