Match openai completions api

main.py

@@ -2,11 +2,14 @@ import fastapi
 import json
 import markdown
 import uvicorn
-from fastapi.responses import HTMLResponse
+from fastapi import HTTPException
+from fastapi.responses import HTMLResponse, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
 from ctransformers import AutoModelForCausalLM
 from pydantic import BaseModel
+from typing import List, Dict, Any
+
 
 llm = AutoModelForCausalLM.from_pretrained("TheBloke/WizardCoder-15B-1.0-GGML",
                                            model_file="WizardCoder-15B-1.0.ggmlv3.q4_0.bin",
@@ -43,11 +46,44 @@ async def index():
 class ChatCompletionRequest(BaseModel):
     prompt: str
 
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequestV2(BaseModel):
+    messages: List[Message]
+    max_tokens: int = 100
+
 @app.post("/v1/completions")
 async def completion(request: ChatCompletionRequest, response_mode=None):
     response = llm(request.prompt)
     return response
 
+@app.post("/v2/chat/completions")
+async def chat(request: ChatCompletionRequestV2):
+    tokens = llm.tokenize([message.content for message in request.messages])
+    
+    try:
+        chat_chunks = llm.generate(tokens, max_tokens=request.max_tokens)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    def format_response(chat_chunks) -> Dict[str, Any]:
+        response = {
+            'choices': []
+        }
+        for chat_chunk in chat_chunks:
+            response['choices'].append({
+                'message': {
+                    'role': 'system',
+                    'content': llm.detokenize(chat_chunk)
+                },
+                'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
+            })
+        return response
+
+    return format_response(chat_chunks)
+
 @app.post("/v1/chat/completions")
 async def chat(request: ChatCompletionRequest, response_mode=None):
     tokens = llm.tokenize(request.prompt)