From 78f7eb93e5748575799b33d3e237b0b57d174109 Mon Sep 17 00:00:00 2001
From: Carson <cpsievert1@gmail.com>
Date: Tue, 3 Dec 2024 10:12:06 -0600
Subject: [PATCH] Fix missing completion types

---
 chatlas/_github.py     | 3 ++-
 chatlas/_groq.py       | 3 ++-
 chatlas/_ollama.py     | 3 ++-
 chatlas/_perplexity.py | 3 ++-
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/chatlas/_github.py b/chatlas/_github.py
index 2b3260a..c622065 100644
--- 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
index 4d0023a..dca414d 100644
--- 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
index 73a7911..e8a3d78 100644
--- 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
index 00bc14e..a59fc7b 100644
--- 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.