Enhance main_text_generation with APIKeyManager and improved provider routing

- Import APIKeyManager for provider key checking
- Use APIKeyManager.get_api_key() instead of get_api_key() function
- Add wavespeed provider to available_providers check
- Add detailed provider preflight logging with flow_type tag
- Improve fallback logic when preferred provider is unavailable

These improvements come from PRs #423-#431 while maintaining the modular textgen_utils structure.
This commit is contained in:
ajaysi
2026-03-22 11:23:38 +05:30
parent a26fa84263
commit e8f282b7a9

View File

@@ -10,6 +10,7 @@ from typing import Optional, Dict, Any, List
from datetime import datetime from datetime import datetime
from loguru import logger from loguru import logger
from fastapi import HTTPException from fastapi import HTTPException
from ..onboarding.api_key_manager import APIKeyManager
from .gemini_provider import gemini_text_response, gemini_structured_json_response from .gemini_provider import gemini_text_response, gemini_structured_json_response
from .huggingface_provider import huggingface_text_response, huggingface_structured_json_response from .huggingface_provider import huggingface_text_response, huggingface_structured_json_response
@@ -133,15 +134,30 @@ def llm_text_gen(
blog_output_format = "markdown" blog_output_format = "markdown"
blog_length = 2000 blog_length = 2000
# Check which providers have API keys available using APIKeyManager
api_key_manager = APIKeyManager()
available_providers = [] available_providers = []
for provider in ("google", "huggingface"): if api_key_manager.get_api_key("gemini"):
if get_api_key(provider, user_id=user_id): available_providers.append("google")
available_providers.append(provider) if api_key_manager.get_api_key("hf_token"):
available_providers.append("huggingface")
if api_key_manager.get_api_key("wavespeed"):
available_providers.append("wavespeed")
logger.info(
f"[llm_text_gen][{flow_tag}] Provider preflight: env_provider='{env_provider or 'auto'}', "
f"provider_list={provider_list}, strict_provider_mode={strict_provider_mode}, "
f"available_providers={available_providers}, preferred_provider={preferred_provider or 'none'}"
)
if gpt_provider not in available_providers: if gpt_provider not in available_providers:
logger.warning(f"[llm_text_gen] Provider {gpt_provider} unavailable for user {user_id}, falling back.") logger.warning(f"[llm_text_gen] Provider {gpt_provider} unavailable for user {user_id}, falling back.")
if available_providers: if "huggingface" in available_providers:
gpt_provider = available_providers[0] gpt_provider = "huggingface"
model = "openai/gpt-oss-120b:cerebras"
elif "google" in available_providers:
gpt_provider = "google"
model = "gemini-2.0-flash-001"
else: else:
logger.error("[llm_text_gen] No API keys found for supported providers.") logger.error("[llm_text_gen] No API keys found for supported providers.")
raise RuntimeError("No LLM API keys configured for tenant or environment defaults.") raise RuntimeError("No LLM API keys configured for tenant or environment defaults.")