Update main.py

main.py

@@ -1,15 +1,18 @@
-import fastapi
 import json
 import markdown
+from typing import List, Dict, Any, Generator
+from functools import partial
+
+import fastapi
 import uvicorn
-from fastapi import HTTPException
-from fastapi.responses import HTMLResponse
+from fastapi import HTTPException, Depends, Request
+from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from sse_starlette.sse import EventSourceResponse
-from starlette.responses import StreamingResponse
-from ctransformers import AutoModelForCausalLM
+from anyio import create_memory_object_stream, run_in_threadpool
+from transformers import AutoModelForCausalLM
 from pydantic import BaseModel
-from typing import List, Dict, Any, Generator
+
 
 
 llm = AutoModelForCausalLM.from_pretrained("TheBloke/WizardCoder-15B-1.0-GGML",
@@ -88,6 +91,37 @@ async def chat(request: ChatCompletionRequest):
 
     return StreamingResponse(format_response(chat_chunks), media_type="text/event-stream")
 
+@app.post("/v2/chat/completions")
+async def chatV2(request: Request, body: ChatCompletionRequest):
+    combined_messages = ' '.join([message.content for message in body.messages])
+    tokens = llm.tokenize(combined_messages)
+
+    send_chan, recv_chan = create_memory_object_stream(10)
+
+    async def event_publisher(inner_send_chan):
+        async with inner_send_chan:
+            try:
+                iterator: Generator = await run_in_threadpool(llm.generate, tokens)
+                for chat_chunk in iterator:
+                    response = {
+                        'choices': [
+                            {
+                                'message': {
+                                    'role': 'system',
+                                    'content': llm.detokenize(chat_chunk)
+                                },
+                                'finish_reason': 'stop' if llm.detokenize(chat_chunk) == "[DONE]" else 'unknown'
+                            }
+                        ]
+                    }
+                    await inner_send_chan.send(f"data: {json.dumps(response)}\n\n")
+                await inner_send_chan.send("event: done\ndata: {}\n\n")
+            except Exception as e:
+                print(f"Exception in event publisher: {str(e)}")
+
+    return StreamingResponse(recv_chan, media_type="text/event-stream", data_sender_callable=partial(event_publisher, send_chan))
+
+
 @app.post("/v0/chat/completions")
 async def chat(request: ChatCompletionRequestV0, response_mode=None):
     tokens = llm.tokenize(request.prompt)