1+ import os
2+ import json
3+ import requests
4+ from flask import Flask , render_template , request , jsonify
5+ from databricks .sdk .core import Config
6+ from databricks .sdk import WorkspaceClient
7+ from databricks .sdk .service .serving import ChatMessage , ChatMessageRole
8+
9+ # Initialize Flask app
10+ app = Flask (__name__ , static_folder = 'static' , template_folder = 'static' )
11+
12+ cfg = Config ()
13+ client = WorkspaceClient ()
14+
15+ # Model serving endpoint name
16+ ENDPOINT_NAME = "agents_ws_vfc_demo-sch_vfc_demo-fashion_merchandising_agent"
17+
18+ @app .route ('/' )
19+ def index ():
20+ """Serve the main HTML page"""
21+ return render_template ('index.html' )
22+
23+ @app .route ('/api/chat' , methods = ['POST' ])
24+ def chat ():
25+ """Handle chat requests and forward to model serving endpoint using Databricks SDK"""
26+ try :
27+ # Get the message history from the request
28+ data = request .get_json ()
29+ messages = data .get ('messages' , [])
30+
31+ if not messages :
32+ return jsonify ({'error' : 'No messages provided' }), 400
33+
34+ # Convert messages to ChatMessage objects
35+ chat_messages = []
36+ for msg in messages :
37+ role = msg .get ('role' , 'user' )
38+ content = msg .get ('content' , '' )
39+
40+ # Map roles to ChatMessageRole enum
41+ if role == 'user' :
42+ chat_role = ChatMessageRole .USER
43+ elif role == 'assistant' :
44+ chat_role = ChatMessageRole .ASSISTANT
45+ elif role == 'system' :
46+ chat_role = ChatMessageRole .SYSTEM
47+ else :
48+ chat_role = ChatMessageRole .USER # Default to user
49+
50+ chat_messages .append (ChatMessage (role = chat_role , content = content ))
51+
52+ # Get user information for logging
53+ user_email = request .headers .get ('X-Forwarded-Email' )
54+ app .logger .info (f"Making request to model endpoint for user: { user_email } )
55+
56+ # Make request to model serving endpoint using Databricks SDK
57+ response = client .serving_endpoints .query (
58+ name = ENDPOINT_NAME ,
59+ messages = chat_messages
60+ )
61+
62+ # Extract the response content
63+ if response and hasattr (response , 'choices' ) and response .choices :
64+ # Get the first choice
65+ choice = response .choices [0 ]
66+ if hasattr (choice , 'message' ) and choice .message :
67+ result_content = choice .message .content
68+ return jsonify ({'content' : result_content })
69+ else :
70+ return jsonify ({'error' : 'No message content in response' }), 500
71+ else :
72+ return jsonify ({'error' : 'No response choices received' }), 500
73+
74+ except Exception as e :
75+ error_msg = str (e )
76+ app .logger .error (f"Error calling model endpoint: { error_msg } )
77+
78+ # Handle specific error types
79+ if "authentication" in error_msg .lower () or "unauthorized" in error_msg .lower ():
80+ return jsonify ({
81+ 'error' : 'Authentication failed. Your Databricks token may have expired.' ,
82+ 'details' : 'Please refresh the page or log in again.'
83+ }), 401
84+ elif "permission" in error_msg .lower () or "forbidden" in error_msg .lower ():
85+ return jsonify ({
86+ 'error' : 'Access denied. You may not have permission to access this model endpoint.' ,
87+ 'details' : 'Please contact your Databricks administrator.'
88+ }), 403
89+ elif "timeout" in error_msg .lower ():
90+ return jsonify ({'error' : 'Request timed out. Please try again.' }), 504
91+ else :
92+ return jsonify ({
93+ 'error' : f'Unexpected error: { error_msg } ,
94+ 'details' : 'Please try again or contact support if the issue persists.'
95+ }), 500
96+
97+ @app .route ('/health' )
98+ def health ():
99+ """Health check endpoint"""
100+ user_email = request .headers .get ('X-Forwarded-Email' )
101+
102+ try :
103+ # Test the client connection
104+ client .current_user .me ()
105+ client_healthy = True
106+ except Exception as e :
107+ client_healthy = False
108+ app .logger .error (f"Client health check failed: { str (e )} )
109+
110+ return jsonify ({
111+ 'status' : 'healthy' ,
112+ 'client_authenticated' : client_healthy ,
113+ 'user_email' : user_email if user_email else 'anonymous' ,
114+ 'endpoint_name' : ENDPOINT_NAME
115+ })
116+
117+ @app .route ('/api/debug' )
118+ def debug ():
119+ """Debug endpoint to check authentication and endpoint status"""
120+ headers_info = {
121+ 'X-Forwarded-Access-Token' : 'present' if request .headers .get ('X-Forwarded-Access-Token' ) else 'missing' ,
122+ 'X-Forwarded-Email' : request .headers .get ('X-Forwarded-Email' , 'not provided' ),
123+ 'User-Agent' : request .headers .get ('User-Agent' , 'not provided' ),
124+ 'Authorization' : 'present' if request .headers .get ('Authorization' ) else 'missing'
125+ }
126+
127+ try :
128+ # Check if we can access the serving endpoint
129+ endpoint_info = client .serving_endpoints .get (name = ENDPOINT_NAME )
130+ endpoint_status = endpoint_info .state .value if endpoint_info .state else 'unknown'
131+ except Exception as e :
132+ endpoint_status = f'error: { str (e )}
133+
134+ try :
135+ # Check current user
136+ current_user = client .current_user .me ()
137+ user_info = current_user .user_name if current_user else 'unknown'
138+ except Exception as e :
139+ user_info = f'error: { str (e )}
140+
141+ return jsonify ({
142+ 'message' : 'Debug information for Databricks App' ,
143+ 'headers' : headers_info ,
144+ 'endpoint_name' : ENDPOINT_NAME ,
145+ 'endpoint_status' : endpoint_status ,
146+ 'current_user' : user_info ,
147+ 'sdk_config' : {
148+ 'host' : cfg .host if hasattr (cfg , 'host' ) else 'not set' ,
149+ 'auth_type' : cfg .auth_type if hasattr (cfg , 'auth_type' ) else 'not set'
150+ }
151+ })
152+
153+ if __name__ == '__main__' :
154+ # Get port from environment variable or default to 8080
155+ port = int (os .environ .get ('PORT' , 8080 ))
156+ app .run (host = '0.0.0.0' , port = port , debug = False )
0 commit comments