Coverage for mcpgateway / routers / llm_proxy_router.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/routers/llm_proxy_router.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5 

6LLM Proxy Router. 

7This module provides OpenAI-compatible API endpoints for the internal 

8LLM proxy service. It routes requests to configured LLM providers. 

9""" 

10 

11# Third-Party 

12from fastapi import APIRouter, Depends, HTTPException, status 

13from fastapi.responses import StreamingResponse 

14from sqlalchemy.orm import Session 

15 

16# First-Party 

17from mcpgateway.config import settings 

18from mcpgateway.db import get_db 

19from mcpgateway.llm_schemas import ChatCompletionRequest, ChatCompletionResponse 

20from mcpgateway.middleware.rbac import get_current_user_with_permissions, require_permission 

21from mcpgateway.services.llm_provider_service import ( 

22 LLMModelNotFoundError, 

23 LLMProviderNotFoundError, 

24) 

25from mcpgateway.services.llm_proxy_service import ( 

26 LLMProxyAuthError, 

27 LLMProxyRequestError, 

28 LLMProxyService, 

29) 

30from mcpgateway.services.logging_service import LoggingService 

31 

32# Initialize logging 

33logging_service = LoggingService() 

34logger = logging_service.get_logger(__name__) 

35 

36# Create router 

37llm_proxy_router = APIRouter() 

38 

39# Initialize service 

40llm_proxy_service = LLMProxyService() 

41 

42 

43@llm_proxy_router.post( 

44 "/chat/completions", 

45 response_model=ChatCompletionResponse, 

46 summary="Chat Completions", 

47 description="Create a chat completion using configured LLM providers. OpenAI-compatible API.", 

48 responses={ 

49 200: {"description": "Chat completion response"}, 

50 400: {"description": "Invalid request"}, 

51 401: {"description": "Authentication required"}, 

52 404: {"description": "Model not found"}, 

53 500: {"description": "Provider error"}, 

54 }, 

55) 

56@require_permission("llm.invoke") 

57async def chat_completions( 

58 request: ChatCompletionRequest, 

59 db: Session = Depends(get_db), 

60 current_user: dict = Depends(get_current_user_with_permissions), 

61): 

62 """Create a chat completion. 

63 

64 This endpoint is compatible with the OpenAI Chat Completions API. 

65 It routes requests to configured LLM providers based on the model ID. 

66 

67 Args: 

68 request: Chat completion request (OpenAI-compatible). 

69 db: Database session. 

70 current_user: Authenticated user. 

71 

72 Returns: 

73 ChatCompletionResponse or StreamingResponse for streaming requests. 

74 

75 Raises: 

76 HTTPException: If model not found, streaming disabled, or provider error. 

77 """ 

78 # Check if streaming is enabled 

79 if request.stream and not settings.llm_streaming_enabled: 

80 raise HTTPException( 

81 status_code=status.HTTP_400_BAD_REQUEST, 

82 detail="Streaming is disabled in gateway configuration", 

83 ) 

84 

85 try: 

86 if request.stream: 

87 # Return streaming response 

88 return StreamingResponse( 

89 llm_proxy_service.chat_completion_stream(db, request), 

90 media_type="text/event-stream", 

91 headers={ 

92 "Cache-Control": "no-cache", 

93 "Connection": "keep-alive", 

94 "X-Accel-Buffering": "no", 

95 }, 

96 ) 

97 else: 

98 # Return regular response 

99 return await llm_proxy_service.chat_completion(db, request) 

100 

101 except LLMModelNotFoundError as e: 

102 logger.warning(f"Model not found: {request.model}") 

103 raise HTTPException( 

104 status_code=status.HTTP_404_NOT_FOUND, 

105 detail=str(e), 

106 ) 

107 except LLMProviderNotFoundError as e: 

108 logger.warning(f"Provider not found for model: {request.model}") 

109 raise HTTPException( 

110 status_code=status.HTTP_404_NOT_FOUND, 

111 detail=str(e), 

112 ) 

113 except LLMProxyAuthError as e: 

114 logger.error(f"Authentication error: {e}") 

115 raise HTTPException( 

116 status_code=status.HTTP_401_UNAUTHORIZED, 

117 detail=str(e), 

118 ) 

119 except LLMProxyRequestError as e: 

120 logger.error(f"Proxy request error: {e}") 

121 raise HTTPException( 

122 status_code=status.HTTP_502_BAD_GATEWAY, 

123 detail=str(e), 

124 ) 

125 except Exception as e: 

126 logger.error(f"Unexpected error in chat completion: {e}") 

127 raise HTTPException( 

128 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

129 detail=f"Internal error: {str(e)}", 

130 ) 

131 

132 

133@llm_proxy_router.get( 

134 "/models", 

135 summary="List Models", 

136 description="List available models from configured providers. OpenAI-compatible API.", 

137) 

138@require_permission("llm.read") 

139async def list_models( 

140 db: Session = Depends(get_db), 

141 current_user: dict = Depends(get_current_user_with_permissions), 

142): 

143 """List available models. 

144 

145 Returns a list of available models in OpenAI-compatible format. 

146 

147 Args: 

148 db: Database session. 

149 current_user: Authenticated user. 

150 

151 Returns: 

152 List of available models. 

153 """ 

154 # First-Party 

155 from mcpgateway.services.llm_provider_service import LLMProviderService 

156 

157 provider_service = LLMProviderService() 

158 models = provider_service.get_gateway_models(db) 

159 

160 # Format as OpenAI-compatible response 

161 model_list = [] 

162 for model in models: 

163 model_list.append( 

164 { 

165 "id": model.model_id, 

166 "object": "model", 

167 "created": 0, 

168 "owned_by": model.provider_name, 

169 } 

170 ) 

171 

172 return { 

173 "object": "list", 

174 "data": model_list, 

175 }