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.auth import get_current_user 

18from mcpgateway.config import settings 

19from mcpgateway.db import get_db 

20from mcpgateway.llm_schemas import ChatCompletionRequest, ChatCompletionResponse 

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) 

56async def chat_completions( 

57 request: ChatCompletionRequest, 

58 db: Session = Depends(get_db), 

59 current_user: dict = Depends(get_current_user), 

60): 

61 """Create a chat completion. 

62 

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

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

65 

66 Args: 

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

68 db: Database session. 

69 current_user: Authenticated user. 

70 

71 Returns: 

72 ChatCompletionResponse or StreamingResponse for streaming requests. 

73 

74 Raises: 

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

76 """ 

77 # Check if streaming is enabled 

78 if request.stream and not settings.llm_streaming_enabled: 

79 raise HTTPException( 

80 status_code=status.HTTP_400_BAD_REQUEST, 

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

82 ) 

83 

84 try: 

85 if request.stream: 

86 # Return streaming response 

87 return StreamingResponse( 

88 llm_proxy_service.chat_completion_stream(db, request), 

89 media_type="text/event-stream", 

90 headers={ 

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

92 "Connection": "keep-alive", 

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

94 }, 

95 ) 

96 else: 

97 # Return regular response 

98 return await llm_proxy_service.chat_completion(db, request) 

99 

100 except LLMModelNotFoundError as e: 

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

102 raise HTTPException( 

103 status_code=status.HTTP_404_NOT_FOUND, 

104 detail=str(e), 

105 ) 

106 except LLMProviderNotFoundError as e: 

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

108 raise HTTPException( 

109 status_code=status.HTTP_404_NOT_FOUND, 

110 detail=str(e), 

111 ) 

112 except LLMProxyAuthError as e: 

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

114 raise HTTPException( 

115 status_code=status.HTTP_401_UNAUTHORIZED, 

116 detail=str(e), 

117 ) 

118 except LLMProxyRequestError as e: 

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

120 raise HTTPException( 

121 status_code=status.HTTP_502_BAD_GATEWAY, 

122 detail=str(e), 

123 ) 

124 except Exception as e: 

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

126 raise HTTPException( 

127 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

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

129 ) 

130 

131 

132@llm_proxy_router.get( 

133 "/models", 

134 summary="List Models", 

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

136) 

137async def list_models( 

138 db: Session = Depends(get_db), 

139 current_user: dict = Depends(get_current_user), 

140): 

141 """List available models. 

142 

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

144 

145 Args: 

146 db: Database session. 

147 current_user: Authenticated user. 

148 

149 Returns: 

150 List of available models. 

151 """ 

152 # First-Party 

153 from mcpgateway.services.llm_provider_service import LLMProviderService 

154 

155 provider_service = LLMProviderService() 

156 models = provider_service.get_gateway_models(db) 

157 

158 # Format as OpenAI-compatible response 

159 model_list = [] 

160 for model in models: 

161 model_list.append( 

162 { 

163 "id": model.model_id, 

164 "object": "model", 

165 "created": 0, 

166 "owned_by": model.provider_name, 

167 } 

168 ) 

169 

170 return { 

171 "object": "list", 

172 "data": model_list, 

173 }