app.py

import gradio as gr
from smolagents import CodeAgent, Tool, InferenceClientModel
import pandas as pd
import requests
import os
import spaces

# ------------------------------------------------------------------------
# 1. Define Tools (Class-Based Style)
# ------------------------------------------------------------------------

class NYCEnergyGradeTool(Tool):
    name = "get_building_energy_grade"
    description = "Retrieves the Local Law 84 Energy Efficiency Letter Grade for a specific NYC Building (BIN). Use this when asked about compliance, grades, or energy scores."
    inputs = {
        "bin_number": {
            "type": "string",
            "description": "The 7-digit NYC Building Identification Number (BIN)."
        }
    }
    output_type = "string"

    def forward(self, bin_number: str) -> str:
        # NYC Open Data Endpoint for LL84 (2023 Data)
        endpoint = "https://data.cityofnewyork.us/resource/gpwd-npar.json"
        
        params = {
            "$limit": 1,
            "nyc_building_identification": str(bin_number).strip()
        }
        
        try:
            response = requests.get(endpoint, params=params)
            data = response.json()
            
            if not data:
                return f"No Local Law 84 data found for BIN {bin_number}."
                
            record = data[0]
            grade = record.get("energy_efficiency_grade", "N/A")
            score = record.get("energy_star_score", "N/A")
            address = record.get("property_name", "Unknown Address")
            
            return f"**NYC Data for BIN {bin_number}:**\n- Address: {address}\n- Energy Grade: {grade}\n- Energy Star Score: {score}"
            
        except Exception as e:
            return f"Error fetching NYC Open Data: {str(e)}"

class SeniorLogSearchTool(Tool):
    name = "search_senior_logs"
    description = "Searches the internal 'tribal knowledge' logs of Senior Engineers. Use this for specific equipment issues (e.g., 'pump', 'chiller', 'tenant complaints')."
    inputs = {
        "query": {
            "type": "string",
            "description": "The keyword to search for in the logs."
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        results = []
        try:
            # In a real app, this would query a Vector DB (FAISS/Chroma)
            with open("knowledge_base.txt", "r") as f:
                lines = f.readlines()
                for line in lines:
                    if query.lower() in line.lower():
                        results.append(line.strip())
            
            if not results:
                return "No specific notes found in the Senior Logs for this issue."
            
            return "Found in Senior Engineer Logs:\n" + "\n".join(results)
        except FileNotFoundError:
            return "System Error: knowledge_base.txt not found."

# ------------------------------------------------------------------------
# 2. Configure the Model & Agent
# ------------------------------------------------------------------------

# Initialize Model using InferenceClientModel (Serverless/Inference API)
model = InferenceClientModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    token=os.environ.get("HF_TOKEN")
)

# Initialize Agent
agent = CodeAgent(
    tools=[NYCEnergyGradeTool(), SeniorLogSearchTool()],
    model=model,
    additional_authorized_imports=["pandas", "requests", "datetime"],
    description="You are 'Chief Joe', a Digital Twin of a Senior Facility Engineer. You utilize both hard data (NYC Codes) and soft data (Senior Logs) to solve building issues."
)

# ------------------------------------------------------------------------
# 3. Define GPU Wrappers (Gradio Logic)
# ------------------------------------------------------------------------

@spaces.GPU(duration=120)
def diagnose_issue(issue_description, bin_input):
    """
    Main diagnostic function. 
    It combines the user's issue description and optional BIN number into a prompt.
    """
    
    # Construct the prompt for the agent
    prompt = f"""
    You are acting as Senior Engineer 'Chief Joe'.
    
    User Issue: "{issue_description}"
    Building BIN: "{bin_input if bin_input else 'Not provided'}"

    INSTRUCTIONS:
    1. If the issue is about equipment (pumps, chillers, boilers), use 'search_senior_logs' to see if this is a known quirks.
    2. If the user provided a BIN and mentions compliance or energy, use 'get_building_energy_grade'.
    3. Synthesize the findings into a helpful, authoritative response. 
       - If you found a log entry, cite it: "According to my logs from Jan 2024..."
       - If you found NYC data, cite it: "City records show..."
    """
    
    try:
        response = agent.run(prompt)
        return response
    except Exception as e:
        return f"System Error: {str(e)}"

# ------------------------------------------------------------------------
# 4. The UI
# ------------------------------------------------------------------------

custom_css = """
body { background-color: #f4f4f9; }
.gradio-container { font-family: 'Roboto', sans-serif; }
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, title="Senior Engineer Digital Twin") as demo:
    gr.Markdown("# 👷‍♂️ Chief Joe: Senior Facility Engineer Digital Twin")
    gr.Markdown("""
    **Concept:** An AI Agent that bridges the gap between "Tribal Knowledge" (Senior Logs) and "Public Data" (NYC Open Data).
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            issue_input = gr.Textbox(
                label="Describe the Issue", 
                placeholder="e.g., 'The domestic water pump is short cycling again.'",
                lines=3
            )
            bin_input = gr.Textbox(
                label="NYC BIN (Optional)", 
                placeholder="e.g., 1088888"
            )
            submit_btn = gr.Button("Ask Chief Joe", variant="primary")
            
        with gr.Column(scale=2):
            output_box = gr.Markdown(label="Chief Joe's Diagnosis")

    # Example Buttons to demo capabilities
    gr.Markdown("### Try these scenarios:")
    with gr.Row():
        ex1 = gr.Button("Scenario 1: Pump Maintenance (Internal Logs)")
        ex2 = gr.Button("Scenario 2: LL97 Compliance (Open Data)")
        ex3 = gr.Button("Scenario 3: Tenant Complaints (Hybrid)")

    # Event Handlers
    submit_btn.click(
        fn=diagnose_issue, 
        inputs=[issue_input, bin_input], 
        outputs=output_box
    )

    # Example Click Handlers
    ex1.click(
        fn=diagnose_issue,
        inputs=[gr.Textbox(value="The domestic water pump is short cycling.", visible=False), gr.Textbox(value="", visible=False)],
        outputs=output_box
    )
    ex2.click(
        fn=diagnose_issue,
        inputs=[gr.Textbox(value="Check our Local Law 84 energy grade.", visible=False), gr.Textbox(value="1000000", visible=False)],
        outputs=output_box
    )
    ex3.click(
        fn=diagnose_issue,
        inputs=[gr.Textbox(value="Tenant on 14th floor complaining about heat.", visible=False), gr.Textbox(value="", visible=False)],
        outputs=output_box
    )

if __name__ == "__main__":
    demo.launch()