sana0721/vertex
0
1"""2OpenAI handler module for creating clients and processing OpenAI Direct mode responses.3This module encapsulates all OpenAI-specific logic that was previously in chat_api.py.4"""5import json6import time7import asyncio8from typing import Dict, Any, AsyncGenerator9 10from fastapi.responses import JSONResponse, StreamingResponse11import openai12from google.auth.transport.requests import Request as AuthRequest13 14from models import OpenAIRequest15from config import VERTEX_REASONING_TAG16import config as app_config17from api_helpers import (18 create_openai_error_response,19 openai_fake_stream_generator,20 StreamingReasoningProcessor21)22from message_processing import extract_reasoning_by_tags23from credentials_manager import _refresh_auth24 25 26class OpenAIDirectHandler:27 """Handles OpenAI Direct mode operations including client creation and response processing."""28 29 def __init__(self, credential_manager):30 self.credential_manager = credential_manager31 self.safety_settings = [32 {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},33 {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},34 {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},35 {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},36 {"category": 'HARM_CATEGORY_CIVIC_INTEGRITY', "threshold": 'OFF'}37 ]38 39 def create_openai_client(self, project_id: str, gcp_token: str, location: str = "global") -> openai.AsyncOpenAI:40 """Create an OpenAI client configured for Vertex AI endpoint."""41 endpoint_url = (42 f"https://aiplatform.googleapis.com/v1beta1/"43 f"projects/{project_id}/locations/{location}/endpoints/openapi"44 )45 46 return openai.AsyncOpenAI(47 base_url=endpoint_url,48 api_key=gcp_token, # OAuth token49 )50 51 def prepare_openai_params(self, request: OpenAIRequest, model_id: str) -> Dict[str, Any]:52 """Prepare parameters for OpenAI API call."""53 params = {54 "model": model_id,55 "messages": [msg.model_dump(exclude_unset=True) for msg in request.messages],56 "temperature": request.temperature,57 "max_tokens": request.max_tokens,58 "top_p": request.top_p,59 "stream": request.stream,60 "stop": request.stop,61 "seed": request.seed,62 "n": request.n,63 }64 # Remove None values65 return {k: v for k, v in params.items() if v is not None}66 67 def prepare_extra_body(self) -> Dict[str, Any]:68 """Prepare extra body parameters for OpenAI API call."""69 return {70 "extra_body": {71 'google': {72 'safety_settings': self.safety_settings,73 'thought_tag_marker': VERTEX_REASONING_TAG74 }75 }76 }77 78 async def handle_streaming_response(79 self, 80 openai_client: openai.AsyncOpenAI,81 openai_params: Dict[str, Any],82 openai_extra_body: Dict[str, Any],83 request: OpenAIRequest84 ) -> StreamingResponse:85 """Handle streaming responses for OpenAI Direct mode."""86 if app_config.FAKE_STREAMING_ENABLED:87 print(f"INFO: OpenAI Fake Streaming (SSE Simulation) ENABLED for model '{request.model}'.")88 return StreamingResponse(89 openai_fake_stream_generator(90 openai_client=openai_client,91 openai_params=openai_params,92 openai_extra_body=openai_extra_body,93 request_obj=request,94 is_auto_attempt=False95 ),96 media_type="text/event-stream"97 )98 else:99 print(f"INFO: OpenAI True Streaming ENABLED for model '{request.model}'.")100 return StreamingResponse(101 self._true_stream_generator(openai_client, openai_params, openai_extra_body, request),102 media_type="text/event-stream"103 )104 105 async def _true_stream_generator(106 self,107 openai_client: openai.AsyncOpenAI,108 openai_params: Dict[str, Any],109 openai_extra_body: Dict[str, Any],110 request: OpenAIRequest111 ) -> AsyncGenerator[str, None]:112 """Generate true streaming response."""113 try:114 # Ensure stream=True is explicitly passed for real streaming115 openai_params_for_stream = {**openai_params, "stream": True}116 stream_response = await openai_client.chat.completions.create(117 **openai_params_for_stream,118 extra_body=openai_extra_body119 )120 121 # Create processor for tag-based extraction across chunks122 reasoning_processor = StreamingReasoningProcessor(VERTEX_REASONING_TAG)123 chunk_count = 0124 has_sent_content = False125 126 async for chunk in stream_response:127 chunk_count += 1128 try:129 chunk_as_dict = chunk.model_dump(exclude_unset=True, exclude_none=True)130 131 choices = chunk_as_dict.get('choices')132 if choices and isinstance(choices, list) and len(choices) > 0:133 delta = choices[0].get('delta')134 if delta and isinstance(delta, dict):135 # Always remove extra_content if present136 if 'extra_content' in delta:137 del delta['extra_content']138 139 content = delta.get('content', '')140 if content:141 # print(f"DEBUG: Chunk {chunk_count} - Raw content: '{content}'")142 # Use the processor to extract reasoning143 processed_content, current_reasoning = reasoning_processor.process_chunk(content)144 145 # Debug logging for processing results146 # if processed_content or current_reasoning:147 # print(f"DEBUG: Chunk {chunk_count} - Processed content: '{processed_content}', Reasoning: '{current_reasoning[:50]}...' if len(current_reasoning) > 50 else '{current_reasoning}'")148 149 # Send chunks for both reasoning and content as they arrive150 chunks_to_send = []151 152 # If we have reasoning content, send it153 if current_reasoning:154 reasoning_chunk = chunk_as_dict.copy()155 reasoning_chunk['choices'][0]['delta'] = {'reasoning_content': current_reasoning}156 chunks_to_send.append(reasoning_chunk)157 158 # If we have regular content, send it159 if processed_content:160 content_chunk = chunk_as_dict.copy()161 content_chunk['choices'][0]['delta'] = {'content': processed_content}162 chunks_to_send.append(content_chunk)163 has_sent_content = True164 165 # Send all chunks166 for chunk_to_send in chunks_to_send:167 yield f"data: {json.dumps(chunk_to_send)}\n\n"168 else:169 # Still yield the chunk even if no content (could have other delta fields)170 yield f"data: {json.dumps(chunk_as_dict)}\n\n"171 else:172 # Yield chunks without choices too (they might contain metadata)173 yield f"data: {json.dumps(chunk_as_dict)}\n\n"174 175 except Exception as chunk_error:176 error_msg = f"Error processing OpenAI chunk for {request.model}: {str(chunk_error)}"177 print(f"ERROR: {error_msg}")178 if len(error_msg) > 1024:179 error_msg = error_msg[:1024] + "..."180 error_response = create_openai_error_response(500, error_msg, "server_error")181 yield f"data: {json.dumps(error_response)}\n\n"182 yield "data: [DONE]\n\n"183 return184 185 # Debug logging for buffer state and chunk count186 # print(f"DEBUG: Stream ended after {chunk_count} chunks. Buffer state - tag_buffer: '{reasoning_processor.tag_buffer}', "187 # f"inside_tag: {reasoning_processor.inside_tag}, "188 # f"reasoning_buffer: '{reasoning_processor.reasoning_buffer[:50]}...' if reasoning_processor.reasoning_buffer else ''")189 190 # Flush any remaining buffered content191 remaining_content, remaining_reasoning = reasoning_processor.flush_remaining()192 193 # Send any remaining reasoning first194 if remaining_reasoning:195 # print(f"DEBUG: Flushing remaining reasoning: '{remaining_reasoning[:50]}...' if len(remaining_reasoning) > 50 else '{remaining_reasoning}'")196 reasoning_chunk = {197 "id": f"chatcmpl-{int(time.time())}",198 "object": "chat.completion.chunk",199 "created": int(time.time()),200 "model": request.model,201 "choices": [{"index": 0, "delta": {"reasoning_content": remaining_reasoning}, "finish_reason": None}]202 }203 yield f"data: {json.dumps(reasoning_chunk)}\n\n"204 205 # Send any remaining content206 if remaining_content:207 # print(f"DEBUG: Flushing remaining content: '{remaining_content}'")208 final_chunk = {209 "id": f"chatcmpl-{int(time.time())}",210 "object": "chat.completion.chunk",211 "created": int(time.time()),212 "model": request.model,213 "choices": [{"index": 0, "delta": {"content": remaining_content}, "finish_reason": None}]214 }215 yield f"data: {json.dumps(final_chunk)}\n\n"216 has_sent_content = True217 218 # Always send a finish reason chunk219 finish_chunk = {220 "id": f"chatcmpl-{int(time.time())}",221 "object": "chat.completion.chunk",222 "created": int(time.time()),223 "model": request.model,224 "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]225 }226 yield f"data: {json.dumps(finish_chunk)}\n\n"227 228 yield "data: [DONE]\n\n"229 230 except Exception as stream_error:231 error_msg = str(stream_error)232 if len(error_msg) > 1024:233 error_msg = error_msg[:1024] + "..."234 error_msg_full = f"Error during OpenAI streaming for {request.model}: {error_msg}"235 print(f"ERROR: {error_msg_full}")236 error_response = create_openai_error_response(500, error_msg_full, "server_error")237 yield f"data: {json.dumps(error_response)}\n\n"238 yield "data: [DONE]\n\n"239 240 async def handle_non_streaming_response(241 self,242 openai_client: openai.AsyncOpenAI,243 openai_params: Dict[str, Any],244 openai_extra_body: Dict[str, Any],245 request: OpenAIRequest246 ) -> JSONResponse:247 """Handle non-streaming responses for OpenAI Direct mode."""248 try:249 # Ensure stream=False is explicitly passed250 openai_params_non_stream = {**openai_params, "stream": False}251 response = await openai_client.chat.completions.create(252 **openai_params_non_stream,253 extra_body=openai_extra_body254 )255 response_dict = response.model_dump(exclude_unset=True, exclude_none=True)256 257 try:258 choices = response_dict.get('choices')259 if choices and isinstance(choices, list) and len(choices) > 0:260 message_dict = choices[0].get('message')261 if message_dict and isinstance(message_dict, dict):262 # Always remove extra_content from the message if it exists263 if 'extra_content' in message_dict:264 del message_dict['extra_content']265 266 # Extract reasoning from content267 full_content = message_dict.get('content')268 actual_content = full_content if isinstance(full_content, str) else ""269 270 if actual_content:271 print(f"INFO: OpenAI Direct Non-Streaming - Applying tag extraction with fixed marker: '{VERTEX_REASONING_TAG}'")272 reasoning_text, actual_content = extract_reasoning_by_tags(actual_content, VERTEX_REASONING_TAG)273 message_dict['content'] = actual_content274 if reasoning_text:275 message_dict['reasoning_content'] = reasoning_text276 # print(f"DEBUG: Tag extraction success. Reasoning len: {len(reasoning_text)}, Content len: {len(actual_content)}")277 # else:278 # print(f"DEBUG: No content found within fixed tag '{VERTEX_REASONING_TAG}'.")279 else:280 print(f"WARNING: OpenAI Direct Non-Streaming - No initial content found in message.")281 message_dict['content'] = ""282 283 except Exception as e_reasoning:284 print(f"WARNING: Error during non-streaming reasoning processing for model {request.model}: {e_reasoning}")285 286 return JSONResponse(content=response_dict)287 288 except Exception as e:289 error_msg = f"Error calling OpenAI client for {request.model}: {str(e)}"290 print(f"ERROR: {error_msg}")291 return JSONResponse(292 status_code=500, 293 content=create_openai_error_response(500, error_msg, "server_error")294 )295 296 async def process_request(self, request: OpenAIRequest, base_model_name: str):297 """Main entry point for processing OpenAI Direct mode requests."""298 print(f"INFO: Using OpenAI Direct Path for model: {request.model}")299 300 # Get credentials301 rotated_credentials, rotated_project_id = self.credential_manager.get_credentials()302 303 if not rotated_credentials or not rotated_project_id:304 error_msg = "OpenAI Direct Mode requires GCP credentials, but none were available or loaded successfully."305 print(f"ERROR: {error_msg}")306 return JSONResponse(307 status_code=500, 308 content=create_openai_error_response(500, error_msg, "server_error")309 )310 311 print(f"INFO: [OpenAI Direct Path] Using credentials for project: {rotated_project_id}")312 gcp_token = _refresh_auth(rotated_credentials)313 314 if not gcp_token:315 error_msg = f"Failed to obtain valid GCP token for OpenAI client (Project: {rotated_project_id})."316 print(f"ERROR: {error_msg}")317 return JSONResponse(318 status_code=500, 319 content=create_openai_error_response(500, error_msg, "server_error")320 )321 322 # Create client and prepare parameters323 openai_client = self.create_openai_client(rotated_project_id, gcp_token)324 model_id = f"google/{base_model_name}"325 openai_params = self.prepare_openai_params(request, model_id)326 openai_extra_body = self.prepare_extra_body()327 328 # Handle streaming vs non-streaming329 if request.stream:330 return await self.handle_streaming_response(331 openai_client, openai_params, openai_extra_body, request332 )333 else:334 return await self.handle_non_streaming_response(335 openai_client, openai_params, openai_extra_body, request336 )