hra commited on
Commit
7254b09
Β·
verified Β·
1 Parent(s): 302a595

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +88 -0
app.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py – MCP wrapper for triggering a Databricks spam check job
2
+
3
+ import os
4
+ import time
5
+ import json
6
+ import requests
7
+ import gradio as gr
8
+
9
+ # ── Environment setup ─────────────────────────────────────────
10
+ DATABRICKS_HOST = os.getenv("DATABRICKS_HOST")
11
+ DATABRICKS_TOKEN = os.getenv("DATABRICKS_TOKEN")
12
+ DATABRICKS_NOTEBOOK_PATH = os.getenv("DATABRICKS_NOTEBOOK_PATH")
13
+ EXISTING_CLUSTER_ID = os.getenv("EXISTING_CLUSTER_ID")
14
+
15
+ HEADERS = {
16
+ "Authorization": f"Bearer {DATABRICKS_TOKEN}",
17
+ "Content-Type": "application/json"
18
+ }
19
+
20
+ # ── MCP Tool ──────────────────────────────────────────────────
21
+ def run_spam_check(phone_number: str):
22
+ # 1. Submit the job to Databricks
23
+ submit_payload = {
24
+ "run_name": f"FraudCheck_{phone_number}",
25
+ "notebook_task": {
26
+ "notebook_path": DATABRICKS_NOTEBOOK_PATH,
27
+ "base_parameters": {
28
+ "phone_number": phone_number
29
+ }
30
+ },
31
+ "existing_cluster_id": EXISTING_CLUSTER_ID
32
+ }
33
+
34
+ response = requests.post(
35
+ f"{DATABRICKS_HOST}/api/2.1/jobs/runs/submit",
36
+ headers=HEADERS,
37
+ json=submit_payload
38
+ )
39
+
40
+ if response.status_code != 200:
41
+ return {"error": "Failed to start job", "details": response.text}
42
+
43
+ run_id = response.json()["run_id"]
44
+
45
+ # 2. Poll for completion
46
+ while True:
47
+ status_response = requests.get(
48
+ f"{DATABRICKS_HOST}/api/2.1/jobs/runs/get?run_id={run_id}",
49
+ headers=HEADERS
50
+ )
51
+ status_json = status_response.json()
52
+ run_state = status_json["state"]["life_cycle_state"]
53
+ if run_state in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
54
+ break
55
+ time.sleep(2)
56
+
57
+ # 3. Get output
58
+ output_response = requests.get(
59
+ f"{DATABRICKS_HOST}/api/2.1/jobs/runs/get-output?run_id={run_id}",
60
+ headers=HEADERS
61
+ )
62
+
63
+ if output_response.status_code != 200:
64
+ return {"error": "Failed to fetch output", "details": output_response.text}
65
+
66
+ output_json = output_response.json()
67
+ notebook_output = output_json.get("notebook_output", {}).get("result")
68
+
69
+ if notebook_output:
70
+ try:
71
+ return json.loads(notebook_output) if isinstance(notebook_output, str) else notebook_output
72
+ except json.JSONDecodeError:
73
+ return {"raw_output": notebook_output}
74
+ else:
75
+ return {"warning": "No result returned in notebook output."}
76
+
77
+ # ── Minimal UI to test ────────────────────────────────────────
78
+ demo = gr.Interface(
79
+ fn=run_spam_check,
80
+ inputs=gr.Textbox(label="Phone Number to Check"),
81
+ outputs="json",
82
+ title="Databricks Spam Check",
83
+ description="Triggers a Databricks job to check for spam/fraud using a phone number.",
84
+ api_name="spam_check"
85
+ )
86
+
87
+ if __name__ == "__main__":
88
+ demo.launch(mcp_server=True)