diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
index b89bda5..f222dfd 100644
--- a/llama_cpp/llama.py
+++ b/llama_cpp/llama.py
@@ -480,7 +480,7 @@ class Llama:
         repeat_penalty: float = 1.1,
         top_k: int = 40,
         stream: bool = False,
-    ):
+    ) -> Union[Completion, Iterator[CompletionChunk]]:
         """Generate text from a prompt.
 
         Args: