Update app.py

app.py

@@ -168,65 +168,58 @@ async def direct_chat(payload: ChatPayload):
         raise HTTPException(status_code=500, detail="Custom API key not configured.")
     
     async def custom_api_streamer():
-        executor = ThreadPoolExecutor(max_workers=1)
         client = None
-        
         try:
             logger.info("Sending request to Custom API for /chat.")
             
-            def create_stream():
-                nonlocal client
-                client = openai.OpenAI(
-                    api_key=custom_api_key_secret, 
-                    base_url=custom_api_base_url,
-                    timeout=30.0
-                )
-                return client.chat.completions.create(
-                    model=custom_api_model,
-                    temperature=payload.temperature,
-                    messages=[{"role": "user", "content": payload.message}],
-                    stream=True
-                )
+            # Use AsyncOpenAI with proper configuration
+            from openai import AsyncOpenAI
+            client = AsyncOpenAI(
+                api_key=custom_api_key_secret,
+                base_url=custom_api_base_url,
+                timeout=60.0  # Longer timeout for gemma3:27b model
+            )
             
-            # Run the blocking call in a thread
-            stream = await asyncio.get_event_loop().run_in_executor(executor, create_stream)
+            stream = await client.chat.completions.create(
+                model=custom_api_model,
+                temperature=payload.temperature,
+                messages=[{"role": "user", "content": payload.message}],
+                stream=True
+            )
             
-            def process_chunk():
+            async for chunk in stream:
                 try:
-                    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
+                    # Exact same logic as your working code
+                    if hasattr(chunk.choices[0].delta, "reasoning_content") and chunk.choices[0].delta.reasoning_content:
+                        yield chunk.choices[0].delta.reasoning_content
+                    elif chunk.choices[0].delta.content is not None:  # Handle None explicitly
+                        yield chunk.choices[0].delta.content
                         
-                        if content_to_yield:
-                            yield content_to_yield
+                except (IndexError, AttributeError) as e:
+                    # Skip malformed chunks silently (some APIs send empty chunks)
+                    continue
                 except Exception as e:
-                    logger.error(f"Error processing stream chunk: {e}")
-                    yield f"Error processing stream: {str(e)}"
-            
-            # Process chunks in the thread pool
-            loop = asyncio.get_event_loop()
-            chunks = await loop.run_in_executor(executor, lambda: list(process_chunk()))
-            
-            for chunk in chunks:
-                yield chunk
-                
+                    logger.warning(f"Skipping chunk due to error: {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)}"
+            
+            # Handle specific connection errors with retry suggestion
+            if "peer closed connection" in str(e) or "incomplete chunked read" in str(e):
+                yield "Connection interrupted. Please try again."
+            else:
+                yield f"Error processing with Custom API: {str(e)}"
+                
         finally:
-            # Cleanup
             if client:
                 try:
-                    client.close()
+                    await client.close()
                 except Exception as cleanup_error:
                     logger.warning(f"Error closing OpenAI client: {cleanup_error}")
-            executor.shutdown(wait=False)
     
     return StreamingResponse(
-        custom_api_streamer(), 
+        custom_api_streamer(),
         media_type="text/plain",
         headers={
             "Cache-Control": "no-cache",