Update private_gpt/server/completions/completions_router.py

private_gpt/server/completions/completions_router.py

@@ -36,92 +36,51 @@ 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"}}