Update private_gpt/server/completions/completions_router.py

private_gpt/server/completions/completions_router.py

@@ -36,50 +36,92 @@ class CompletionsBody(BaseModel):
     }
 
 
-@completions_router.post(
-    "/completions",
-    response_model=None,
-    summary="Completion",
-    responses={200: {"model": OpenAICompletion}},
-    tags=["Contextual Completions"],
-)
-def prompt_completion(
-    request: Request, body: CompletionsBody
-) -> OpenAICompletion | StreamingResponse:
-    """We recommend most users use our Chat completions API.
+# @completions_router.post(
+#     "/completions",
+#     response_model=None,
+#     summary="Completion",
+#     responses={200: {"model": OpenAICompletion}},
+#     tags=["Contextual Completions"],
+# )
+# def prompt_completion(
+#     request: Request, body: CompletionsBody
+# ) -> OpenAICompletion | StreamingResponse:
+#     """We recommend most users use our Chat completions API.
+
+#     Given a prompt, the model will return one predicted completion.
 
-    Given a prompt, the model will return one predicted completion.
+#     Optionally include a `system_prompt` to influence the way the LLM answers.
 
-    Optionally include a `system_prompt` 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.
 
-    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 `'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}]}
+#     ```
+#     """
+#     messages = [OpenAIMessage(content=body.prompt, role="user")]
+#     # If system prompt is passed, create a fake message with the system prompt.
+#     if body.system_prompt:
+#         messages.insert(0, OpenAIMessage(content=body.system_prompt, role="system"))
 
-    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}]}
-    ```
-    """
-    messages = [OpenAIMessage(content=body.prompt, role="user")]
-    # If system prompt is passed, create a fake message with the system prompt.
-    if body.system_prompt:
-        messages.insert(0, OpenAIMessage(content=body.system_prompt, role="system"))
+#     chat_body = ChatBody(
+#         messages=messages,
+#         use_context=body.use_context,
+#         stream=body.stream,
+#         include_sources=body.include_sources,
+#         context_filter=body.context_filter,
+#     )
+#     return chat_completion(request, chat_body)
 
-    chat_body = ChatBody(
-        messages=messages,
-        use_context=body.use_context,
-        stream=body.stream,
-        include_sources=body.include_sources,
-        context_filter=body.context_filter,
-    )
-    return chat_completion(request, chat_body)
+@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."""
+    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"}}