Fix return type of _get_model_by_name

Flagro · Oct 14, 2024 · 19f4075 · 19f4075
1 parent 55d1c4a
commit 19f4075
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bot/rp_bot/ai_agent/ai.py b/bot/rp_bot/ai_agent/ai.py
@@ -1,6 +1,6 @@
 import tiktoken
 import io
-from typing import AsyncIterator, Literal, Optional
+from typing import AsyncIterator, Literal
 from openai import OpenAI
 from langchain_openai import OpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -57,7 +57,7 @@ def _get_default_model_name(
 
     def _get_model_by_name(
         self, model_name: str, model_type: Literal["text", "vision", "image_generation"]
-    ) -> Optional[Model]:
+    ) -> Model:
         models_dict = {
             "text": self.ai_config.TextGeneration.Models,
             "vision": self.ai_config.TextGeneration.Models,