add typedef

HumanSignal · Nov 26, 2024 · 9b78da0 · 9b78da0
1 parent 7131627
commit 9b78da0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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: