1 files changed, 9 insertions, 10 deletions
diff --git a/g4f/local/__init__.py b/g4f/local/__init__.py
index c9d3d74a..d13c8c1d 100644
--- a/g4f/local/__init__.py
+++ b/g4f/local/__init__.py
@@ -1,17 +1,17 @@
-from ..typing import Union, Iterator, Messages
-from ..stubs  import ChatCompletion, ChatCompletionChunk
-from ._engine import LocalProvider
-from ._models import models
-from ..client import iter_response, filter_none, IterResponse
+from ..typing import Union, Messages
+from ..locals.provider import LocalProvider
+from ..locals.models import get_models
+from ..client.client import iter_response, filter_none
+from ..client.types import IterResponse
 
 class LocalClient():
     def __init__(self, **kwargs) -> None:
         self.chat: Chat = Chat(self)
-    
+
     @staticmethod
     def list_models():
-        return list(models.keys())
-        
+        return list(get_models())
+
 class Completions():
     def __init__(self, client: LocalClient):
         self.client: LocalClient = client
@@ -25,8 +25,7 @@ class Completions():
         max_tokens: int = None,
         stop: Union[list[str], str] = None,
         **kwargs
-    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
-
+    ) -> IterResponse:
         stop = [stop] if isinstance(stop, str) else stop
         response = LocalProvider.create_completion(
             model, messages, stream,