Fix missing completion types

posit-dev · Dec 3, 2024 · 78f7eb9 · 78f7eb9
1 parent ad96e6c
commit 78f7eb9
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/chatlas/_github.py b/chatlas/_github.py
@@ -9,6 +9,7 @@
 from ._utils import MISSING, MISSING_TYPE, inform_model_default
 
 if TYPE_CHECKING:
+    from ._openai import ChatCompletion
     from .types.openai import ChatClientArgs, SubmitInputArgs
 
 
@@ -21,7 +22,7 @@ def ChatGithub(
     base_url: str = "https://models.inference.ai.azure.com/",
     seed: Optional[int] | MISSING_TYPE = MISSING,
     kwargs: Optional["ChatClientArgs"] = None,
-) -> Chat["SubmitInputArgs"]:
+) -> Chat["SubmitInputArgs", ChatCompletion]:
     """
     Chat with a model hosted on the GitHub model marketplace.
 

diff --git a/chatlas/_groq.py b/chatlas/_groq.py
@@ -9,6 +9,7 @@
 from ._utils import MISSING, MISSING_TYPE, inform_model_default
 
 if TYPE_CHECKING:
+    from ._openai import ChatCompletion
     from .types.openai import ChatClientArgs, SubmitInputArgs
 
 
@@ -21,7 +22,7 @@ def ChatGroq(
     base_url: str = "https://api.groq.com/openai/v1",
     seed: Optional[int] | MISSING_TYPE = MISSING,
     kwargs: Optional["ChatClientArgs"] = None,
-) -> Chat["SubmitInputArgs"]:
+) -> Chat["SubmitInputArgs", ChatCompletion]:
     """
     Chat with a model hosted on Groq.
 

diff --git a/chatlas/_ollama.py b/chatlas/_ollama.py
@@ -10,6 +10,7 @@
 from ._turn import Turn
 
 if TYPE_CHECKING:
+    from ._openai import ChatCompletion
     from .types.openai import ChatClientArgs, SubmitInputArgs
 
 
@@ -21,7 +22,7 @@ def ChatOllama(
     base_url: str = "http://localhost:11434/v1",
     seed: Optional[int] = None,
     kwargs: Optional["ChatClientArgs"] = None,
-) -> Chat["SubmitInputArgs"]:
+) -> Chat["SubmitInputArgs", ChatCompletion]:
     """
     Chat with a local Ollama model.
 

diff --git a/chatlas/_perplexity.py b/chatlas/_perplexity.py
@@ -9,6 +9,7 @@
 from ._utils import MISSING, MISSING_TYPE, inform_model_default
 
 if TYPE_CHECKING:
+    from ._openai import ChatCompletion
     from .types.openai import ChatClientArgs, SubmitInputArgs
 
 
@@ -21,7 +22,7 @@ def ChatPerplexity(
     base_url: str = "https://api.perplexity.ai/",
     seed: Optional[int] | MISSING_TYPE = MISSING,
     kwargs: Optional["ChatClientArgs"] = None,
-) -> Chat["SubmitInputArgs"]:
+) -> Chat["SubmitInputArgs", ChatCompletion]:
     """
     Chat with a model hosted on perplexity.ai.