Handle prompt list

abetlen · abetlen · commit 55279b679df4 · 2023-04-06T21:07:35.000-04:00
diff --git a/llama_cpp/server/__main__.py b/llama_cpp/server/__main__.py
@@ -60,7 +60,7 @@ class Settings(BaseSettings):
 
 
 class CreateCompletionRequest(BaseModel):
-    prompt: str
+    prompt: Union[str, List[str]]
     suffix: Optional[str] = Field(None)
     max_tokens: int = 16
     temperature: float = 0.8
@@ -100,10 +100,10 @@ class Config:
     response_model=CreateCompletionResponse,
 )
 def create_completion(request: CreateCompletionRequest):
-    if request.stream:
-        chunks: Iterator[llama_cpp.CompletionChunk] = llama(**request.dict())  # type: ignore
-        return EventSourceResponse(dict(data=json.dumps(chunk)) for chunk in chunks)
-    return llama(
+    if isinstance(request.prompt, list):
+        request.prompt = "".join(request.prompt)
+
+    completion_or_chunks = llama(
         **request.dict(
             exclude={
                 "model",
@@ -117,6 +117,11 @@ def create_completion(request: CreateCompletionRequest):
             }
         )
     )
+    if request.stream:
+        chunks: Iterator[llama_cpp.CompletionChunk] = completion_or_chunks  # type: ignore
+        return EventSourceResponse(dict(data=json.dumps(chunk)) for chunk in chunks)
+    completion: llama_cpp.Completion = completion_or_chunks  # type: ignore
+    return completion
 
 
 class CreateEmbeddingRequest(BaseModel):
@@ -259,4 +264,6 @@ def get_models() -> ModelList:
     import os
     import uvicorn
 
-    uvicorn.run(app, host=os.getenv("HOST", "localhost"), port=int(os.getenv("PORT", 8000)))
+    uvicorn.run(
+        app, host=os.getenv("HOST", "localhost"), port=int(os.getenv("PORT", 8000))
+    )
