Update main.py

main.py

@@ -1,8 +1,10 @@
+import json
 from typing import List
 import fastapi
 import markdown
 import uvicorn
 from ctransformers import AutoModelForCausalLM
+from fastapi import HTTPException
 from fastapi.responses import HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
@@ -57,6 +59,38 @@ class ChatCompletionRequest(BaseModel):
 system_message = "Below is a conversation between a human user and a helpful AI coding assistant."
 
 @app.post("/v1/chat/completions")
+async def chat(request: ChatCompletionRequest):
+    kwargs = request.dict()
+    dialogue_template = DialogueTemplate(
+        system=system_message, messages=kwargs['messages']
+    )
+    prompt = dialogue_template.get_inference_prompt()
+    tokens = llm.tokenize(combined_messages)
+    
+    try:
+        chat_chunks = llm.generate(tokens)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    async def format_response(chat_chunks: Generator) -> Any:
+        for chat_chunk in chat_chunks:
+            response = {
+                'choices': [
+                    {
+                        'message': {
+                            'role': 'system',
+                            'content': llm.detokenize(chat_chunk)
+                        },
+                        'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
+                    }
+                ]
+            }
+            yield f"data: {json.dumps(response)}\n\n"
+        yield "event: done\ndata: {}\n\n"
+
+    return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
+
+@app.post("/v0/chat/completions")
 async def chat(request: ChatCompletionRequest, response_mode=None):
     kwargs = request.dict()
     dialogue_template = DialogueTemplate(