Update app.py

app.py

@@ -160,75 +160,64 @@ async def direct_chat(payload: ChatPayload):
     custom_api_key_secret = os.getenv("CUSTOM_API_SECRET_KEY")
     custom_api_base_url = os.getenv("CUSTOM_API_BASE_URL", CUSTOM_API_BASE_URL_DEFAULT)
     custom_api_model = os.getenv("CUSTOM_API_MODEL", CUSTOM_API_MODEL_DEFAULT)
-
+    
     if not custom_api_key_secret:
         logger.error("Custom API key ('CUSTOM_API_SECRET_KEY') is not configured for /chat.")
         raise HTTPException(status_code=500, detail="Custom API key not configured.")
-
-    client = openai.OpenAI(api_key=custom_api_key_secret, base_url=custom_api_base_url)
-
+    
     async def custom_api_streamer():
+        client = None
         try:
             logger.info("Sending request to Custom API for /chat.")
-            stream = client.chat.completions.create(
+            
+            # Use AsyncOpenAI for proper async handling
+            from openai import AsyncOpenAI
+            client = AsyncOpenAI(
+                api_key=custom_api_key_secret, 
+                base_url=custom_api_base_url,
+                timeout=30.0  # Add timeout to prevent hanging connections
+            )
+            
+            stream = await client.chat.completions.create(
                 model=custom_api_model,
                 temperature=payload.temperature,
                 messages=[{"role": "user", "content": payload.message}],
                 stream=True
             )
-            for chunk in stream:
-                content_to_yield = None
-                if hasattr(chunk.choices[0].delta, "reasoning_content") and chunk.choices[0].delta.reasoning_content:
-                    content_to_yield = chunk.choices[0].delta.reasoning_content
-                elif chunk.choices[0].delta and chunk.choices[0].delta.content:
-                    content_to_yield = chunk.choices[0].delta.content
-                if content_to_yield:
-                    yield content_to_yield
+            
+            async for chunk in stream:
+                try:
+                    content_to_yield = None
+                    if hasattr(chunk.choices[0].delta, "reasoning_content") and chunk.choices[0].delta.reasoning_content:
+                        content_to_yield = chunk.choices[0].delta.reasoning_content
+                    elif chunk.choices[0].delta and chunk.choices[0].delta.content:
+                        content_to_yield = chunk.choices[0].delta.content
+                    
+                    if content_to_yield:
+                        yield content_to_yield
+                        
+                except (IndexError, AttributeError) as e:
+                    logger.warning(f"Skipping malformed chunk: {e}")
+                    continue
+                    
         except Exception as e:
             logger.error(f"Error during Custom API call for /chat: {e}", exc_info=True)
             yield f"Error processing with Custom API: {str(e)}"
-    return StreamingResponse(custom_api_streamer(), media_type="text/plain")
-
-
-@app.post("/gemini/submit_task", response_model=TaskSubmissionResponse)
-async def submit_gemini_task(request: GeminiTaskRequest, background_tasks: BackgroundTasks):
-    task_id = str(uuid.uuid4())
-    logger.info(f"Received Gemini task submission. Assigning Task ID: {task_id}. Message: '{request.message[:50]}...'")
-
-    gemini_api_key_from_request = request.api_key
-    gemini_api_key_secret = os.getenv("GEMINI_API_KEY")
-    key_to_use = gemini_api_key_from_request or gemini_api_key_secret
-
-    if not key_to_use:
-        logger.error(f"[Task {task_id}] Gemini API Key missing for task submission.")
-        raise HTTPException(status_code=400, detail="Gemini API Key required.")
-
-    requested_model = request.gemini_model or DEFAULT_GEMINI_MODEL
+        finally:
+            # Ensure proper cleanup of the client
+            if client:
+                try:
+                    await client.close()
+                except Exception as cleanup_error:
+                    logger.warning(f"Error closing OpenAI client: {cleanup_error}")
     
-    current_time = datetime.now(timezone.utc)
-    tasks_db[task_id] = {
-        "status": "PENDING",
-        "result": None,
-        "error": None,
-        "submitted_at": current_time,
-        "last_updated_at": current_time,
-        "request_params": request.model_dump() # Store original request
-    }
-
-    background_tasks.add_task(
-        process_gemini_request_background,
-        task_id,
-        request.message,
-        request.url,
-        requested_model,
-        key_to_use
-    )
-    
-    logger.info(f"[Task {task_id}] Task submitted to background processing.")
-    return TaskSubmissionResponse(
-        task_id=task_id,
-        status="PENDING",
-        task_detail_url=f"/gemini/task/{task_id}" # Provide the URL to poll
+    return StreamingResponse(
+        custom_api_streamer(), 
+        media_type="text/plain",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        }
     )