Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-relay-gateway/src/lexigram/ai/relay/gateway/operations/upstream.py: 39%

38 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Upstream invocation helpers for the relay gateway service. 

2 

3Builds endpoint URLs and outbound model names from channel state, wraps 

4the upstream HTTP call with structured events, and records consecutive 

5failures against the failover tracker. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import Any 

11 

12from lexigram.ai.relay.gateway.config import RelayGatewayConfig 

13from lexigram.ai.relay.gateway.operations.failover import RelayFailoverTracker 

14from lexigram.contracts.ai.relay import ( 

15 RelayChannel, 

16 RelayFormat, 

17 RelayGatewayError, 

18 RelayGatewayRequest, 

19 RelayUpstreamProtocol, 

20 UpstreamRequest, 

21 UpstreamResponse, 

22) 

23from lexigram.contracts.core.result import Result 

24from lexigram.logging import get_logger 

25 

26__all__ = [ 

27 "call_upstream", 

28 "note_failure", 

29 "note_success", 

30 "outbound_model", 

31 "should_track_upstream_failure", 

32 "upstream_url", 

33] 

34 

35logger = get_logger(__name__) 

36 

37 

38def upstream_url(channel: RelayChannel, model: str) -> str: 

39 """Build the endpoint URL for *channel*'s target format. 

40 

41 Args: 

42 channel: The selected channel. 

43 model: Outbound model alias (embedded in the Gemini path). 

44 

45 Returns: 

46 The standard endpoint path for the channel's target format 

47 joined onto the channel's base URL. 

48 

49 Raises: 

50 ValueError: The channel's target format is not one of the 

51 four relay wire formats. Unreachable via registry 

52 validation. 

53 """ 

54 base = channel.upstream_base_url.rstrip("/") 

55 if channel.target_format == RelayFormat.OPENAI_CHAT: 

56 return f"{base}/v1/chat/completions" 

57 if channel.target_format == RelayFormat.OPENAI_RESPONSES: 

58 return f"{base}/v1/responses" 

59 if channel.target_format == RelayFormat.CLAUDE: 

60 return f"{base}/v1/messages" 

61 if channel.target_format == RelayFormat.GEMINI: 

62 return f"{base}/v1beta/models/{model}:generateContent" 

63 raise ValueError(f"unsupported target format: {channel.target_format}") 

64 

65 

66def outbound_model( 

67 config: RelayGatewayConfig, channel: RelayChannel, alias: str 

68) -> str: 

69 """Resolve the upstream model name for *alias* on *channel*. 

70 

71 The channel's ``model_map`` wins when it carries the alias; 

72 otherwise the alias is sent as-is. The channel's configured 

73 suffix (e.g. ``":thinking"``) is appended after the mapping. 

74 

75 Args: 

76 config: Gateway configuration carrying the model suffix table. 

77 channel: The selected channel. 

78 alias: The client-visible model alias from the request. 

79 

80 Returns: 

81 The model name sent to the channel's upstream. 

82 """ 

83 return channel.resolve_model(alias) + config.model_suffix.get(channel.name, "") 

84 

85 

86def should_track_upstream_failure(code: str) -> bool: 

87 """Return whether *code* counts toward a channel failover ban. 

88 

89 Transport-level upstream failures count; a cancelled client 

90 request and a malformed but delivered 2xx body do not. 

91 

92 Args: 

93 code: The gateway error code of the failed upstream call. 

94 

95 Returns: 

96 ``True`` when the failure should count, ``False`` otherwise. 

97 """ 

98 return code in { 

99 "UPSTREAM_ERROR", 

100 "UPSTREAM_TIMEOUT", 

101 "UPSTREAM_FAILED", 

102 } 

103 

104 

105def note_failure(failover: RelayFailoverTracker | None, channel_name: str) -> None: 

106 """Count one upstream failure against *channel_name*. 

107 

108 Args: 

109 failover: The failover tracker; ``None`` disables accounting. 

110 channel_name: The channel that failed upstream. 

111 """ 

112 if failover is not None: 

113 failover.record_failure(channel_name) 

114 

115 

116def note_success(failover: RelayFailoverTracker | None, channel_name: str) -> None: 

117 """Reset *channel_name*'s failures and restore it when banned. 

118 

119 Args: 

120 failover: The failover tracker; ``None`` disables accounting. 

121 channel_name: The channel that succeeded upstream. 

122 """ 

123 if failover is not None: 

124 failover.record_success(channel_name) 

125 

126 

127async def call_upstream( 

128 upstream: RelayUpstreamProtocol, 

129 channel: RelayChannel, 

130 outbound_model: str, 

131 payload: dict[str, Any], 

132 request: RelayGatewayRequest, 

133) -> Result[UpstreamResponse, RelayGatewayError]: 

134 """Send the converted payload to the selected channel's endpoint. 

135 

136 Args: 

137 upstream: The upstream transport adapter. 

138 channel: The selected channel. 

139 outbound_model: Model alias with the channel's suffix applied. 

140 payload: Converted request payload dict. 

141 request: The original gateway request. 

142 

143 Returns: 

144 ``Ok(UpstreamResponse)`` or ``Err`` as returned by the 

145 adapter; the adapter already normalizes transport failures. 

146 """ 

147 url = upstream_url(channel, outbound_model) 

148 logger.info( 

149 "relay_gateway_upstream_started", 

150 request_id=request.request_id, 

151 channel=channel.name, 

152 method="POST", 

153 url=url, 

154 ) 

155 result = await upstream.request( 

156 UpstreamRequest( 

157 request_id=request.request_id, 

158 method="POST", 

159 url=url, 

160 headers={"content-type": "application/json"}, 

161 payload=payload, 

162 timeout_seconds=channel.timeout_seconds, 

163 channel_name=channel.name, 

164 ) 

165 ) 

166 if result.is_err(): 

167 error = result.unwrap_err() 

168 logger.warning( 

169 "relay_gateway_upstream_failed", 

170 request_id=request.request_id, 

171 channel=channel.name, 

172 code=error.code, 

173 status_code=error.status_code, 

174 error=str(error), 

175 ) 

176 return result