yuxh1996 commited on
Commit
98727e9
Β·
1 Parent(s): c1359fb

Enhance API key retrieval and validation for MCP clients

Browse files
Files changed (1) hide show
  1. app.py +27 -13
app.py CHANGED
@@ -16,38 +16,50 @@ def get_api_client():
16
  """Get API client with current API key"""
17
  # Try to get API key from multiple sources
18
  api_key = None
 
19
 
20
  # 1. Try from request headers (for MCP clients)
21
  try:
22
  request = gr.request()
23
  if request and hasattr(request, 'headers'):
24
- api_key = get_api_key_from_headers(dict(request.headers))
25
- except:
26
- pass
 
 
 
 
 
27
 
28
  # 2. Check if running on Hugging Face Space
29
  is_space = os.getenv("SPACE_ID") is not None
30
  space_api_key = os.getenv("A1D_API_KEY")
 
 
31
 
32
  # 3. Determine if this is a web browser request or MCP client request
33
  is_web_request = False
34
- try:
35
- request = gr.request()
36
- if request and hasattr(request, 'headers'):
37
- user_agent = request.headers.get('user-agent', '').lower()
38
- # Web browsers typically have 'mozilla' in user agent
39
- is_web_request = 'mozilla' in user_agent or 'chrome' in user_agent or 'safari' in user_agent
40
- except:
41
- is_web_request = True # Default to web request if we can't determine
 
 
 
 
42
 
43
  # 4. Use Space API key ONLY for web browser requests on Hugging Face Space
44
  if not api_key and is_space and space_api_key and is_web_request:
45
  print("πŸ“‘ Using API key from Space environment variable (web demo)")
46
  return A1DAPIClient(space_api_key)
47
 
48
- # 5. For MCP clients, API key is mandatory
49
  if not api_key:
50
- raise ValueError(
51
  "πŸ”‘ API key is required for MCP clients!\n\n"
52
  "Please provide API_KEY in request headers.\n"
53
  "Get your API key at https://a1d.ai\n\n"
@@ -69,6 +81,8 @@ def get_api_client():
69
  ' }\n'
70
  '}'
71
  )
 
 
72
 
73
  print("πŸ”‘ Using API key from MCP client headers")
74
  return A1DAPIClient(api_key)
 
16
  """Get API client with current API key"""
17
  # Try to get API key from multiple sources
18
  api_key = None
19
+ user_agent = ""
20
 
21
  # 1. Try from request headers (for MCP clients)
22
  try:
23
  request = gr.request()
24
  if request and hasattr(request, 'headers'):
25
+ headers = dict(request.headers)
26
+ api_key = get_api_key_from_headers(headers)
27
+ user_agent = headers.get('user-agent', '')
28
+ print(f"πŸ” Request headers found - User-Agent: {user_agent}")
29
+ print(
30
+ f"πŸ” API key from headers: {'Found' if api_key else 'Not found'}")
31
+ except Exception as e:
32
+ print(f"πŸ” No request context available: {e}")
33
 
34
  # 2. Check if running on Hugging Face Space
35
  is_space = os.getenv("SPACE_ID") is not None
36
  space_api_key = os.getenv("A1D_API_KEY")
37
+ print(
38
+ f"πŸ” Environment check - Is Space: {is_space}, Space API key: {'Found' if space_api_key else 'Not found'}")
39
 
40
  # 3. Determine if this is a web browser request or MCP client request
41
  is_web_request = False
42
+ if user_agent:
43
+ user_agent_lower = user_agent.lower()
44
+ # Web browsers typically have 'mozilla' in user agent
45
+ is_web_request = ('mozilla' in user_agent_lower or
46
+ 'chrome' in user_agent_lower or
47
+ 'safari' in user_agent_lower or
48
+ 'edge' in user_agent_lower)
49
+ print(f"πŸ” Request type detection - Is web request: {is_web_request}")
50
+ else:
51
+ # If no user agent, assume it's NOT a web request (likely MCP client)
52
+ is_web_request = False
53
+ print("πŸ” No User-Agent found - assuming MCP client request")
54
 
55
  # 4. Use Space API key ONLY for web browser requests on Hugging Face Space
56
  if not api_key and is_space and space_api_key and is_web_request:
57
  print("πŸ“‘ Using API key from Space environment variable (web demo)")
58
  return A1DAPIClient(space_api_key)
59
 
60
+ # 5. For MCP clients or when no Space API key available, user API key is mandatory
61
  if not api_key:
62
+ error_msg = (
63
  "πŸ”‘ API key is required for MCP clients!\n\n"
64
  "Please provide API_KEY in request headers.\n"
65
  "Get your API key at https://a1d.ai\n\n"
 
81
  ' }\n'
82
  '}'
83
  )
84
+ print(f"❌ API key validation failed: {error_msg}")
85
+ raise ValueError(error_msg)
86
 
87
  print("πŸ”‘ Using API key from MCP client headers")
88
  return A1DAPIClient(api_key)