Update private_gpt/server/chat/chat_router.py

private_gpt/server/chat/chat_router.py

@@ -49,6 +49,64 @@ class ChatBody(BaseModel):
     }
 
 
+# @chat_router.post(
+#     "/chat/completions",
+#     response_model=None,
+#     responses={200: {"model": OpenAICompletion}},
+#     tags=["Contextual Completions"],
+# )
+# def chat_completion(
+#     request: Request, body: ChatBody
+# ) -> OpenAICompletion | StreamingResponse:
+#     """Given a list of messages comprising a conversation, return a response.
+
+#     Optionally include an initial `role: system` message to influence the way
+#     the LLM answers.
+
+#     If `use_context` is set to `true`, the model will use context coming
+#     from the ingested documents to create the response. The documents being used can
+#     be filtered using the `context_filter` and passing the document IDs to be used.
+#     Ingested documents IDs can be found using `/ingest/list` endpoint. If you want
+#     all ingested documents to be used, remove `context_filter` altogether.
+
+#     When using `'include_sources': true`, the API will return the source Chunks used
+#     to create the response, which come from the context provided.
+
+#     When using `'stream': true`, the API will return data chunks following [OpenAI's
+#     streaming model](https://platform.openai.com/docs/api-reference/chat/streaming):
+#     ```
+#     {"id":"12345","object":"completion.chunk","created":1694268190,
+#     "model":"private-gpt","choices":[{"index":0,"delta":{"content":"Hello"},
+#     "finish_reason":null}]}
+#     ```
+#     """
+#     service = request.state.injector.get(ChatService)
+#     all_messages = [
+#         ChatMessage(content=m.content, role=MessageRole(m.role)) for m in body.messages
+#     ]
+#     if body.stream:
+#         completion_gen = service.stream_chat(
+#             messages=all_messages,
+#             use_context=body.use_context,
+#             context_filter=body.context_filter,
+#         )
+#         return StreamingResponse(
+#             to_openai_sse_stream(
+#                 completion_gen.response,
+#                 completion_gen.sources if body.include_sources else None,
+#             ),
+#             media_type="text/event-stream",
+#         )
+#     else:
+#         completion = service.chat(
+#             messages=all_messages,
+#             use_context=body.use_context,
+#             context_filter=body.context_filter,
+#         )
+#         return to_openai_response(
+#             completion.response, completion.sources if body.include_sources else None
+#         )
+
 @chat_router.post(
     "/chat/completions",
     response_model=None,
@@ -58,51 +116,36 @@ class ChatBody(BaseModel):
 def chat_completion(
     request: Request, body: ChatBody
 ) -> OpenAICompletion | StreamingResponse:
-    """Given a list of messages comprising a conversation, return a response.
-
-    Optionally include an initial `role: system` message to influence the way
-    the LLM answers.
-
-    If `use_context` is set to `true`, the model will use context coming
-    from the ingested documents to create the response. The documents being used can
-    be filtered using the `context_filter` and passing the document IDs to be used.
-    Ingested documents IDs can be found using `/ingest/list` endpoint. If you want
-    all ingested documents to be used, remove `context_filter` altogether.
-
-    When using `'include_sources': true`, the API will return the source Chunks used
-    to create the response, which come from the context provided.
+    """Given a list of messages comprising a conversation, return a response."""
+    try:
+        service = request.state.injector.get(ChatService)
+        all_messages = [
+            ChatMessage(content=m.content, role=MessageRole(m.role)) for m in body.messages
+        ]
+        if body.stream:
+            completion_gen = service.stream_chat(
+                messages=all_messages,
+                use_context=body.use_context,
+                context_filter=body.context_filter,
+            )
+            return StreamingResponse(
+                to_openai_sse_stream(
+                    completion_gen.response,
+                    completion_gen.sources if body.include_sources else None,
+                ),
+                media_type="text/event-stream",
+            )
+        else:
+            completion = service.chat(
+                messages=all_messages,
+                use_context=body.use_context,
+                context_filter=body.context_filter,
+            )
+            return to_openai_response(
+                completion.response, completion.sources if body.include_sources else None
+            )
+    except Exception as e:
+        # Log the exception details for debugging
+        print(f"Error processing chat completion: {e}")
+        return {"error": {"message": "Internal server error"}}
 
-    When using `'stream': true`, the API will return data chunks following [OpenAI's
-    streaming model](https://platform.openai.com/docs/api-reference/chat/streaming):
-    ```
-    {"id":"12345","object":"completion.chunk","created":1694268190,
-    "model":"private-gpt","choices":[{"index":0,"delta":{"content":"Hello"},
-    "finish_reason":null}]}
-    ```
-    """
-    service = request.state.injector.get(ChatService)
-    all_messages = [
-        ChatMessage(content=m.content, role=MessageRole(m.role)) for m in body.messages
-    ]
-    if body.stream:
-        completion_gen = service.stream_chat(
-            messages=all_messages,
-            use_context=body.use_context,
-            context_filter=body.context_filter,
-        )
-        return StreamingResponse(
-            to_openai_sse_stream(
-                completion_gen.response,
-                completion_gen.sources if body.include_sources else None,
-            ),
-            media_type="text/event-stream",
-        )
-    else:
-        completion = service.chat(
-            messages=all_messages,
-            use_context=body.use_context,
-            context_filter=body.context_filter,
-        )
-        return to_openai_response(
-            completion.response, completion.sources if body.include_sources else None
-        )