fix type errors

pyproject.toml

@@ -23,7 +23,6 @@ dependencies = [
     "langchain==0.1.17",
     "langchain-openai==0.1.6",
     "ollama==0.1.9",
-    "types-requests==2.31.0.20240406",
 ]
 
 [project.urls]
@@ -35,7 +34,13 @@ issues = "https://github.com/ricopinazo/comai/issues"
 comai = "comai.cli:app"
 
 [project.optional-dependencies]
-test = ["pytest", "hatchling"]
+test = [
+    "pytest",
+    "hatchling",
+    "types-requests==2.31.0.20240406",
+    "mypy==1.9.0",
+    "black==24.4.0",
+]
 
 [tool.hatch.version]
 path = "src/comai/__init__.py"

src/comai/animations.py

@@ -68,8 +68,7 @@ def print_command_prompt(command: str):
     message = [
         ("class:mark", ANSWER_PROMPT),
     ]
-    # return "ls"
-    return prompt(message, default="%s" % command, style=style)
+    return prompt(message, default="%s" % command, style=style)  # type: ignore
 
 
 def hide_cursor() -> None:

src/comai/chain.py

@@ -6,6 +6,7 @@
 import uuid
 from dataclasses import dataclass
 from typing import Iterable, Iterator, List
+from pydantic import BaseModel, SecretStr
 
 from langchain.globals import set_debug, set_verbose
 from langchain_community.chat_message_histories import SQLChatMessageHistory
@@ -87,15 +88,27 @@ def attatch_history(
     )
 
 
+class OpenaiSecrets(BaseModel):
+    default_key: SecretStr | None
+    comai_key: SecretStr | None
+
+
+def extract_secret(key: str | None):
+    if key is None:
+        return None
+    else:
+        return SecretStr(key)
+
+
 def create_chain_stream(settings: Settings, context: Context):
     prompt = create_prompt(context)
     if settings.provider == "ollama":
         model = ChatOllama(model=settings.model, temperature=0)
     elif settings.provider == "openai":
-        default_key = os.environ.get("OPENAI_API_KEY")
-        comai_key = os.environ.get("COMAI_OPENAI_API_KEY")
+        default_key = extract_secret(os.environ.get("OPENAI_API_KEY"))
+        comai_key = extract_secret(os.environ.get("COMAI_OPENAI_API_KEY"))
         api_key = comai_key if comai_key is not None else default_key
-        model = ChatOpenAI(model=settings.model, temperature=0, api_key=api_key)
+        model = ChatOpenAI(model=settings.model, temperature=0, api_key=api_key)  # type: ignore
     base_chain = prompt | model
 
     if context.session_id is not None:

src/comai/cli.py

@@ -21,7 +21,6 @@
     Settings,
 )
 from comai.context import get_context
-from comai.menu import get_option_from_menu, MenuOption
 from comai.animations import (
     print_command_token,
     query_animation,
@@ -31,7 +30,9 @@
     print_command_prompt,
 )
 
-typer.core.rich = None  # this is to avoid using right for the help panel
+# this is to avoid using right for the help panel
+typer.core.rich = None  # type: ignore
+
 app = typer.Typer(pretty_exceptions_enable=False, add_completion=False)
 
 

src/comai/prompt.py

@@ -18,7 +18,7 @@ def prompt_str(question: str, default: str) -> str:
         ("class:mark", "? "),
         ("class:question", f"{question.strip()} "),
     ]
-    return prompt(message, default="%s" % default, style=style)
+    return prompt(message, default="%s" % default, style=style)  # type: ignore
 
 
 def prompt_options(question: str, options: list[str], default: str) -> str: