Mirrowel commited on
Commit
2c9c7d5
·
1 Parent(s): c48abee

feat(setup): ensure proxy key validation runs after credential creation

Browse files

Delay mandatory PROXY_API_KEY validation until the server runtime block to facilitate initial environment setup.

This adjustment allows the application to proceed with the `--add-credential` flow even if the proxy key is initially unset:
- Imports and runs `ensure_env_defaults()` to generate missing keys or files.
- Reloads environment variables (dotenv) immediately after generation.
- Validates the PROXY_API_KEY only before starting uvicorn.

Also adds support for CHUTES_API_KEY in the credential tool.

src/proxy_app/main.py CHANGED
@@ -133,8 +133,7 @@ ENABLE_REQUEST_LOGGING = args.enable_request_logging
133
  if ENABLE_REQUEST_LOGGING:
134
  logging.info("Request logging is enabled.")
135
  PROXY_API_KEY = os.getenv("PROXY_API_KEY")
136
- if not PROXY_API_KEY:
137
- raise ValueError("PROXY_API_KEY environment variable not set.")
138
 
139
  # Discover API keys from environment variables
140
  api_keys = {}
@@ -713,7 +712,15 @@ async def token_count(
713
 
714
  if __name__ == "__main__":
715
  if args.add_credential:
 
 
 
 
 
716
  run_credential_tool()
717
  else:
 
 
 
718
  import uvicorn
719
  uvicorn.run(app, host=args.host, port=args.port)
 
133
  if ENABLE_REQUEST_LOGGING:
134
  logging.info("Request logging is enabled.")
135
  PROXY_API_KEY = os.getenv("PROXY_API_KEY")
136
+ # Note: PROXY_API_KEY validation moved to server startup to allow credential tool to run first
 
137
 
138
  # Discover API keys from environment variables
139
  api_keys = {}
 
712
 
713
  if __name__ == "__main__":
714
  if args.add_credential:
715
+ # Import and call ensure_env_defaults to create .env and PROXY_API_KEY if needed
716
+ from rotator_library.credential_tool import ensure_env_defaults
717
+ ensure_env_defaults()
718
+ # Reload environment variables after ensure_env_defaults creates/updates .env
719
+ load_dotenv(override=True)
720
  run_credential_tool()
721
  else:
722
+ # Validate PROXY_API_KEY before starting the server
723
+ if not PROXY_API_KEY:
724
+ raise ValueError("PROXY_API_KEY environment variable not set. Please run with --add-credential to set up your environment.")
725
  import uvicorn
726
  uvicorn.run(app, host=args.host, port=args.port)
src/rotator_library/credential_tool.py CHANGED
@@ -47,7 +47,8 @@ async def setup_api_key():
47
  "OpenAI": "OPENAI_API_KEY", "Anthropic": "ANTHROPIC_API_KEY",
48
  "Google AI Studio (Gemini)": "GEMINI_API_KEY", "Azure OpenAI": "AZURE_API_KEY",
49
  "Vertex AI": "GOOGLE_API_KEY", "AWS Bedrock": "AWS_ACCESS_KEY_ID",
50
- "Cohere": "COHERE_API_KEY", "Mistral AI": "MISTRAL_API_KEY",
 
51
  "Codestral (Mistral)": "CODESTRAL_API_KEY", "Groq": "GROQ_API_KEY",
52
  "Perplexity": "PERPLEXITYAI_API_KEY", "xAI": "XAI_API_KEY",
53
  "Together AI": "TOGETHERAI_API_KEY", "Fireworks AI": "FIREWORKS_AI_API_KEY",
 
47
  "OpenAI": "OPENAI_API_KEY", "Anthropic": "ANTHROPIC_API_KEY",
48
  "Google AI Studio (Gemini)": "GEMINI_API_KEY", "Azure OpenAI": "AZURE_API_KEY",
49
  "Vertex AI": "GOOGLE_API_KEY", "AWS Bedrock": "AWS_ACCESS_KEY_ID",
50
+ "Cohere": "COHERE_API_KEY", "Chutes": "CHUTES_API_KEY",
51
+ "Mistral AI": "MISTRAL_API_KEY",
52
  "Codestral (Mistral)": "CODESTRAL_API_KEY", "Groq": "GROQ_API_KEY",
53
  "Perplexity": "PERPLEXITYAI_API_KEY", "xAI": "XAI_API_KEY",
54
  "Together AI": "TOGETHERAI_API_KEY", "Fireworks AI": "FIREWORKS_AI_API_KEY",