Update apis/chat_api.py

apis/chat_api.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import sys
 import uvicorn
+import traceback
 
 from fastapi import FastAPI, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -76,28 +77,33 @@ class ChatAPIApp:
     def chat_completions(
         self, item: ChatCompletionsPostItem, api_key: str = Depends(extract_api_key)
     ):
-        streamer = MessageStreamer(model=item.model)
-        composer = MessageComposer(model=item.model)
-        composer.merge(messages=item.messages)
-        # streamer.chat = stream_chat_mock
-
-        stream_response = streamer.chat_response(
-            prompt=composer.merged_str,
-            temperature=item.temperature,
-            max_new_tokens=item.max_tokens,
-            api_key=api_key,
-        )
-        if item.stream:
-            event_source_response = EventSourceResponse(
-                streamer.chat_return_generator(stream_response),
-                media_type="text/event-stream",
-                ping=2000,
-                ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
+        try:
+            streamer = MessageStreamer(model=item.model)
+            composer = MessageComposer(model=item.model)
+            composer.merge(messages=item.messages)
+
+            stream_response = streamer.chat_response(
+                prompt=composer.merged_str,
+                temperature=item.temperature,
+                max_new_tokens=item.max_tokens,
+                api_key=api_key,
             )
-            return event_source_response
-        else:
-            data_response = streamer.chat_return_dict(stream_response)
-            return data_response
+
+            if item.stream:
+                event_source_response = EventSourceResponse(
+                    streamer.chat_return_generator(stream_response),
+                    media_type="text/event-stream",
+                    ping=2000,
+                    ping_message_factory=lambda: ServerSentEvent(**{"comment": ""}),
+                )
+                return event_source_response
+            else:
+                data_response = streamer.chat_return_dict(stream_response)
+                return data_response
+        except Exception as e:
+            logger.error(f"Error in chat_completions: {str(e)}")
+            logger.error(traceback.format_exc())
+            return {"error": "Internal server error"}
 
     def setup_routes(self):
         for prefix in ["", "/v1"]: