1- """Unified AI Service Provider — routes all calls through AiGateway.
1+ """Unified AI Service Provider — thin wrapper around AiGateway client .
22
33Architecture::
44
55 QuantStrategyLifecycle AiGateway (VPS, single service)
66 ─────────────────── ─────────────────────────────────
77 AiServiceClient │
8- ├─ review() ── task=analyze ──────┤──▶ LlmAdapter (Claude/GPT API )
9- ├─ verify() ── task=execute ──────┤──▶ CodexAdapter (codex exec )
10- └─ execute()── task=execute ──────┤──▶ CodexAdapter (codex exec )
8+ ├─ review() ────────────────────── ┤──▶ POST /v1/ai/review (multi-model )
9+ ├─ verify() ────────────────────── ┤──▶ POST /v1/ai/execute/jobs (async )
10+ └─ execute()────────────────────── ┤──▶ POST /v1/ai/execute/jobs (async )
1111
1212No API keys in this repo — all AI backends accessed through AiGateway.
13- Only `CODEX_AUDIT_SERVICE_URL` is required.
13+ Only `` CODEX_AUDIT_SERVICE_URL` ` is required.
1414
15- Benefits:
16- - API keys live on the VPS (one place), not in N repos
17- - New backends = new adapter on the gateway, callers unchanged
18- - REPAIR/SAFETY patterns unchanged, just the transport is unified
15+ This module is a backward-compatible wrapper. New code should use
16+ ``ai_gateway_client.AiGatewayClient`` directly when available.
1917"""
2018
2119from __future__ import annotations
2220
2321import enum
24- import json
2522import os
26- import time
27- import urllib .error
28- import urllib .request
2923from collections .abc import Sequence
3024from dataclasses import dataclass
3125from typing import Any
3226
27+ # Try to import the unified client; fall back to local implementation
28+ try :
29+ from ai_gateway_client import AiGatewayClient , GatewayConfig , AiResult
30+ _HAS_GATEWAY_CLIENT = True
31+ except ImportError :
32+ _HAS_GATEWAY_CLIENT = False
33+
3334
3435class AiProviderId (str , enum .Enum ):
3536 CODEX_VPS = "codex_vps"
@@ -46,14 +47,11 @@ class AiPattern(str, enum.Enum):
4647class AiProviderConfig :
4748 provider : AiProviderId
4849 label : str
49- model : str = "" # sent as "model" to gateway
50+ model : str = ""
5051 task : str = "analyze" # analyze (API) or execute (Codex)
5152 can_execute_code : bool = False
5253 can_analyze : bool = True
5354
54- def resolve_service_url (self ) -> str | None :
55- return os .environ .get ("CODEX_AUDIT_SERVICE_URL" , "" ).strip () or None
56-
5755 @classmethod
5856 def claude (cls ) -> "AiProviderConfig" :
5957 return cls (provider = AiProviderId .CLAUDE , label = "Claude" ,
@@ -105,52 +103,111 @@ def from_env(cls) -> "AiServiceConfig":
105103
106104
107105class AiServiceClient :
106+ """Backward-compatible wrapper around AiGatewayClient.
107+
108+ New code should use ``AiGatewayClient`` directly. This class exists
109+ to keep existing ``codex_integration.py`` and ``ai_reviewer.py`` working
110+ without changes.
111+ """
112+
108113 def __init__ (self , config : AiServiceConfig ):
109114 self .config = config
115+ if _HAS_GATEWAY_CLIENT :
116+ gw_config = GatewayConfig .from_env ()
117+ else :
118+ gw_config = None
119+ self ._gw_config = gw_config
110120
111121 def review (self , prompt : str , * , timeout : float = 120.0 ) -> list ["AiCallResult" ]:
112- """Run all reviewers (analyze/sync)."""
113- import concurrent .futures
122+ """Run all reviewers concurrently via AiGateway."""
114123 if not self .config .reviewers :
115124 return []
116- with concurrent .futures .ThreadPoolExecutor (max_workers = min (len (self .config .reviewers ), 3 )) as pool :
117- futures = {pool .submit (self ._call , c , prompt , timeout ): c for c in self .config .reviewers }
118- results = []
119- for f in concurrent .futures .as_completed (futures ):
120- try :
121- results .append (f .result ())
122- except Exception as exc :
123- results .append (AiCallResult .unavailable (futures [f ].label , str (exc )))
124- return results
125+
126+ if _HAS_GATEWAY_CLIENT and self ._gw_config :
127+ client = AiGatewayClient (self ._gw_config )
128+ reviewers_list = [
129+ self ._map_provider_label (c )
130+ for c in self .config .reviewers
131+ ]
132+ result = client .review (
133+ prompt ,
134+ reviewers = reviewers_list ,
135+ verifier = "codex" if self .config .verifier else None ,
136+ timeout = timeout ,
137+ )
138+ return [
139+ AiCallResult (
140+ provider = r .provider , success = r .success ,
141+ output = r .output , note = r .error if not r .success else "" ,
142+ )
143+ for r in result .results
144+ ]
145+
146+ # Fallback: local implementation (no gateway client available)
147+ return self ._review_local (prompt , timeout )
125148
126149 def verify (self , prompt : str , * , timeout : float = 600.0 ) -> "AiCallResult | None" :
127150 if self .config .verifier is None :
128151 return None
129- return self ._call (self .config .verifier , prompt , timeout )
152+
153+ if _HAS_GATEWAY_CLIENT and self ._gw_config :
154+ client = AiGatewayClient (self ._gw_config )
155+ r = client .execute (prompt , mode = "review_only" , timeout = timeout )
156+ return AiCallResult (provider = r .provider , success = r .success , output = r .output , note = r .error )
157+
158+ return self ._call_local (self .config .verifier , prompt , timeout )
130159
131160 def execute (self , prompt : str , * , timeout : float = 600.0 ) -> "AiCallResult" :
132161 if self .config .primary is not None :
133- r = self ._call (self .config .primary , prompt , timeout )
162+ r = self ._call_single (self .config .primary , prompt , timeout )
134163 if r .success :
135164 return r
136165 for fb in self .config .fallback :
137- r = self ._call (fb , prompt , timeout )
166+ r = self ._call_single (fb , prompt , timeout )
138167 if r .success :
139- return AiCallResult (provider = r .provider , success = True , output = r .output , raw = r . raw ,
168+ return AiCallResult (provider = r .provider , success = True , output = r .output ,
140169 note = "Fallback after primary failed" )
141170 return AiCallResult .unavailable ("all" , "All providers exhausted" )
142171
143- def _call (self , provider : AiProviderConfig , prompt : str , timeout : float ) -> "AiCallResult" :
144- """Call the AiGateway — all providers use the same endpoint."""
145- service_url = provider .resolve_service_url ()
172+ def _call_single (self , provider : AiProviderConfig , prompt : str , timeout : float ) -> "AiCallResult" :
173+ if _HAS_GATEWAY_CLIENT and self ._gw_config :
174+ client = AiGatewayClient (self ._gw_config )
175+ if provider .task == "analyze" :
176+ r = client .analyze (prompt , model = provider .model , timeout = timeout )
177+ else :
178+ r = client .execute (prompt , mode = "review_only" , timeout = timeout )
179+ return AiCallResult (provider = r .provider , success = r .success , output = r .output , note = r .error )
180+ return self ._call_local (provider , prompt , timeout )
181+
182+ def _review_local (self , prompt : str , timeout : float ) -> list ["AiCallResult" ]:
183+ """Local fallback when gateway client is not installed."""
184+ import concurrent .futures
185+ with concurrent .futures .ThreadPoolExecutor (max_workers = min (len (self .config .reviewers ), 3 )) as pool :
186+ futures = {pool .submit (self ._call_local , c , prompt , timeout ): c for c in self .config .reviewers }
187+ results = []
188+ for f in concurrent .futures .as_completed (futures ):
189+ try :
190+ results .append (f .result ())
191+ except Exception as exc :
192+ results .append (AiCallResult .unavailable (futures [f ].label , str (exc )))
193+ return results
194+
195+ def _call_local (self , provider : AiProviderConfig , prompt : str , timeout : float ) -> "AiCallResult" :
196+ """Direct HTTP call to AiGateway — used when client library not installed."""
197+ import json as _json
198+ import urllib .error as _urllib_err
199+ import urllib .request as _urllib_req
200+ import time as _time
201+
202+ service_url = os .environ .get ("CODEX_AUDIT_SERVICE_URL" , "" ).strip ()
146203 if not service_url :
147204 return AiCallResult .unavailable (provider .label , "CODEX_AUDIT_SERVICE_URL not configured" )
148205
149206 try :
150207 token = _fetch_oidc_token ()
151208 base_url = service_url .rstrip ("/" )
152209
153- payload = json .dumps ({
210+ payload = _json .dumps ({
154211 "task" : provider .task ,
155212 "model" : provider .model ,
156213 "prompt" : prompt ,
@@ -160,41 +217,31 @@ def _call(self, provider: AiProviderConfig, prompt: str, timeout: float) -> "AiC
160217 "mode" : "review_only" ,
161218 }).encode ("utf-8" )
162219
163- sync = provider .task == "analyze"
164- req = urllib .request .Request (
165- f"{ base_url } /v1/codex-audit/jobs" , data = payload , method = "POST" ,
220+ req = _urllib_req .Request (
221+ f"{ base_url } /v1/ai/execute/jobs" , data = payload , method = "POST" ,
166222 headers = {"Authorization" : f"Bearer { token } " , "Content-Type" : "application/json" ,
167223 "Accept" : "application/json" , "User-Agent" : "quant-strategy-lifecycle" },
168224 )
169- with urllib . request .urlopen (req , timeout = 30 ) as resp :
170- result = json .loads (resp .read ().decode ("utf-8" ))
225+ with _urllib_req .urlopen (req , timeout = 30 ) as resp :
226+ result = _json .loads (resp .read ().decode ("utf-8" ))
171227
172- if sync :
173- # Analyze returns result inline
174- status = result .get ("status" )
175- if status == "succeeded" :
176- return AiCallResult (provider = provider .label , success = True ,
177- output = str (result .get ("output" , "" )), raw = result )
178- return AiCallResult (provider = provider .label , success = False ,
179- output = result .get ("error" , "unknown" ), raw = result )
180-
181- # Execute returns async job_id → poll
228+ # async job → poll
182229 job_id = result .get ("job_id" )
183230 if not isinstance (job_id , str ) or not job_id :
184231 return AiCallResult .unavailable (provider .label , "No job_id from gateway" )
185232
186- deadline = time .time () + timeout + 60
187- while time .time () < deadline :
188- time .sleep (5 )
189- req2 = urllib . request .Request (
190- f"{ base_url } /v1/codex-audit /jobs/{ job_id } " , method = "GET" ,
233+ deadline = _time .time () + timeout + 60
234+ while _time .time () < deadline :
235+ _time .sleep (5 )
236+ req2 = _urllib_req .Request (
237+ f"{ base_url } /v1/ai/execute /jobs/{ job_id } " , method = "GET" ,
191238 headers = {"Authorization" : f"Bearer { token } " , "Accept" : "application/json" ,
192239 "User-Agent" : "quant-strategy-lifecycle" },
193240 )
194241 try :
195- with urllib . request .urlopen (req2 , timeout = 30 ) as resp2 :
196- job = json .loads (resp2 .read ().decode ("utf-8" ))
197- except urllib . error .HTTPError :
242+ with _urllib_req .urlopen (req2 , timeout = 30 ) as resp2 :
243+ job = _json .loads (resp2 .read ().decode ("utf-8" ))
244+ except _urllib_err .HTTPError :
198245 continue
199246 status = job .get ("status" )
200247 if status == "succeeded" :
@@ -207,6 +254,14 @@ def _call(self, provider: AiProviderConfig, prompt: str, timeout: float) -> "AiC
207254 except Exception as exc :
208255 return AiCallResult .unavailable (provider .label , str (exc ))
209256
257+ @staticmethod
258+ def _map_provider_label (config : AiProviderConfig ) -> str :
259+ if config .provider == AiProviderId .CLAUDE :
260+ return "claude"
261+ if config .provider == AiProviderId .GPT :
262+ return "gpt"
263+ return "codex"
264+
210265
211266@dataclass (frozen = True )
212267class AiCallResult :
@@ -222,13 +277,15 @@ def unavailable(cls, provider: str, reason: str) -> "AiCallResult":
222277
223278
224279def _fetch_oidc_token (audience : str = "quant-codex-audit" ) -> str :
280+ import json as _json
281+ import urllib .request as _urllib_req
282+
225283 token_url = os .environ .get ("ACTIONS_ID_TOKEN_REQUEST_URL" , "" )
226284 token_bearer = os .environ .get ("ACTIONS_ID_TOKEN_REQUEST_TOKEN" , "" )
227285 if token_url and token_bearer :
228- req = urllib .request .Request (
229- f"{ token_url } &audience={ audience } " ,
230- headers = {"Authorization" : f"Bearer { token_bearer } " },
231- )
232- with urllib .request .urlopen (req , timeout = 10 ) as resp :
233- return str (json .loads (resp .read ().decode ("utf-8" )).get ("value" , "" ))
286+ separator = "&" if "?" in token_url else "?"
287+ url = f"{ token_url } { separator } audience={ _urllib_req .quote (audience , safe = '' )} "
288+ req = _urllib_req .Request (url , headers = {"Authorization" : f"Bearer { token_bearer } " })
289+ with _urllib_req .urlopen (req , timeout = 10 ) as resp :
290+ return str (_json .loads (resp .read ().decode ("utf-8" )).get ("value" , "" ))
234291 return os .environ .get ("CODEX_AUDIT_SERVICE_TOKEN" , "" ).strip ()
0 commit comments