matthoffner commited on
Commit
a7653ed
1 Parent(s): 7d42f36

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +10 -1
main.py CHANGED
@@ -9,6 +9,7 @@ from sse_starlette.sse import EventSourceResponse
9
  from ctransformers.langchain import CTransformers
10
  from pydantic import BaseModel
11
  from typing import List, Any
 
12
 
13
  llm = AutoModelForCausalLM.from_pretrained("NeoDim/starchat-alpha-GGML",
14
  model_file="starchat-alpha-ggml-q4_0.bin",
@@ -30,8 +31,16 @@ async def index():
30
  html_content = markdown.markdown(md_template_string)
31
  return HTMLResponse(content=html_content, status_code=200)
32
 
 
 
 
 
 
 
33
  class ChatCompletionRequest(BaseModel):
34
- messages: List[Any]
 
 
35
 
36
  @app.post("/v1/chat/completions")
37
  async def chat(request: ChatCompletionRequest, response_mode=None):
 
9
  from ctransformers.langchain import CTransformers
10
  from pydantic import BaseModel
11
  from typing import List, Any
12
+ from typing_extensions import TypedDict, Literal
13
 
14
  llm = AutoModelForCausalLM.from_pretrained("NeoDim/starchat-alpha-GGML",
15
  model_file="starchat-alpha-ggml-q4_0.bin",
 
31
  html_content = markdown.markdown(md_template_string)
32
  return HTMLResponse(content=html_content, status_code=200)
33
 
34
+ class ChatCompletionRequestMessage(BaseModel):
35
+ role: Literal["system", "user", "assistant"] = Field(
36
+ default="user", description="The role of the message."
37
+ )
38
+ content: str = Field(default="", description="The content of the message.")
39
+
40
  class ChatCompletionRequest(BaseModel):
41
+ messages: List[ChatCompletionRequestMessage] = Field(
42
+ default=[], description="A list of messages to generate completions for."
43
+ )
44
 
45
  @app.post("/v1/chat/completions")
46
  async def chat(request: ChatCompletionRequest, response_mode=None):