Update app.py

app.py

@@ -54,7 +54,7 @@ def get_api_key(call: str = 'api_key') -> str:
 
 def encode_chat(messages: List[dict]) -> str:
     encoded = "\n".join(
-        f"<|im_start|>{msg['role']}{f' [{msg['name']}]' if 'name' in msg else ''}\n{msg['content']}<|end_of_text|>"
+        f"<|im_start|>{msg['role']}{' [' + msg['name'] + ']' if 'name' in msg else ''}\n{msg['content']}<|end_of_text|>"
         for msg in messages
     )
     logger.debug(f"Encoded chat: {encoded}")
@@ -170,26 +170,32 @@ async def stream_chat(params: dict):
                 async with session.post(
                     url,
                     headers={
-                        "Authorization": f"Bearer {get_api_key('api_key')}",
+                        "Authorization": f"Bearer {get_api_key('api_key' if attempt == 1 else 'oai_api_key')}",
                         "Content-Type": "application/json"
                     },
-                    json=params
+                    json=params,
+                    timeout=30
                 ) as resp:
                     resp.raise_for_status()
-                    async for line in resp.content:
-                        if line:
-                            line_str = line.decode('utf-8').strip()
-                            if line_str.startswith("data: "):
-                                line_str = line_str[6:]
-                            if line_str == "[DONE]":
-                                break
-                            try:
-                                message = json.loads(line_str)
-                                yield message
-                            except json.JSONDecodeError:
-                                logger.warning("Failed to decode JSON from line.")
-                                continue
-                break  # Successful request, exit the loop
+                    buffer = ""
+                    async for chunk in resp.content:
+                        if chunk:
+                            buffer += chunk.decode('utf-8')
+                            while '\n' in buffer:
+                                line, buffer = buffer.split('\n', 1)
+                                line = line.strip()
+                                if line.startswith("data: "):
+                                    line = line[6:].strip()
+                                if line == "[DONE]":
+                                    return
+                                if not line:
+                                    continue
+                                try:
+                                    message = json.loads(line)
+                                    yield message
+                                except json.JSONDecodeError:
+                                    continue
+                break
             except aiohttp.ClientError as e:
                 logger.error(f"Stream chat request failed on attempt {attempt}: {e}")
                 if attempt == 2:
@@ -242,13 +248,21 @@ async def respond(
         "stream": True
     }
 
-    response_text = ""
-    async for token in stream_chat(params):
-        if token and 'choices' in token:
-            delta = token['choices'][0].get('delta', {})
-            content = delta.get("content", delta.get("refusal", ""))
-            response_text += content
-            yield response_text
+    try:
+        response_text = ""
+        async for token in stream_chat(params):
+            if token and 'choices' in token and len(token['choices']) > 0:
+                delta = token['choices'][0].get('delta', {})
+                content = delta.get("content", delta.get("refusal", ""))
+                response_text += content
+                yield response_text
+        
+        if not response_text:
+            yield "I apologize, but I was unable to generate a response. Please try again."
+            
+    except Exception as e:
+        logger.error(f"Error during chat response generation: {e}")
+        yield "I encountered an error while processing your request. Please try again later."
 
 def create_gradio_interface() -> gr.ChatInterface:
     return gr.ChatInterface(
@@ -309,7 +323,7 @@ def main():
     args = ArgParser().args
 
     uvicorn.run(
-        "main:app",
+        app,
         host=args.server,
         port=args.port,
         reload=args.dev