bibibi12345 commited on
Commit
9119b63
·
1 Parent(s): a87d626

fixed token expire at startup issue

Browse files
Files changed (1) hide show
  1. gemini_proxy.py +35 -6
gemini_proxy.py CHANGED
@@ -421,6 +421,20 @@ def get_credentials():
421
  credentials = Credentials.from_authorized_user_info(creds_data, SCOPES)
422
  print("Loaded credentials from GOOGLE_APPLICATION_CREDENTIALS.")
423
  print(f"DEBUG: Env credentials - Token: {credentials.token[:20] if credentials.token else 'None'}..., Expired: {credentials.expired}, Expiry: {credentials.expiry}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
424
  return credentials
425
  except Exception as e:
426
  print(f"Could not load credentials from GOOGLE_APPLICATION_CREDENTIALS: {e}")
@@ -474,6 +488,19 @@ def get_credentials():
474
  credentials._expiry = expiry_utc
475
  return credentials
476
 
 
 
 
 
 
 
 
 
 
 
 
 
 
477
  return credentials
478
  except Exception as e:
479
  print(f"Could not load cached credentials: {e}. Starting new login.")
@@ -836,15 +863,15 @@ async def proxy_request(request: Request, full_path: str, username: str = Depend
836
 
837
  if resp.status_code == 401:
838
  print("[STREAM] Still getting 401 after token refresh.")
839
- yield f'data: {{"error": {{"message": "Authentication failed even after token refresh. Please restart the proxy to re-authenticate."}}}}\n\n'
840
  return
841
  except Exception as refresh_error:
842
  print(f"[STREAM] Token refresh failed: {refresh_error}")
843
- yield f'data: {{"error": {{"message": "Token refresh failed. Please restart the proxy to re-authenticate."}}}}\n\n'
844
  return
845
  else:
846
  print("[STREAM] No refresh token available.")
847
- yield f'data: {{"error": {{"message": "Authentication failed. Please restart the proxy to re-authenticate."}}}}\n\n'
848
  return
849
 
850
  with resp:
@@ -874,7 +901,9 @@ async def proxy_request(request: Request, full_path: str, username: str = Depend
874
  response_chunk = obj["response"]
875
  # Output in standard Gemini streaming format
876
  response_json = json.dumps(response_chunk, separators=(',', ':'))
877
- yield f"data: {response_json}\n\n"
 
 
878
  except json.JSONDecodeError:
879
  # Skip invalid JSON
880
  continue
@@ -882,11 +911,11 @@ async def proxy_request(request: Request, full_path: str, username: str = Depend
882
  except requests.exceptions.RequestException as e:
883
  print(f"Error during streaming request: {e}")
884
  # Format error as real Gemini API would
885
- yield f'data: {{"error": {{"message": "Upstream request failed: {str(e)}"}}}}\n\n'
886
  except Exception as e:
887
  print(f"An unexpected error occurred during streaming: {e}")
888
  # Format error as real Gemini API would
889
- yield f'data: {{"error": {{"message": "An unexpected error occurred: {str(e)}"}}}}\n\n'
890
 
891
  # Create the streaming response with headers matching real Gemini API
892
  response_headers = {
 
421
  credentials = Credentials.from_authorized_user_info(creds_data, SCOPES)
422
  print("Loaded credentials from GOOGLE_APPLICATION_CREDENTIALS.")
423
  print(f"DEBUG: Env credentials - Token: {credentials.token[:20] if credentials.token else 'None'}..., Expired: {credentials.expired}, Expiry: {credentials.expiry}")
424
+
425
+ # Always refresh tokens at startup when loading from file to avoid issues
426
+ if credentials.refresh_token:
427
+ print("Refreshing environment credentials at startup for reliability...")
428
+ try:
429
+ credentials.refresh(GoogleAuthRequest())
430
+ # Note: We don't save environment credentials back to the env file
431
+ print("Startup token refresh successful for environment credentials.")
432
+ except Exception as refresh_error:
433
+ print(f"Startup token refresh failed for environment credentials: {refresh_error}. Credentials may be stale.")
434
+ # Continue with existing credentials - they might still work
435
+ else:
436
+ print("No refresh token available in environment credentials - using as-is.")
437
+
438
  return credentials
439
  except Exception as e:
440
  print(f"Could not load credentials from GOOGLE_APPLICATION_CREDENTIALS: {e}")
 
488
  credentials._expiry = expiry_utc
489
  return credentials
490
 
491
+ # Always refresh tokens at startup when loading from file to avoid issues
492
+ if credentials.refresh_token:
493
+ print("Refreshing tokens at startup for reliability...")
494
+ try:
495
+ credentials.refresh(GoogleAuthRequest())
496
+ save_credentials(credentials)
497
+ print("Startup token refresh successful.")
498
+ except Exception as refresh_error:
499
+ print(f"Startup token refresh failed: {refresh_error}. Credentials may be stale.")
500
+ # Continue with existing credentials - they might still work
501
+ else:
502
+ print("No refresh token available - using cached credentials as-is.")
503
+
504
  return credentials
505
  except Exception as e:
506
  print(f"Could not load cached credentials: {e}. Starting new login.")
 
863
 
864
  if resp.status_code == 401:
865
  print("[STREAM] Still getting 401 after token refresh.")
866
+ yield f'data: {{"error": {{"message": "Authentication failed even after token refresh. Please restart the proxy to re-authenticate."}}}}\n\n'.encode('utf-8')
867
  return
868
  except Exception as refresh_error:
869
  print(f"[STREAM] Token refresh failed: {refresh_error}")
870
+ yield f'data: {{"error": {{"message": "Token refresh failed. Please restart the proxy to re-authenticate."}}}}\n\n'.encode('utf-8')
871
  return
872
  else:
873
  print("[STREAM] No refresh token available.")
874
+ yield f'data: {{"error": {{"message": "Authentication failed. Please restart the proxy to re-authenticate."}}}}\n\n'.encode('utf-8')
875
  return
876
 
877
  with resp:
 
901
  response_chunk = obj["response"]
902
  # Output in standard Gemini streaming format
903
  response_json = json.dumps(response_chunk, separators=(',', ':'))
904
+ # Encode back to UTF-8 bytes to match exactly what real Gemini API sends
905
+ response_line = f"data: {response_json}\n\n"
906
+ yield response_line.encode('utf-8')
907
  except json.JSONDecodeError:
908
  # Skip invalid JSON
909
  continue
 
911
  except requests.exceptions.RequestException as e:
912
  print(f"Error during streaming request: {e}")
913
  # Format error as real Gemini API would
914
+ yield f'data: {{"error": {{"message": "Upstream request failed: {str(e)}"}}}}\n\n'.encode('utf-8')
915
  except Exception as e:
916
  print(f"An unexpected error occurred during streaming: {e}")
917
  # Format error as real Gemini API would
918
+ yield f'data: {{"error": {{"message": "An unexpected error occurred: {str(e)}"}}}}\n\n'.encode('utf-8')
919
 
920
  # Create the streaming response with headers matching real Gemini API
921
  response_headers = {