From 64b69a22799ca53d4a7f6096a9f6b6d8f2fa33fd Mon Sep 17 00:00:00 2001
From: Pedro Rico Pinazo <ricopinazo@gmail.com>
Date: Tue, 21 May 2024 07:19:38 +0100
Subject: [PATCH] fix type errors

---
 pyproject.toml          |  9 +++++++--
 src/comai/animations.py |  3 +--
 src/comai/chain.py      | 19 ++++++++++++++++---
 src/comai/cli.py        |  5 +++--
 src/comai/menu.py       | 32 --------------------------------
 src/comai/prompt.py     |  2 +-
 6 files changed, 28 insertions(+), 42 deletions(-)
 delete mode 100644 src/comai/menu.py

diff --git a/pyproject.toml b/pyproject.toml
index 632850a..945f1ac 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/src/comai/animations.py b/src/comai/animations.py
index 0b18047..f7773a5 100644
--- a/src/comai/animations.py
+++ b/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:
diff --git a/src/comai/chain.py b/src/comai/chain.py
index 1f64c4f..43b59bd 100644
--- a/src/comai/chain.py
+++ b/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:
diff --git a/src/comai/cli.py b/src/comai/cli.py
index 5ec73e2..160dcb6 100755
--- a/src/comai/cli.py
+++ b/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)
 
 
diff --git a/src/comai/menu.py b/src/comai/menu.py
deleted file mode 100644
index 52016b3..0000000
--- a/src/comai/menu.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import click
-from enum import Enum
-
-from rich import print
-from rich.markup import escape
-
-from comai.settings import Settings
-from comai.animations import show_cursor
-
-
-class MenuOption(str, Enum):
-    run = "r"
-    cancel = "c"
-
-
-DEFAULT_OPTION = escape("[r]")
-MENU_PROMPT = f"[bright_black] ➜ [underline bold]r[/underline bold]un | [underline bold]c[/underline bold]ancel {DEFAULT_OPTION}:[/bright_black]"
-
-
-def get_option_from_menu(settings: Settings) -> MenuOption:
-    if settings.verbose:
-        print(MENU_PROMPT, end="", flush=True)
-    show_cursor()
-    option = click.prompt(
-        "",
-        prompt_suffix="",
-        type=MenuOption,
-        default=MenuOption.run,
-        show_default=False,
-        show_choices=False,
-    )
-    return option
diff --git a/src/comai/prompt.py b/src/comai/prompt.py
index d1a051a..7d94068 100644
--- a/src/comai/prompt.py
+++ b/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: