improve /completions api compatibility

josStorer · Jul 10, 2023 · f9f1d5c · f9f1d5c
1 parent 0edec68
commit f9f1d5c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/backend-python/routes/completion.py b/backend-python/routes/completion.py
@@ -44,7 +44,7 @@ class Config:
 
 
 class CompletionBody(ModelConfigBody):
- prompt: str
+ prompt: str or List[str]
  model: str = "rwkv"
  stream: bool = False
  stop: str = None
@@ -306,9 +306,12 @@ async def completions(body: CompletionBody, request: Request):
  if model is None:
  raise HTTPException(status.HTTP_400_BAD_REQUEST, "model not loaded")
 
- if body.prompt is None or body.prompt == "":
+ if body.prompt is None or body.prompt == "" or body.prompt == []:
  raise HTTPException(status.HTTP_400_BAD_REQUEST, "prompt not found")
 
+ if type(body.prompt) == list:
+ body.prompt = body.prompt[0] # TODO: support multiple prompts
+
  if body.stream:
  return EventSourceResponse(
  eval_rwkv(model, request, body, body.prompt, body.stream, body.stop, False)