Added support for local LLM models via Ollama

.env.example

@@ -0,0 +1,3 @@
+SIRJI_OPENAI_API_KEY=sk-none
+SIRJI_USE_OLLAMA=true
+SIRJI_OLLAMA_MODEL=gemma:2b

.gitignore

@@ -11,5 +11,6 @@ __pycache__/
 
 # Virtual environment
 venv
+.venv
 
 .DS_Store

README.md

@@ -16,7 +16,7 @@ Follow these [setup steps](./docs/setup-steps.md) to start Sirji.
 
 ## Capabilities<a name="capabilities"></a>
 
-Sirji is being built using Python. It currently uses OpenAI chat completions API and OpenAI assistants API.
+Sirji is being built using Python. It currently uses OpenAI chat completions API and OpenAI assistants API. Sirji works with either an OpenAI API key to use GPT 4, or a local model via Ollama.
 
 To begin with, it will be equipped with:
 
@@ -46,7 +46,7 @@ The sequence diagram below shows how the user initiates Sirji by giving a proble
 
 ## Supported AI Models<a name="supported-ai-models"></a>
 
-We are planning to use the `gpt-4-turbo-preview` model for the dogfood release. But the package will be designed to be composable for supporting other AI models too.
+We are planning to use the `gpt-4-turbo-preview` model for the dogfood release. ~~But the package will be designed to be composable for supporting other AI models too.~~ It should now work with local models via Ollama.
 
 ## Contribution
 

main.py

@@ -1,35 +1,34 @@
-import sys
-import uuid
 import os
+import queue
 import shutil
 import textwrap
-import re
 import threading
-import queue
 import time
 
-from sirji.view.terminal import open_terminal_and_run_command
-from sirji.view.screen import get_screen_resolution
-from sirji.view.chat import send_external_system_message, run_chat_app, disable_chat_send_button, enable_chat_send_button
-
-from sirji.tools.logger import coder as cLogger
-from sirji.tools.logger import researcher as rLogger
-from sirji.tools.logger import planner as pLogger
-from sirji.tools.logger import executor as eLogger
-from sirji.tools.logger import sirji as sLogger
-from sirji.tools.logger import user as uLogger
+from dotenv import load_dotenv
 
 from sirji.agents.coder import Coder
+from sirji.agents.executor import Executor
 from sirji.agents.planner import Planner
 from sirji.agents.researcher import Researcher
-from sirji.agents.executor import Executor
 from sirji.agents.user import User
-
 from sirji.messages.parser import MessageParser
+from sirji.tools.logger import coder as cLogger
+from sirji.tools.logger import executor as eLogger
+from sirji.tools.logger import planner as pLogger
+from sirji.tools.logger import researcher as rLogger
+from sirji.tools.logger import sirji as sLogger
+from sirji.tools.logger import user as uLogger
+from sirji.view.chat import (
+    disable_chat_send_button,
+    enable_chat_send_button,
+    run_chat_app,
+    send_external_system_message,
+)
+from sirji.view.screen import get_screen_resolution
+from sirji.view.terminal import open_terminal_and_run_command
 
-
-from sirji.messages.problem_statement import ProblemStatementMessage
-last_recipient = ''
+last_recipient = ""
 
 # Global queue to safely exchange messages between threads
 messages_queue = queue.Queue()
@@ -40,7 +39,7 @@ def empty_workspace():
 
     # List all files and directories in the workspace directory
     for item in os.listdir(workspace_dir):
-        if (item == "logs"):
+        if item == "logs":
             continue  # Skip the code directory
 
         item_path = os.path.join(workspace_dir, item)
@@ -56,14 +55,14 @@ def empty_workspace():
             print(f"Error removing {item_path}: {e}")
 
 
-class Main():
+class Main:
     def __init__(self):
         self.problem_statement = None  # Placeholder
         self.initialize_logs()
 
         self.coder = Coder()
         self.planner = Planner()
-        self.researcher = Researcher('openai_assistant', 'openai_assistant')
+        self.researcher = Researcher("openai_assistant", "openai_assistant")
         self.executor = Executor()
         self.user = User()
 
@@ -78,40 +77,49 @@ def truncate_logs(self):
                     # Check if the file exists at the given location
                     if os.path.exists(logger.filepath):
                         # File exists, so open it in 'w' mode to clear or create it, then immediately close it
-                        open(logger.filepath, 'w').close()
+                        open(logger.filepath, "w").close()
                     else:
                         # If the file does not exist, handle accordingly (e.g., log this situation or take other actions)
                         print(f"File does not exist at {logger.filepath}.")
 
                 except AttributeError:
                     # Handle the case where the logger doesn't have a 'filepath' attribute
-                    print(
-                        f"Logger {logger} does not have a 'filepath' attribute.")
+                    print(f"Logger {logger} does not have a 'filepath' attribute.")
 
                 except IOError as e:
                     # Handle other I/O errors, such as permission errors
                     print(
-                        f"Failed to open file for logger {logger} at {logger.filepath}: {e}")
+                        f"Failed to open file for logger {logger} at {logger.filepath}: {e}"
+                    )
 
     def initialize_logs(self):
         # Truncate the logs
         self.truncate_logs()
 
         # Initialize the logs
         cLogger.initialize_logs(
-            "Coder: Specializing in generating and modifying code, this agent is skilled in various programming languages and is equipped to handle tasks ranging from quick fixes to developing complex algorithms.\n\n\n")
+            "Coder: Specializing in generating and modifying code, this agent is skilled in various programming languages and is equipped to handle tasks ranging from quick fixes to developing complex algorithms.\n\n\n"
+        )
 
-        rLogger.initialize_logs("Researcher: dives into vast pools of information to find answers, evidence, or data that support the task at hand. Whether it's through browsing the web, accessing databases, or consulting academic journals, this agent is adept at gathering and synthesizing relevant information to aid in problem-solving.\n\n\n")
+        rLogger.initialize_logs(
+            "Researcher: dives into vast pools of information to find answers, evidence, or data that support the task at hand. Whether it's through browsing the web, accessing databases, or consulting academic journals, this agent is adept at gathering and synthesizing relevant information to aid in problem-solving.\n\n\n"
+        )
 
-        pLogger.initialize_logs("Planner: planner is tasked with orchestrating the overall strategy for solving user queries. Assesses the problem statement and determines the most effective sequence of actions, delegating tasks to other agents and tools as necessary. This agent ensures that Sirji's workflow is efficient and goal-oriented.\n\n\n")
+        pLogger.initialize_logs(
+            "Planner: planner is tasked with orchestrating the overall strategy for solving user queries. Assesses the problem statement and determines the most effective sequence of actions, delegating tasks to other agents and tools as necessary. This agent ensures that Sirji's workflow is efficient and goal-oriented.\n\n\n"
+        )
 
-        eLogger.initialize_logs("Executor: responsible for running code or scripts in a controlled environment, allowing for executing and testing activities. Executor verifies the correctness and efficacy of solutions before they are finalized and implements automated tasks as defined by the Planner.\n\n\n")
+        eLogger.initialize_logs(
+            "Executor: responsible for running code or scripts in a controlled environment, allowing for executing and testing activities. Executor verifies the correctness and efficacy of solutions before they are finalized and implements automated tasks as defined by the Planner.\n\n\n"
+        )
 
         sLogger.initialize_logs(
-            "Sirji: Sirji will automatically create a plan to solve the problem statement, prioritize it, organize research, write code, execute it, and fix issues.\n\n\n")
+            "Sirji: Sirji will automatically create a plan to solve the problem statement, prioritize it, organize research, write code, execute it, and fix issues.\n\n\n"
+        )
 
         uLogger.initialize_logs(
-            "User: The user is the person who interacts with Sirji. The user can ask questions, provide problem statements, and receive solutions from Sirji.\n\n\n")
+            "User: The user is the person who interacts with Sirji. The user can ask questions, provide problem statements, and receive solutions from Sirji.\n\n\n"
+        )
 
     def open_views(self):
         screen_width, screen_height = get_screen_resolution()
@@ -123,18 +131,21 @@ def open_views(self):
             (f"watch -n 1 'cat {pLogger.filepath}'", "Plan Progress"),
             (f"tail -f {rLogger.filepath}", "Research Agent"),
             (f"tail -f {cLogger.filepath}", "Coding Agent"),
-            (f"tail -f {eLogger.filepath}", "Execution Agent")
+            (f"tail -f {eLogger.filepath}", "Execution Agent"),
         ]
 
         current_directory = os.getcwd()
 
         # Prepend `cd {current_directory} &&` to each command to ensure it runs in the desired directory
-        command_title_pairs = [(f"cd {current_directory} && {command}", title)
-                               for command, title in command_title_pairs]
+        command_title_pairs = [
+            (f"cd {current_directory} && {command}", title)
+            for command, title in command_title_pairs
+        ]
 
         for i, (command, title) in enumerate(command_title_pairs):
             open_terminal_and_run_command(
-                command, title, i, window_width, window_height)
+                command, title, i, window_width, window_height
+            )
 
     def _parse_response(self, response_str):
         """
@@ -156,7 +167,7 @@ def handle_response(self, message):
             recipient = response.get("TO").strip()
             sender = response.get("FROM").strip()
             action = response.get("ACTION").strip()
-        except Exception as e:
+        except Exception:
             recipient = last_recipient
             message = textwrap.dedent(f"""
               ```
@@ -172,31 +183,30 @@ def handle_response(self, message):
             action = response.get("ACTION").strip()
 
         send_external_system_message(
-            f"Forwarding message from {sender} to {recipient} for action: {action}")
-
-        response_message = ''
+            f"Forwarding message from {sender} to {recipient} for action: {action}"
+        )
+
+        response_message = ""
 
-        if (action == "question" or action == "inform"):
+        if action == "question" or action == "inform":
             details = response.get("DETAILS")
             send_external_system_message(f"Question: {details}")
             enable_chat_send_button()
 
             user_input = self.wait_for_user_input()
 
-            ans_message = self.user.generate_answer_message(
-                user_input, 'Coder')
+            ans_message = self.user.generate_answer_message(user_input, "Coder")
 
             disable_chat_send_button()
             self.handle_response(ans_message)
-        elif (action == 'solution-complete'):
+        elif action == "solution-complete":
             details = response.get("DETAILS")
             send_external_system_message(f" Solution complete: {details}")
             enable_chat_send_button()
 
             user_input = self.wait_for_user_input()
 
-            ans_message = self.user.generate_feedback_message(
-                user_input, 'Coder')
+            ans_message = self.user.generate_feedback_message(user_input, "Coder")
             disable_chat_send_button()
             self.handle_response(ans_message)
 
@@ -223,7 +233,8 @@ def start(self):
         self.open_views()
 
         send_external_system_message(
-            "Hello! Please provide the problem you need me to solve.")
+            "Hello! Please provide the problem you need me to solve."
+        )
 
         while True:
             # Wait for a new message to become available
@@ -234,7 +245,8 @@ def start(self):
         disable_chat_send_button()
 
         ps_message = self.user.generate_problem_statement_message(
-            self.problem_statement, 'Coder')
+            self.problem_statement, "Coder"
+        )
 
         response_message = self.coder.message(ps_message)
 
@@ -258,6 +270,7 @@ def perform_non_gui_tasks():
 
 if __name__ == "__main__":
     empty_workspace()
+    load_dotenv()
 
     background_thread = threading.Thread(target=perform_non_gui_tasks)
     background_thread.start()

requirements.txt

@@ -4,4 +4,5 @@ beautifulsoup4==4.12.3
 pypdf2==3.0.1
 markdownify==0.11.6
 pyobjc==10.2
-playwright==1.42.0
+playwright==1.42.0
+python-dotenv==1.0.1

sirji/agents/coder.py

@@ -1,8 +1,5 @@
-from openai import OpenAI
-import os
-
+from sirji.config.model import ModelClient
 from sirji.prompts.coder import CoderPrompt
-
 from sirji.tools.logger import coder as cLogger
 
 
@@ -11,6 +8,7 @@ class SingletonMeta(type):
     This is a metaclass that will be used to create a Singleton class.
     It ensures that only one instance of the Singleton class exists.
     """
+
     _instances = {}
 
     def __call__(cls, *args, **kwargs):
@@ -24,37 +22,34 @@ class Coder(metaclass=SingletonMeta):
     def __init__(self):
         # Initialize conversation
         self.conversation = [
-            {'role': 'system', 'content': CoderPrompt().system_prompt()}]
+            {"role": "system", "content": CoderPrompt().system_prompt()}
+        ]
 
-        # Fetch OpenAI API key from environment variable
-        api_key = os.environ.get("SIRJI_OPENAI_API_KEY")
+        # Initialize ModelClient
+        model_client = ModelClient()
 
-        if api_key is None:
-            raise ValueError(
-                "OpenAI API key is not set as an environment variable")
-
-        # Initialize OpenAI client
-        self.client = OpenAI(api_key=api_key)
+        # Set self.client as the value of self.client set in ModelClient class
+        self.client = model_client.client
+        self.model = model_client.model
 
     def message(self, input_message):
         # Append user's input message to the conversation
-        self.conversation.append({'role': 'user', 'content': input_message})
+        self.conversation.append({"role": "user", "content": input_message})
 
         cLogger.info(f"Incoming: \n{input_message}")
 
         cLogger.info("Calling OpenAI Chat Completions API")
 
         chat_completion = self.client.chat.completions.create(
             messages=self.conversation,
-            model="gpt-4-turbo-preview",
+            model=self.model,
             temperature=0,
             max_tokens=4095,
         )
 
         response_message = chat_completion.choices[0].message.content
 
-        self.conversation.append(
-            {'role': 'assistant', 'content': response_message})
-
+        self.conversation.append({"role": "assistant", "content": response_message})
+
         cLogger.info(f"Outgoing: \n{response_message}")
         return response_message