From 7d8c9c855391a21ca4fb3091a6b0728f1bbc746b Mon Sep 17 00:00:00 2001
From: piEsposito <piero.skywalker@gmail.com>
Date: Fri, 19 Jul 2024 16:02:41 -0300
Subject: [PATCH] allow setting assistant model

---
 pi_ai_os.py    | 33 ++++++++++++++++++++++++++++-----
 pyproject.toml |  2 +-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/pi_ai_os.py b/pi_ai_os.py
index 01cf819..158a097 100644
--- a/pi_ai_os.py
+++ b/pi_ai_os.py
@@ -124,6 +124,7 @@ def divide(a: int, b: int) -> int:
 the FILES tag must contain the name of the relevant files, one per line, without anything else.
 You can only return files that appeared in the context tags, even if they mention other files or directories that were not written there.
 If a file does not exist (and the user asks for something that will create it), do not include it in the list of dependencies.
+If there are no relevant files, just return the tags without any content.
 """
 
 
@@ -158,10 +159,12 @@ def extract_dependencies(raw_text):
 
 
 def select_relevant_files_on_chdir(
-    user_prompt: str, target_dirs: list[str] | None = None
+    user_prompt: str,
+    target_dirs: list[str] | None = None,
+    assistant_model: str = ASSISTANT_MODEL_NAME,
 ):
     assistant = AI(
-        ASSISTANT_MODEL_NAME,
+        assistant_model,
         system=ASSISTANT_MODEL_SYSTEM_PROMPT,
         max_new_tokens=1024,
         temperature=0,
@@ -194,8 +197,16 @@ def extract_code(text: str) -> str:
     return match.group(1).strip() if match else None
 
 
-def pi_ai_os(model: str, initial_message: str, config_file: str, assistant: bool):
+def pi_ai_os(
+    model: str,
+    initial_message: str,
+    config_file: str,
+    assistant: bool,
+    assistant_model: str,
+):
     print(f"Welcome to AIOS. Model: {model}\n")
+    if assistant:
+        print(f"Assistant model: {assistant_model}")
 
     config_file = config_file or DEFAULT_CONFIG_FILE
     if not os.path.exists(config_file):
@@ -251,7 +262,7 @@ def pi_ai_os(model: str, initial_message: str, config_file: str, assistant: bool
 
         if assistant:
             dependencies_context = select_relevant_files_on_chdir(
-                user_message, target_dirs
+                user_message, target_dirs, assistant_model
             )
             user_message = f"{user_message}\n\nYou are also provided following context: {dependencies_context}"
 
@@ -309,6 +320,12 @@ def main():
         help="Triggers using an assistant model for context gathering",
         action="store_true",
     )
+    parser.add_argument(
+        "-am",
+        "--assistant-model",
+        help="Specify the assistant model name",
+        default=ASSISTANT_MODEL_NAME,
+    )
     parser.add_argument(
         "-v",
         "--verbose",
@@ -323,7 +340,13 @@ def main():
 
     initial_message = " ".join(args.message) if args.message else None
     try:
-        pi_ai_os(args.model, initial_message, args.config, args.assistant)
+        pi_ai_os(
+            args.model,
+            initial_message,
+            args.config,
+            args.assistant,
+            args.assistant_model,
+        )
     except KeyboardInterrupt:
         print("\nExiting...")
         sys.exit(0)
diff --git a/pyproject.toml b/pyproject.toml
index d1f302d..0bb8d2c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pi-ai-os"
-version = "0.0.3"
+version = "0.0.4"
 description = "Integrate your OS with an LLM."
 authors = ["piEsposito <piero.skywalker@gmail.com>"]
 license = "Apache 2.0"