Update app/main.py
Browse files- app/main.py +388 -143
app/main.py
CHANGED
@@ -66,6 +66,53 @@ async def get_api_key(authorization: Optional[str] = Header(None)):
|
|
66 |
|
67 |
return api_key
|
68 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
69 |
# Credential Manager for handling multiple service accounts
|
70 |
class CredentialManager:
|
71 |
def __init__(self, default_credentials_dir="/app/credentials"):
|
@@ -75,79 +122,224 @@ class CredentialManager:
|
|
75 |
self.current_index = 0
|
76 |
self.credentials = None
|
77 |
self.project_id = None
|
78 |
-
|
79 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
80 |
def load_credentials_list(self):
|
81 |
"""Load the list of available credential files"""
|
82 |
# Look for all .json files in the credentials directory
|
83 |
pattern = os.path.join(self.credentials_dir, "*.json")
|
84 |
self.credentials_files = glob.glob(pattern)
|
85 |
-
|
86 |
if not self.credentials_files:
|
87 |
# print(f"No credential files found in {self.credentials_dir}")
|
88 |
-
return False
|
89 |
-
|
90 |
-
|
91 |
-
|
92 |
-
|
|
|
|
|
93 |
def refresh_credentials_list(self):
|
94 |
-
"""Refresh the list of credential files
|
95 |
-
|
96 |
-
self.load_credentials_list()
|
97 |
-
|
98 |
-
|
99 |
-
if
|
100 |
-
print(f"Credential files updated: {
|
101 |
-
|
102 |
-
|
103 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
104 |
def get_next_credentials(self):
|
105 |
-
"""
|
106 |
-
|
107 |
-
|
108 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
try:
|
114 |
-
credentials = service_account.Credentials.from_service_account_file(file_path,scopes=['https://www.googleapis.com/auth/cloud-platform'])
|
115 |
-
project_id = credentials.project_id
|
116 |
-
print(f"Loaded credentials from {file_path} for project: {project_id}")
|
117 |
-
self.credentials = credentials
|
118 |
-
self.project_id = project_id
|
119 |
-
return credentials, project_id
|
120 |
-
except Exception as e:
|
121 |
-
print(f"Error loading credentials from {file_path}: {e}")
|
122 |
-
# Try the next file if this one fails
|
123 |
-
if len(self.credentials_files) > 1:
|
124 |
-
print("Trying next credential file...")
|
125 |
-
return self.get_next_credentials()
|
126 |
return None, None
|
127 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
128 |
def get_random_credentials(self):
|
129 |
-
"""Get a random credential file and load it"""
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
# Choose a random credential file
|
134 |
-
file_path = random.choice(self.credentials_files)
|
135 |
-
|
136 |
-
try:
|
137 |
-
credentials = service_account.Credentials.from_service_account_file(file_path,scopes=['https://www.googleapis.com/auth/cloud-platform'])
|
138 |
-
project_id = credentials.project_id
|
139 |
-
print(f"Loaded credentials from {file_path} for project: {project_id}")
|
140 |
-
self.credentials = credentials
|
141 |
-
self.project_id = project_id
|
142 |
-
return credentials, project_id
|
143 |
-
except Exception as e:
|
144 |
-
print(f"Error loading credentials from {file_path}: {e}")
|
145 |
-
# Try another random file if this one fails
|
146 |
-
if len(self.credentials_files) > 1:
|
147 |
-
print("Trying another credential file...")
|
148 |
-
return self.get_random_credentials()
|
149 |
return None, None
|
150 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
151 |
# Initialize the credential manager
|
152 |
credential_manager = CredentialManager()
|
153 |
|
@@ -190,59 +382,130 @@ class OpenAIRequest(BaseModel):
|
|
190 |
def init_vertex_ai():
|
191 |
global client # This will hold the fallback client if initialized
|
192 |
try:
|
193 |
-
# Priority 1: Check for credentials JSON content in environment variable
|
194 |
credentials_json_str = os.environ.get("GOOGLE_CREDENTIALS_JSON")
|
|
|
|
|
195 |
if credentials_json_str:
|
|
|
196 |
try:
|
197 |
-
#
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
-
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
-
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
216 |
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
228 |
try:
|
229 |
-
# Initialize
|
230 |
if client is None:
|
231 |
-
client = genai.Client(vertexai=True, credentials=
|
232 |
-
print(f"INFO: Initialized fallback Vertex AI client using
|
|
|
233 |
else:
|
234 |
-
|
235 |
-
|
236 |
-
|
237 |
-
except Exception as
|
238 |
-
print(f"ERROR: Failed to initialize
|
239 |
-
|
240 |
-
|
241 |
-
|
242 |
-
|
243 |
-
|
244 |
-
|
245 |
-
#
|
|
|
246 |
# Priority 2: Try to use the credential manager to get credentials from files
|
247 |
# We call get_next_credentials here mainly to validate it works and log the first file found
|
248 |
# The actual rotation happens per-request
|
@@ -1374,46 +1637,27 @@ async def chat_completions(request: OpenAIRequest, api_key: str = Depends(get_ap
|
|
1374 |
base_model_name = request.model.replace("-openai", "") # Extract base model name
|
1375 |
UNDERLYING_MODEL_ID = f"google/{base_model_name}" # Add google/ prefix
|
1376 |
|
1377 |
-
# --- Determine Credentials for OpenAI Client
|
|
|
|
|
1378 |
credentials_to_use = None
|
1379 |
project_id_to_use = None
|
1380 |
credential_source = "unknown"
|
1381 |
|
1382 |
-
|
1383 |
-
|
1384 |
-
|
1385 |
-
|
1386 |
-
|
1387 |
-
|
1388 |
-
|
1389 |
-
|
1390 |
-
|
1391 |
-
|
1392 |
-
|
1393 |
-
|
1394 |
-
|
1395 |
-
|
1396 |
-
project_id_to_use = project_id
|
1397 |
-
credential_source = "GOOGLE_CREDENTIALS_JSON env var"
|
1398 |
-
print(f"INFO: [OpenAI Path] Using credentials from {credential_source} for project: {project_id_to_use}")
|
1399 |
-
except Exception as e:
|
1400 |
-
print(f"WARNING: [OpenAI Path] Error processing GOOGLE_CREDENTIALS_JSON: {e}. Trying next method.")
|
1401 |
-
credentials_to_use = None # Ensure reset if failed
|
1402 |
-
|
1403 |
-
# Priority 2: Credential Manager (Rotated Files)
|
1404 |
-
if credentials_to_use is None:
|
1405 |
-
print(f"INFO: [OpenAI Path] Checking Credential Manager (directory: {credential_manager.credentials_dir})")
|
1406 |
-
rotated_credentials, rotated_project_id = credential_manager.get_next_credentials()
|
1407 |
-
if rotated_credentials and rotated_project_id:
|
1408 |
-
credentials_to_use = rotated_credentials
|
1409 |
-
project_id_to_use = rotated_project_id
|
1410 |
-
credential_source = f"Credential Manager file (Index: {credential_manager.current_index -1 if credential_manager.current_index > 0 else len(credential_manager.credentials_files) - 1})"
|
1411 |
-
print(f"INFO: [OpenAI Path] Using credentials from {credential_source} for project: {project_id_to_use}")
|
1412 |
-
else:
|
1413 |
-
print(f"INFO: [OpenAI Path] No credentials loaded via Credential Manager.")
|
1414 |
-
|
1415 |
-
# Priority 3: GOOGLE_APPLICATION_CREDENTIALS (File Path in Env Var)
|
1416 |
-
if credentials_to_use is None:
|
1417 |
file_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
|
1418 |
if file_path:
|
1419 |
print(f"INFO: [OpenAI Path] Checking GOOGLE_APPLICATION_CREDENTIALS file path: {file_path}")
|
@@ -1432,9 +1676,10 @@ async def chat_completions(request: OpenAIRequest, api_key: str = Depends(get_ap
|
|
1432 |
else:
|
1433 |
print(f"ERROR: [OpenAI Path] GOOGLE_APPLICATION_CREDENTIALS file does not exist at path: {file_path}")
|
1434 |
|
|
|
1435 |
# Error if no credentials found after all checks
|
1436 |
if credentials_to_use is None or project_id_to_use is None:
|
1437 |
-
error_msg = "No valid credentials found for OpenAI client path.
|
1438 |
print(f"ERROR: {error_msg}")
|
1439 |
error_response = create_openai_error_response(500, error_msg, "server_error")
|
1440 |
return JSONResponse(status_code=500, content=error_response)
|
|
|
66 |
|
67 |
return api_key
|
68 |
|
69 |
+
# Helper function to parse multiple JSONs from a string
|
70 |
+
def parse_multiple_json_credentials(json_str: str) -> List[Dict[str, Any]]:
|
71 |
+
"""
|
72 |
+
Parse multiple JSON objects from a string separated by commas.
|
73 |
+
Format expected: {json_object1},{json_object2},...
|
74 |
+
Returns a list of parsed JSON objects.
|
75 |
+
"""
|
76 |
+
credentials_list = []
|
77 |
+
nesting_level = 0
|
78 |
+
current_object_start = -1
|
79 |
+
str_length = len(json_str)
|
80 |
+
|
81 |
+
for i, char in enumerate(json_str):
|
82 |
+
if char == '{':
|
83 |
+
if nesting_level == 0:
|
84 |
+
current_object_start = i
|
85 |
+
nesting_level += 1
|
86 |
+
elif char == '}':
|
87 |
+
if nesting_level > 0:
|
88 |
+
nesting_level -= 1
|
89 |
+
if nesting_level == 0 and current_object_start != -1:
|
90 |
+
# Found a complete top-level JSON object
|
91 |
+
json_object_str = json_str[current_object_start : i + 1]
|
92 |
+
try:
|
93 |
+
credentials_info = json.loads(json_object_str)
|
94 |
+
# Basic validation for service account structure
|
95 |
+
required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]
|
96 |
+
if all(field in credentials_info for field in required_fields):
|
97 |
+
credentials_list.append(credentials_info)
|
98 |
+
print(f"DEBUG: Successfully parsed a JSON credential object.")
|
99 |
+
else:
|
100 |
+
print(f"WARNING: Parsed JSON object missing required fields: {json_object_str[:100]}...")
|
101 |
+
except json.JSONDecodeError as e:
|
102 |
+
print(f"ERROR: Failed to parse JSON object segment: {json_object_str[:100]}... Error: {e}")
|
103 |
+
current_object_start = -1 # Reset for the next object
|
104 |
+
else:
|
105 |
+
# Found a closing brace without a matching open brace in scope, might indicate malformed input
|
106 |
+
print(f"WARNING: Encountered unexpected '}}' at index {i}. Input might be malformed.")
|
107 |
+
|
108 |
+
|
109 |
+
if nesting_level != 0:
|
110 |
+
print(f"WARNING: JSON string parsing ended with non-zero nesting level ({nesting_level}). Check for unbalanced braces.")
|
111 |
+
|
112 |
+
print(f"DEBUG: Parsed {len(credentials_list)} credential objects from the input string.")
|
113 |
+
return credentials_list
|
114 |
+
|
115 |
+
|
116 |
# Credential Manager for handling multiple service accounts
|
117 |
class CredentialManager:
|
118 |
def __init__(self, default_credentials_dir="/app/credentials"):
|
|
|
122 |
self.current_index = 0
|
123 |
self.credentials = None
|
124 |
self.project_id = None
|
125 |
+
# New: Store credentials loaded directly from JSON objects
|
126 |
+
self.in_memory_credentials: List[Dict[str, Any]] = []
|
127 |
+
self.load_credentials_list() # Load file-based credentials initially
|
128 |
+
|
129 |
+
def add_credential_from_json(self, credentials_info: Dict[str, Any]) -> bool:
|
130 |
+
"""
|
131 |
+
Add a credential from a JSON object to the manager's in-memory list.
|
132 |
+
|
133 |
+
Args:
|
134 |
+
credentials_info: Dict containing service account credentials
|
135 |
+
|
136 |
+
Returns:
|
137 |
+
bool: True if credential was added successfully, False otherwise
|
138 |
+
"""
|
139 |
+
try:
|
140 |
+
# Validate structure again before creating credentials object
|
141 |
+
required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]
|
142 |
+
if not all(field in credentials_info for field in required_fields):
|
143 |
+
print(f"WARNING: Skipping JSON credential due to missing required fields.")
|
144 |
+
return False
|
145 |
+
|
146 |
+
credentials = service_account.Credentials.from_service_account_info(
|
147 |
+
credentials_info,
|
148 |
+
scopes=['https://www.googleapis.com/auth/cloud-platform']
|
149 |
+
)
|
150 |
+
project_id = credentials.project_id
|
151 |
+
print(f"DEBUG: Successfully created credentials object from JSON for project: {project_id}")
|
152 |
+
|
153 |
+
# Store the credentials object and project ID
|
154 |
+
self.in_memory_credentials.append({
|
155 |
+
'credentials': credentials,
|
156 |
+
'project_id': project_id,
|
157 |
+
'source': 'json_string' # Add source for clarity
|
158 |
+
})
|
159 |
+
print(f"INFO: Added credential for project {project_id} from JSON string to Credential Manager.")
|
160 |
+
return True
|
161 |
+
except Exception as e:
|
162 |
+
print(f"ERROR: Failed to create credentials from parsed JSON object: {e}")
|
163 |
+
return False
|
164 |
+
|
165 |
+
def load_credentials_from_json_list(self, json_list: List[Dict[str, Any]]) -> int:
|
166 |
+
"""
|
167 |
+
Load multiple credentials from a list of JSON objects into memory.
|
168 |
+
|
169 |
+
Args:
|
170 |
+
json_list: List of dicts containing service account credentials
|
171 |
+
|
172 |
+
Returns:
|
173 |
+
int: Number of credentials successfully loaded
|
174 |
+
"""
|
175 |
+
# Avoid duplicates if called multiple times
|
176 |
+
existing_projects = {cred['project_id'] for cred in self.in_memory_credentials}
|
177 |
+
success_count = 0
|
178 |
+
newly_added_projects = set()
|
179 |
+
|
180 |
+
for credentials_info in json_list:
|
181 |
+
project_id = credentials_info.get('project_id')
|
182 |
+
# Check if this project_id from JSON exists in files OR already added from JSON
|
183 |
+
is_duplicate_file = any(os.path.basename(f) == f"{project_id}.json" for f in self.credentials_files) # Basic check
|
184 |
+
is_duplicate_mem = project_id in existing_projects or project_id in newly_added_projects
|
185 |
+
|
186 |
+
if project_id and not is_duplicate_file and not is_duplicate_mem:
|
187 |
+
if self.add_credential_from_json(credentials_info):
|
188 |
+
success_count += 1
|
189 |
+
newly_added_projects.add(project_id)
|
190 |
+
elif project_id:
|
191 |
+
print(f"DEBUG: Skipping duplicate credential for project {project_id} from JSON list.")
|
192 |
+
|
193 |
+
|
194 |
+
if success_count > 0:
|
195 |
+
print(f"INFO: Loaded {success_count} new credentials from JSON list into memory.")
|
196 |
+
return success_count
|
197 |
+
|
198 |
def load_credentials_list(self):
|
199 |
"""Load the list of available credential files"""
|
200 |
# Look for all .json files in the credentials directory
|
201 |
pattern = os.path.join(self.credentials_dir, "*.json")
|
202 |
self.credentials_files = glob.glob(pattern)
|
203 |
+
|
204 |
if not self.credentials_files:
|
205 |
# print(f"No credential files found in {self.credentials_dir}")
|
206 |
+
pass # Don't return False yet, might have in-memory creds
|
207 |
+
else:
|
208 |
+
print(f"Found {len(self.credentials_files)} credential files: {[os.path.basename(f) for f in self.credentials_files]}")
|
209 |
+
|
210 |
+
# Check total credentials
|
211 |
+
return self.get_total_credentials() > 0
|
212 |
+
|
213 |
def refresh_credentials_list(self):
|
214 |
+
"""Refresh the list of credential files and return if any credentials exist"""
|
215 |
+
old_file_count = len(self.credentials_files)
|
216 |
+
self.load_credentials_list() # Reloads file list
|
217 |
+
new_file_count = len(self.credentials_files)
|
218 |
+
|
219 |
+
if old_file_count != new_file_count:
|
220 |
+
print(f"Credential files updated: {old_file_count} -> {new_file_count}")
|
221 |
+
|
222 |
+
# Total credentials = files + in-memory
|
223 |
+
total_credentials = self.get_total_credentials()
|
224 |
+
print(f"DEBUG: Refresh check - Total credentials available: {total_credentials}")
|
225 |
+
return total_credentials > 0
|
226 |
+
|
227 |
+
def get_total_credentials(self):
|
228 |
+
"""Returns the total number of credentials (file + in-memory)."""
|
229 |
+
return len(self.credentials_files) + len(self.in_memory_credentials)
|
230 |
+
|
231 |
def get_next_credentials(self):
|
232 |
+
"""
|
233 |
+
Rotate to the next credential (file or in-memory) and return it.
|
234 |
+
"""
|
235 |
+
total_credentials = self.get_total_credentials()
|
236 |
+
|
237 |
+
if total_credentials == 0:
|
238 |
+
print("WARNING: No credentials available in Credential Manager (files or in-memory).")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
239 |
return None, None
|
240 |
+
|
241 |
+
# Determine which credential (file or in-memory) to use based on the current index
|
242 |
+
# Use a temporary index for calculation to avoid modifying self.current_index prematurely
|
243 |
+
effective_index_to_use = self.current_index % total_credentials
|
244 |
+
num_files = len(self.credentials_files)
|
245 |
+
|
246 |
+
# Advance the main index *after* deciding which one to use for this call
|
247 |
+
self.current_index = (self.current_index + 1) % total_credentials
|
248 |
+
|
249 |
+
if effective_index_to_use < num_files:
|
250 |
+
# It's a file-based credential
|
251 |
+
file_path = self.credentials_files[effective_index_to_use]
|
252 |
+
print(f"DEBUG: Attempting to load credential from file: {os.path.basename(file_path)} (Index {effective_index_to_use})")
|
253 |
+
try:
|
254 |
+
credentials = service_account.Credentials.from_service_account_file(
|
255 |
+
file_path,
|
256 |
+
scopes=['https://www.googleapis.com/auth/cloud-platform']
|
257 |
+
)
|
258 |
+
project_id = credentials.project_id
|
259 |
+
print(f"INFO: Rotated to credential file: {os.path.basename(file_path)} for project: {project_id}")
|
260 |
+
self.credentials = credentials # Cache last used
|
261 |
+
self.project_id = project_id # Cache last used
|
262 |
+
return credentials, project_id
|
263 |
+
except Exception as e:
|
264 |
+
print(f"ERROR: Failed loading credentials from file {os.path.basename(file_path)}: {e}. Skipping.")
|
265 |
+
# Try the next available credential recursively IF there are others available
|
266 |
+
if total_credentials > 1:
|
267 |
+
print("DEBUG: Attempting to get next credential after file load error...")
|
268 |
+
# The index is already advanced, so calling again should try the next one
|
269 |
+
# Need to ensure we don't get stuck in infinite loop if all fail
|
270 |
+
# Let's limit recursion depth or track failed indices (simpler: rely on index advance)
|
271 |
+
# The index was already advanced, so calling again will try the next one
|
272 |
+
return self.get_next_credentials()
|
273 |
+
else:
|
274 |
+
print("ERROR: Only one credential (file) available and it failed to load.")
|
275 |
+
return None, None # No more credentials to try
|
276 |
+
else:
|
277 |
+
# It's an in-memory credential
|
278 |
+
in_memory_index = effective_index_to_use - num_files
|
279 |
+
if in_memory_index < len(self.in_memory_credentials):
|
280 |
+
cred_info = self.in_memory_credentials[in_memory_index]
|
281 |
+
credentials = cred_info['credentials']
|
282 |
+
project_id = cred_info['project_id']
|
283 |
+
print(f"INFO: Rotated to in-memory credential for project: {project_id} (Index {in_memory_index})")
|
284 |
+
# TODO: Add handling for expired in-memory credentials if needed (refresh?)
|
285 |
+
# For now, assume they are valid when loaded
|
286 |
+
self.credentials = credentials # Cache last used
|
287 |
+
self.project_id = project_id # Cache last used
|
288 |
+
return credentials, project_id
|
289 |
+
else:
|
290 |
+
# This case should not happen with correct modulo arithmetic, but added defensively
|
291 |
+
print(f"ERROR: Calculated in-memory index {in_memory_index} is out of bounds.")
|
292 |
+
return None, None
|
293 |
+
|
294 |
+
|
295 |
def get_random_credentials(self):
|
296 |
+
"""Get a random credential (file or in-memory) and load it"""
|
297 |
+
total_credentials = self.get_total_credentials()
|
298 |
+
if total_credentials == 0:
|
299 |
+
print("WARNING: No credentials available for random selection.")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
300 |
return None, None
|
301 |
|
302 |
+
random_index = random.randrange(total_credentials)
|
303 |
+
num_files = len(self.credentials_files)
|
304 |
+
|
305 |
+
if random_index < num_files:
|
306 |
+
# Selected a file-based credential
|
307 |
+
file_path = self.credentials_files[random_index]
|
308 |
+
print(f"DEBUG: Randomly selected file: {os.path.basename(file_path)}")
|
309 |
+
try:
|
310 |
+
credentials = service_account.Credentials.from_service_account_file(
|
311 |
+
file_path,
|
312 |
+
scopes=['https://www.googleapis.com/auth/cloud-platform']
|
313 |
+
)
|
314 |
+
project_id = credentials.project_id
|
315 |
+
print(f"INFO: Loaded random credential from file {os.path.basename(file_path)} for project: {project_id}")
|
316 |
+
self.credentials = credentials # Cache last used
|
317 |
+
self.project_id = project_id # Cache last used
|
318 |
+
return credentials, project_id
|
319 |
+
except Exception as e:
|
320 |
+
print(f"ERROR: Failed loading random credentials file {os.path.basename(file_path)}: {e}. Trying again.")
|
321 |
+
# Try another random credential if this one fails and others exist
|
322 |
+
if total_credentials > 1:
|
323 |
+
return self.get_random_credentials() # Recursive call
|
324 |
+
else:
|
325 |
+
print("ERROR: Only one credential (file) available and it failed to load.")
|
326 |
+
return None, None
|
327 |
+
else:
|
328 |
+
# Selected an in-memory credential
|
329 |
+
in_memory_index = random_index - num_files
|
330 |
+
if in_memory_index < len(self.in_memory_credentials):
|
331 |
+
cred_info = self.in_memory_credentials[in_memory_index]
|
332 |
+
credentials = cred_info['credentials']
|
333 |
+
project_id = cred_info['project_id']
|
334 |
+
print(f"INFO: Loaded random in-memory credential for project: {project_id}")
|
335 |
+
self.credentials = credentials # Cache last used
|
336 |
+
self.project_id = project_id # Cache last used
|
337 |
+
return credentials, project_id
|
338 |
+
else:
|
339 |
+
# Defensive case
|
340 |
+
print(f"ERROR: Calculated random in-memory index {in_memory_index} is out of bounds.")
|
341 |
+
return None, None
|
342 |
+
|
343 |
# Initialize the credential manager
|
344 |
credential_manager = CredentialManager()
|
345 |
|
|
|
382 |
def init_vertex_ai():
|
383 |
global client # This will hold the fallback client if initialized
|
384 |
try:
|
385 |
+
# Priority 1: Check for credentials JSON content in environment variable
|
386 |
credentials_json_str = os.environ.get("GOOGLE_CREDENTIALS_JSON")
|
387 |
+
json_loaded_successfully = False # Flag to track if we succeed via JSON string(s)
|
388 |
+
|
389 |
if credentials_json_str:
|
390 |
+
print("INFO: Found GOOGLE_CREDENTIALS_JSON environment variable. Attempting to load.")
|
391 |
try:
|
392 |
+
# --- Attempt 1: Parse as multiple JSON objects ---
|
393 |
+
json_objects = parse_multiple_json_credentials(credentials_json_str)
|
394 |
+
|
395 |
+
if json_objects:
|
396 |
+
print(f"DEBUG: Parsed {len(json_objects)} potential credential objects from GOOGLE_CREDENTIALS_JSON.")
|
397 |
+
# Add all valid credentials to the credential manager's in-memory list
|
398 |
+
success_count = credential_manager.load_credentials_from_json_list(json_objects)
|
399 |
+
|
400 |
+
if success_count > 0:
|
401 |
+
print(f"INFO: Successfully loaded {success_count} credentials from GOOGLE_CREDENTIALS_JSON into manager.")
|
402 |
+
# Initialize the fallback client with the first *successfully loaded* in-memory credential if needed
|
403 |
+
if client is None and credential_manager.in_memory_credentials:
|
404 |
+
try:
|
405 |
+
first_cred_info = credential_manager.in_memory_credentials[0]
|
406 |
+
first_credentials = first_cred_info['credentials']
|
407 |
+
first_project_id = first_cred_info['project_id']
|
408 |
+
client = genai.Client(
|
409 |
+
vertexai=True,
|
410 |
+
credentials=first_credentials,
|
411 |
+
project=first_project_id,
|
412 |
+
location="us-central1"
|
413 |
+
)
|
414 |
+
print(f"INFO: Initialized fallback Vertex AI client using first credential from GOOGLE_CREDENTIALS_JSON (Project: {first_project_id})")
|
415 |
+
json_loaded_successfully = True
|
416 |
+
except Exception as client_init_err:
|
417 |
+
print(f"ERROR: Failed to initialize genai.Client from first GOOGLE_CREDENTIALS_JSON object: {client_init_err}")
|
418 |
+
# Don't return yet, let it fall through to other methods if client init failed
|
419 |
+
elif client is not None:
|
420 |
+
print("INFO: Fallback client already initialized. GOOGLE_CREDENTIALS_JSON validated.")
|
421 |
+
json_loaded_successfully = True
|
422 |
+
# If client is None but loading failed to add any to manager, json_loaded_successfully remains False
|
423 |
+
|
424 |
+
# If we successfully loaded JSON creds AND initialized/validated the client, we are done with Priority 1
|
425 |
+
if json_loaded_successfully:
|
426 |
+
return True # Exit early, Priority 1 succeeded
|
427 |
+
|
428 |
+
# --- Attempt 2: If multiple parsing didn't yield results, try parsing as a single JSON object ---
|
429 |
+
if not json_loaded_successfully: # Or if json_objects was empty
|
430 |
+
print("DEBUG: Multi-JSON parsing did not yield usable credentials or failed client init. Attempting single JSON parse...")
|
431 |
+
try:
|
432 |
+
credentials_info = json.loads(credentials_json_str)
|
433 |
+
# Check structure (redundant with add_credential_from_json, but good defense)
|
434 |
+
if not isinstance(credentials_info, dict):
|
435 |
+
raise ValueError("Credentials JSON must be a dictionary")
|
436 |
+
required_fields = ["type", "project_id", "private_key_id", "private_key", "client_email"]
|
437 |
+
if not all(field in credentials_info for field in required_fields):
|
438 |
+
raise ValueError(f"Credentials JSON missing required fields")
|
439 |
+
|
440 |
+
# Add this single credential to the manager
|
441 |
+
if credential_manager.add_credential_from_json(credentials_info):
|
442 |
+
print("INFO: Successfully loaded single credential from GOOGLE_CREDENTIALS_JSON into manager.")
|
443 |
+
# Initialize client if needed, using the newly added credential
|
444 |
+
if client is None and credential_manager.in_memory_credentials: # Should have 1 now
|
445 |
+
try:
|
446 |
+
# Get the last added credential (which is the first/only one here)
|
447 |
+
last_cred_info = credential_manager.in_memory_credentials[-1]
|
448 |
+
single_credentials = last_cred_info['credentials']
|
449 |
+
single_project_id = last_cred_info['project_id']
|
450 |
+
client = genai.Client(
|
451 |
+
vertexai=True,
|
452 |
+
credentials=single_credentials,
|
453 |
+
project=single_project_id,
|
454 |
+
location="us-central1"
|
455 |
+
)
|
456 |
+
print(f"INFO: Initialized fallback Vertex AI client using single credential from GOOGLE_CREDENTIALS_JSON (Project: {single_project_id})")
|
457 |
+
json_loaded_successfully = True
|
458 |
+
except Exception as client_init_err:
|
459 |
+
print(f"ERROR: Failed to initialize genai.Client from single GOOGLE_CREDENTIALS_JSON object: {client_init_err}")
|
460 |
+
elif client is not None:
|
461 |
+
print("INFO: Fallback client already initialized. Single GOOGLE_CREDENTIALS_JSON validated.")
|
462 |
+
json_loaded_successfully = True
|
463 |
+
|
464 |
+
# If successful, exit
|
465 |
+
if json_loaded_successfully:
|
466 |
+
return True # Exit early, Priority 1 succeeded (as single JSON)
|
467 |
+
|
468 |
+
except Exception as single_json_err:
|
469 |
+
print(f"WARNING: GOOGLE_CREDENTIALS_JSON could not be parsed as single valid JSON: {single_json_err}. Proceeding to other methods.")
|
470 |
|
471 |
+
except Exception as e:
|
472 |
+
# Catch errors during multi-JSON parsing or loading
|
473 |
+
print(f"WARNING: Error processing GOOGLE_CREDENTIALS_JSON (multi-parse/load attempt): {e}. Will try other methods.")
|
474 |
+
# Ensure flag is False and fall through
|
475 |
+
|
476 |
+
# If GOOGLE_CREDENTIALS_JSON didn't exist or failed to yield a usable client...
|
477 |
+
if not json_loaded_successfully:
|
478 |
+
print(f"INFO: GOOGLE_CREDENTIALS_JSON did not provide usable credentials. Checking filesystem via Credential Manager (directory: {credential_manager.credentials_dir}).")
|
479 |
+
|
480 |
+
# Priority 2: Try Credential Manager (files from directory)
|
481 |
+
# Refresh file list AND check if *any* credentials (file or pre-loaded JSON) exist
|
482 |
+
if credential_manager.refresh_credentials_list(): # Checks total count now
|
483 |
+
# Attempt to get the *next* available credential (could be file or JSON loaded earlier)
|
484 |
+
# We call get_next_credentials here mainly to validate it works and log the first valid one found
|
485 |
+
# The actual rotation happens per-request
|
486 |
+
cm_credentials, cm_project_id = credential_manager.get_next_credentials()
|
487 |
+
|
488 |
+
if cm_credentials and cm_project_id:
|
489 |
try:
|
490 |
+
# Initialize global client ONLY if it hasn't been set by Priority 1
|
491 |
if client is None:
|
492 |
+
client = genai.Client(vertexai=True, credentials=cm_credentials, project=cm_project_id, location="us-central1")
|
493 |
+
print(f"INFO: Initialized fallback Vertex AI client using Credential Manager (Source: {'File' if credential_manager.current_index <= len(credential_manager.credentials_files) else 'JSON'}) for project: {cm_project_id}")
|
494 |
+
return True # Successfully initialized global client via Cred Manager
|
495 |
else:
|
496 |
+
# Client was already initialized (likely by JSON string), but we validated CM works too.
|
497 |
+
print(f"INFO: Fallback client already initialized. Credential Manager source validated for project: {cm_project_id}")
|
498 |
+
# Don't return True here if client was already set, let it fall through to check GAC if needed (though unlikely needed now)
|
499 |
+
except Exception as e:
|
500 |
+
print(f"ERROR: Failed to initialize client with credentials from Credential Manager source: {e}")
|
501 |
+
else:
|
502 |
+
# This might happen if get_next_credentials itself failed internally
|
503 |
+
print(f"INFO: Credential Manager get_next_credentials() returned None.")
|
504 |
+
else:
|
505 |
+
print("INFO: No credentials found via Credential Manager (files or JSON string).")
|
506 |
+
|
507 |
+
# Priority 3: Fall back to GOOGLE_APPLICATION_CREDENTIALS environment variable (file path)
|
508 |
+
# This should only run if client is STILL None after JSON and CM attempts
|
509 |
# Priority 2: Try to use the credential manager to get credentials from files
|
510 |
# We call get_next_credentials here mainly to validate it works and log the first file found
|
511 |
# The actual rotation happens per-request
|
|
|
1637 |
base_model_name = request.model.replace("-openai", "") # Extract base model name
|
1638 |
UNDERLYING_MODEL_ID = f"google/{base_model_name}" # Add google/ prefix
|
1639 |
|
1640 |
+
# --- Determine Credentials for OpenAI Client using Credential Manager ---
|
1641 |
+
# The init_vertex_ai function already loaded JSON credentials into the manager if available.
|
1642 |
+
# Now, we just need to get the next available credential using the manager's rotation.
|
1643 |
credentials_to_use = None
|
1644 |
project_id_to_use = None
|
1645 |
credential_source = "unknown"
|
1646 |
|
1647 |
+
print(f"INFO: [OpenAI Path] Attempting to get next credential from Credential Manager...")
|
1648 |
+
# This will rotate through file-based and JSON-based credentials loaded during startup
|
1649 |
+
rotated_credentials, rotated_project_id = credential_manager.get_next_credentials()
|
1650 |
+
|
1651 |
+
if rotated_credentials and rotated_project_id:
|
1652 |
+
credentials_to_use = rotated_credentials
|
1653 |
+
project_id_to_use = rotated_project_id
|
1654 |
+
# Determine if it came from file or JSON (crude check based on structure)
|
1655 |
+
source_type = "In-Memory JSON" if hasattr(rotated_credentials, '_service_account_email') else "File" # Heuristic
|
1656 |
+
credential_source = f"Credential Manager ({source_type})"
|
1657 |
+
print(f"INFO: [OpenAI Path] Using credentials from {credential_source} for project: {project_id_to_use}")
|
1658 |
+
else:
|
1659 |
+
print(f"INFO: [OpenAI Path] Credential Manager did not provide credentials. Checking GOOGLE_APPLICATION_CREDENTIALS fallback.")
|
1660 |
+
# Priority 3 (Fallback): GOOGLE_APPLICATION_CREDENTIALS (File Path in Env Var)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1661 |
file_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
|
1662 |
if file_path:
|
1663 |
print(f"INFO: [OpenAI Path] Checking GOOGLE_APPLICATION_CREDENTIALS file path: {file_path}")
|
|
|
1676 |
else:
|
1677 |
print(f"ERROR: [OpenAI Path] GOOGLE_APPLICATION_CREDENTIALS file does not exist at path: {file_path}")
|
1678 |
|
1679 |
+
|
1680 |
# Error if no credentials found after all checks
|
1681 |
if credentials_to_use is None or project_id_to_use is None:
|
1682 |
+
error_msg = "No valid credentials found for OpenAI client path. Checked Credential Manager (JSON/Files) and GOOGLE_APPLICATION_CREDENTIALS."
|
1683 |
print(f"ERROR: {error_msg}")
|
1684 |
error_response = create_openai_error_response(500, error_msg, "server_error")
|
1685 |
return JSONResponse(status_code=500, content=error_response)
|