Simply reuse parent's init function instead of overriding private att…

…ributes
rogeriochaves · Jul 25, 2023 · f95a0aa · f95a0aa
1 parent aff3242
commit f95a0aa
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/litechain/contrib/llms/gpt4all_chain.py b/litechain/contrib/llms/gpt4all_chain.py
@@ -1,5 +1,5 @@
 import asyncio
-from typing import AsyncGenerator, Callable, Iterable, Optional, TypeVar
+from typing import AsyncGenerator, Callable, Iterable, Optional, TypeVar, cast
 
 from gpt4all import GPT4All
 
@@ -65,10 +65,9 @@ def __init__(
         n_batch=8,
         n_threads: Optional[int] = None,
     ) -> None:
-        self.name = name
         gpt4all = GPT4All(model, n_threads=n_threads)
 
-        async def generate(prompt: str) -> AsyncGenerator[str, None]:
+        async def generate(prompt: str) -> AsyncGenerator[U, None]:
             loop = asyncio.get_event_loop()
 
             def get_outputs() -> Iterable[str]:
@@ -87,6 +86,6 @@ def get_outputs() -> Iterable[str]:
             outputs = await loop.run_in_executor(None, get_outputs)
 
             for output in outputs:
-                yield output
+                yield cast(U, output)
 
-        self._call = lambda input: generate(call(input))
+        super().__init__(name, lambda input: generate(call(input)))
diff --git a/litechain/contrib/llms/open_ai.py b/litechain/contrib/llms/open_ai.py
@@ -65,9 +65,7 @@ def __init__(
         temperature: Optional[float] = 0,
         max_tokens: Optional[int] = None,
     ) -> None:
-        self.name = name
-
-        async def completion(prompt: str) -> AsyncGenerator[str, None]:
+        async def completion(prompt: str) -> AsyncGenerator[U, None]:
             loop = asyncio.get_event_loop()
 
             def get_completions():
@@ -88,7 +86,7 @@ def get_completions():
                         if "text" in output["choices"][0]:
                             yield output["choices"][0]["text"]
 
-        self._call = lambda input: completion(call(input))
+        super().__init__(name, lambda input: completion(call(input)))
 
 
 @dataclass
@@ -263,7 +261,7 @@ class OpenAIChatChain(Chain[T, U]):
     """
 
     def __init__(
-        self: "OpenAIChatChain[T, Union[OpenAIChatDelta, V]]",
+        self: "OpenAIChatChain[T, OpenAIChatDelta]",
         name: str,
         call: Callable[
             [T],
@@ -275,11 +273,9 @@ def __init__(
         temperature: Optional[float] = 0,
         max_tokens: Optional[int] = None,
     ) -> None:
-        self.name = name
-
         async def chat_completion(
             messages: List[OpenAIChatMessage],
-        ) -> AsyncGenerator[ChainOutput[Union[OpenAIChatDelta, V]], Any]:
+        ) -> AsyncGenerator[ChainOutput[OpenAIChatDelta], None]:
             loop = asyncio.get_event_loop()
 
             def get_completions():
@@ -349,4 +345,7 @@ def get_completions():
                 yield self._output_wrap(pending_function_call)
                 pending_function_call = None
 
-        self._call = lambda input: chat_completion(call(input))
+        super().__init__(
+            name,
+            lambda input: cast(AsyncGenerator[U, None], chat_completion(call(input))),
+        )