openrag/src/api/models.py
2025-12-01 15:01:10 -06:00

200 lines
7 KiB
Python

from starlette.responses import JSONResponse
from utils.logging_config import get_logger
from config.settings import get_openrag_config
logger = get_logger(__name__)
async def get_openai_models(request, models_service, session_manager):
"""Get available OpenAI models"""
try:
# Get API key from request body
api_key = None
try:
body = await request.json()
api_key = body.get("api_key") if body else None
except Exception:
# Body might be empty or invalid JSON, continue to fallback
pass
# If no API key provided, try to get it from stored configuration
if not api_key:
try:
config = get_openrag_config()
api_key = config.providers.openai.api_key
logger.info(
f"Retrieved OpenAI API key from config: {'yes' if api_key else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not api_key:
return JSONResponse(
{
"error": "OpenAI API key is required either in request body or in configuration"
},
status_code=400,
)
models = await models_service.get_openai_models(api_key=api_key)
return JSONResponse(models)
except Exception as e:
logger.error(f"Failed to get OpenAI models: {str(e)}")
return JSONResponse(
{"error": f"Failed to retrieve OpenAI models: {str(e)}"}, status_code=500
)
async def get_anthropic_models(request, models_service, session_manager):
"""Get available Anthropic models"""
try:
# Get API key from request body
api_key = None
try:
body = await request.json()
api_key = body.get("api_key") if body else None
except Exception:
# Body might be empty or invalid JSON, continue to fallback
pass
# If no API key provided, try to get it from stored configuration
if not api_key:
try:
config = get_openrag_config()
api_key = config.providers.anthropic.api_key
logger.info(
f"Retrieved Anthropic API key from config: {'yes' if api_key else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not api_key:
return JSONResponse(
{
"error": "Anthropic API key is required either in request body or in configuration"
},
status_code=400,
)
models = await models_service.get_anthropic_models(api_key=api_key)
return JSONResponse(models)
except Exception as e:
logger.error(f"Failed to get Anthropic models: {str(e)}")
return JSONResponse(
{"error": f"Failed to retrieve Anthropic models: {str(e)}"}, status_code=500
)
async def get_ollama_models(request, models_service, session_manager):
"""Get available Ollama models"""
try:
# Get endpoint from query parameters if provided
query_params = dict(request.query_params)
endpoint = query_params.get("endpoint")
# If no endpoint provided, try to get it from stored configuration
if not endpoint:
try:
config = get_openrag_config()
endpoint = config.providers.ollama.endpoint
logger.info(
f"Retrieved Ollama endpoint from config: {'yes' if endpoint else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not endpoint:
return JSONResponse(
{
"error": "Endpoint is required either as query parameter or in configuration"
},
status_code=400,
)
models = await models_service.get_ollama_models(endpoint=endpoint)
return JSONResponse(models)
except Exception as e:
logger.error(f"Failed to get Ollama models: {str(e)}")
return JSONResponse(
{"error": f"Failed to retrieve Ollama models: {str(e)}"}, status_code=500
)
async def get_ibm_models(request, models_service, session_manager):
"""Get available IBM Watson models"""
try:
# Get parameters from request body if provided
endpoint = None
api_key = None
project_id = None
try:
body = await request.json()
if body:
endpoint = body.get("endpoint")
api_key = body.get("api_key")
project_id = body.get("project_id")
except Exception:
# Body might be empty or invalid JSON, continue to fallback
pass
config = get_openrag_config()
# If no API key provided, try to get it from stored configuration
if not api_key:
try:
api_key = config.providers.watsonx.api_key
logger.info(
f"Retrieved WatsonX API key from config: {'yes' if api_key else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not api_key:
return JSONResponse(
{
"error": "WatsonX API key is required either in request body or in configuration"
},
status_code=400,
)
if not endpoint:
try:
endpoint = config.providers.watsonx.endpoint
logger.info(
f"Retrieved WatsonX endpoint from config: {'yes' if endpoint else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not endpoint:
return JSONResponse(
{
"error": "Endpoint is required either in request body or in configuration"
},
status_code=400,
)
if not project_id:
try:
project_id = config.providers.watsonx.project_id
logger.info(
f"Retrieved WatsonX project ID from config: {'yes' if project_id else 'no'}"
)
except Exception as e:
logger.error(f"Failed to get config: {e}")
if not project_id:
return JSONResponse(
{
"error": "Project ID is required either in request body or in configuration"
},
status_code=400,
)
models = await models_service.get_ibm_models(
endpoint=endpoint, api_key=api_key, project_id=project_id
)
return JSONResponse(models)
except Exception as e:
logger.error(f"Failed to get IBM models: {str(e)}")
return JSONResponse(
{"error": f"Failed to retrieve IBM models: {str(e)}"}, status_code=500
)