Update backend_api.py

backend_api.py

@@ -3,7 +3,7 @@ FastAPI backend for AnyCoder - provides REST API endpoints
 """
 from fastapi import FastAPI, HTTPException, Header, WebSocket, WebSocketDisconnect, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
-import google.generativeai as genai
+from google import genai
 from fastapi.responses import StreamingResponse, RedirectResponse, JSONResponse
 from pydantic import BaseModel
 from typing import Optional, List, Dict, AsyncGenerator
@@ -97,8 +97,21 @@ def get_gemini_model():
 
 # Define models and languages here to avoid importing Gradio UI
 AVAILABLE_MODELS = [
-    {"name": "Gemini 1.5 Flash", "id": "gemini-1.5-flash", "description": "Tốc độ nhanh, xử lý code dài cực tốt"},
-    {"name": "Gemini 1.5 Pro", "id": "gemini-1.5-pro", "description": "Thông minh nhất, hỗ trợ logic phức tạp"}
+    {
+        "name": "Gemini 3.5 Flash", 
+        "id": "gemini-3.5-flash", 
+        "description": "Tốt nhất để viết code dài, tốc độ nhanh, hỗ trợ context lớn"
+    },
+    {
+        "name": "Gemini 3.1 Pro Preview", 
+        "id": "gemini-3.1-pro-preview", 
+        "description": "Thông minh nhất, độ chính xác cao cho logic phức tạp"
+    },
+    {
+        "name": "Gemini 2.0 Flash", 
+        "id": "gemini-2.0-flash", 
+        "description": "Cân bằng giữa hiệu năng và ổn định"
+    }
 ]
 
 # Cache model lookup for faster access (built after AVAILABLE_MODELS is defined)
@@ -738,29 +751,22 @@ async def generate_code(request: CodeGenerationRequest, authorization: Optional[
     if not model_id.startswith("models/"):
         model_id = f"models/{model_id}"
     
-    async def event_stream() -> AsyncGenerator[str, None]:
+    async def event_stream():
         try:
-            # 1. Khởi tạo model với ID đã chuẩn hóa
-            model = genai.GenerativeModel(model_name=model_id)
-            
-            # 2. Định nghĩa System Prompt
-            system_prompt = SYSTEM_PROMPT_CACHE.get(language, GENERIC_SYSTEM_PROMPT.format(language=language))
-            full_prompt = f"{system_prompt}\n\nTask: {query}"
-            
-            # 3. Gọi API Gemini stream
-            response = model.generate_content(full_prompt, stream=True)
+            # Dùng client mới để gọi model
+            # Lưu ý: Thư viện mới dùng model_id trực tiếp, không cần tiền tố 'models/' nếu gọi qua client
+            response = client.models.generate_content_stream(
+                model=request.model_id, 
+                contents=full_prompt
+            )
             
-            # 4. Stream phản hồi về frontend
             for chunk in response:
                 if chunk.text:
                     yield f"data: {json.dumps({'type': 'chunk', 'content': chunk.text})}\n\n"
-                    await asyncio.sleep(0.001)
-
-            # Gửi tín hiệu hoàn tất
+            
             yield f"data: {json.dumps({'type': 'complete'})}\n\n"
-
+            
         except Exception as e:
-            # Gửi thông báo lỗi chi tiết để debug
             yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
             
     return StreamingResponse(event_stream(), media_type="text/event-stream")