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