function_app.py

import azure.functions as func
from agents import Agent, Runner, function_tool, WebSearchTool, FileSearchTool
import os
import json
import requests
from pathlib import Path
import logging

DATA_PATH = Path(__file__).parent / "data" / "team_stats_data_model.json"
try:
    with DATA_PATH.open("r", encoding="utf-8") as f:
        DATA_MODEL = json.load(f)  # cached in memory for warm runs
except Exception as e:
    DATA_MODEL = None

app = func.FunctionApp()




# Databricks config via environment
DATABRICKS_HOST = os.environ.get("DATABRICKS_HOST", "").rstrip("/")
DATABRICKS_TOKEN = os.environ.get("DATABRICKS_TOKEN", "")
DATABRICKS_WAREHOUSE_ID = os.environ.get("DATABRICKS_WAREHOUSE_ID", "")

AGENT_INSTRUCTIONS = f"""
You are an agent that is an expert in NFL stats,
Your role is to answer user questions by running SQL queries on Databricks.
You must understand the user's intent:
1. If the question is about rules, handoff to the rules agent
2. If the question is about stats, generate the appropriate SQL query to answer the question
3. Write a short answer to the question, don't explain your thinking or thought process, be concise and use charts to answer

- You can only answer using the tool call output from the query_databricks tool.
- Do not modify the table names or column names, don't try to add prefixes like dbo.
Here are the tables you have access to:

{DATA_MODEL}
"""

@function_tool
def query_databricks(sql: str) -> str:
    """Run a SQL query on Databricks and return JSON rows.

    Args:
        sql: The SQL statement to execute.
        wait_timeout: How long to wait for results (e.g. "40s").

    """
    print(f"Running SQL query: {sql}")
    if not DATABRICKS_HOST or not DATABRICKS_TOKEN:
        return (
            "Databricks not configured. Set DATABRICKS_HOST and DATABRICKS_TOKEN (and optionally DATABRICKS_WAREHOUSE_ID)."
        )

    body: dict[str, object] = {
        "statement": sql,
        "disposition": "INLINE",
        "format": "JSON_ARRAY",
        "wait_timeout": "30s",
        "on_wait_timeout": "CANCEL",
        "catalog": "nfl",
        "schema": "default",
    }
    if DATABRICKS_WAREHOUSE_ID:
        body["warehouse_id"] = DATABRICKS_WAREHOUSE_ID

    headers = {"Authorization": f"Bearer {DATABRICKS_TOKEN}", "Content-Type": "application/json"}
    try:
        r = requests.post(
            f"{DATABRICKS_HOST}/api/2.0/sql/statements/", headers=headers, json=body, timeout=60
        )
        r.raise_for_status()
        data = r.json()
    except Exception as e:
        return f"Databricks request failed: {e}"

    result = data.get("result") if isinstance(data, dict) else None
    if not isinstance(result, dict):
        # Return raw payload if unexpected
        return json.dumps(data)

    manifest = result.get("manifest") if isinstance(result, dict) else None
    cols = []
    if isinstance(manifest, dict):
        columns = manifest.get("columns")
        if isinstance(columns, list):
            cols = [c.get("name") for c in columns if isinstance(c, dict)]

    rows = result.get("data_array") if isinstance(result, dict) else None
    if isinstance(rows, list) and cols:
        records = [dict(zip(cols, row)) for row in rows if isinstance(row, list)]
        # Limit to avoid huge payloads
        return json.dumps(records[:500])

    # Fallback: return whatever we got
    return json.dumps(result)


@app.function_name(name="AgentInvoke")
@app.route(route="agent_invoke", methods=["GET", "POST"], auth_level=func.AuthLevel.ANONYMOUS)
async def agent_invoke(req: func.HttpRequest) -> func.HttpResponse:
    user_input = req.params.get("user_input")
    if not user_input:
        try:
            body = req.get_json()
            user_input = body.get("user_input") if isinstance(body, dict) else None
        except ValueError:
            user_input = None

    if not user_input:
        return func.HttpResponse(
            "Missing user_input. Send as query '?user_input=...' or JSON body {\"user_input\": \"...\"}.",
            status_code=400,
        )
    try:
        orchestrator_agent = Agent(
            name="Orchestrator Agent",
            model="gpt-4.1",
            instructions="Your role is to understand the user's intent and handoff to the appropriate agent.",
            tools=[],
        )

        rules_agent = Agent(
            name="NFL Rulebook Agent",
            model="gpt-5-nano",
            instructions="You are a helpful assistant that is an expert in the fules of Football, you must use the tool to search the rules of Football.",
            tools=[FileSearchTool(vector_store_ids=["vs_68c99c9b04ac81919b8ba98c907e2cc9"], include_search_results=True)],
        )

        stats_agent = Agent(
            name="NFL Stats Agent",
            model="gpt-5-nano",
            instructions=AGENT_INSTRUCTIONS,
            tools=[query_databricks],
        )
        orchestrator_agent.handoffs = [rules_agent, stats_agent]

        result = await Runner.run(orchestrator_agent, user_input)
        last_agent = result.last_agent.name
        return func.HttpResponse(
            json.dumps({
                "last_agent": last_agent,
                "final_output": result.final_output or ""
            }),
            mimetype="application/json"
        )
    except Exception as e:
        logging.exception("Unhandled error in agent_invoke")
        return func.HttpResponse(
            json.dumps({
                "error": "Function execution failed",
                "message": str(e)
            }),
            status_code=500,
            mimetype="application/json"
        )