diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
index 2df2b14..7c654cf 100644
--- a/adala/runtimes/_litellm.py
+++ b/adala/runtimes/_litellm.py
@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Type, Union, Literal, TypedDict, Iterable
+from typing import Any, Dict, List, Optional, Type, Union, Literal, TypedDict, Iterable, Generator
 from functools import cached_property
 from enum import Enum
 import litellm
@@ -651,7 +651,7 @@ def add_to_current_chunk(
             return chunk
 
     # Build chunks by iterating through parsed template parts
-    def build_chunks(parsed: Iterable[TemplateChunks]):
+    def build_chunks(parsed: Iterable[TemplateChunks]) -> Generator[MessageChunk, None, None]:
         current_chunk: Optional[MessageChunk] = None
 
         for part in parsed: