Update app.py

app.py

@@ -4,6 +4,8 @@ import logging
 import uuid
 from datetime import datetime, timezone
 from typing import Optional, Dict, Any
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 
 from fastapi import FastAPI, HTTPException, Body, BackgroundTasks, Path
 from fastapi.responses import StreamingResponse
@@ -166,50 +168,62 @@ 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.")
             
-            # 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
-            )
+            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
+                )
             
-            stream = await client.chat.completions.create(
-                model=custom_api_model,
-                temperature=payload.temperature,
-                messages=[{"role": "user", "content": payload.message}],
-                stream=True
-            )
+            # Run the blocking call in a thread
+            stream = await asyncio.get_event_loop().run_in_executor(executor, create_stream)
             
-            async for chunk in stream:
+            def process_chunk():
                 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
+                    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
                         
-                except (IndexError, AttributeError) as e:
-                    logger.warning(f"Skipping malformed chunk: {e}")
-                    continue
-                    
+                        if content_to_yield:
+                            yield content_to_yield
+                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
+                
         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)}"
         finally:
-            # Ensure proper cleanup of the client
+            # Cleanup
             if client:
                 try:
-                    await client.close()
+                    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(),