Add loacalAi

nervousapps · Sep 18, 2023 · 2ab123c · 2ab123c
1 parent de5f4b4
commit 2ab123c
Show file tree

Hide file tree

Showing 11 changed files with 142 additions and 37 deletions.
diff --git a/README.md b/README.md
@@ -59,7 +59,7 @@ python --version
 python3 --version
 ```
 
-- [Open API key](https://platform.openai.com/account/api-keys)
+- [Open API key](https://platform.openai.com/account/api-keys), or a [LoacalAI](https://localai.io/) instance, or any other GenAI instance compatible with the openAi python package.
 
 ## :surfing_woman: Installation
 1 - Clone this repository
@@ -179,6 +179,8 @@ An example can be given with the keyword "EXAMPLE:", it will be added to employe
 | CEO_GUIDELINES                 | "As a CEO, I want ... PRODUCT ... "                                         | Default guidelines |
 | INTERACTIVE                    | Wait for user entry (basic for now), no employee interaction if asynchronous mode is activated                                        | "no" |
 | ASYNCHRONOUS                    | Asynchronous mode, employees can work in parallel                                        | "yes" |
+| LOCAL_AI_URL                    | Url of a local instance of an api compatible with openAi python package                                        | "" |
+| MODEL_NAME                    | Model name to use                                        | gpt-3.5-turbo-16k |
 
 ## :roller_coaster: Going further
 - rework managers prompt :curly_loop:

diff --git a/config b/config
@@ -1,6 +1,7 @@
 # OpenAI stuff
 KEYFILE="./openai_key.txt"
-GPT_VERSION="gpt-3.5-turbo-16k"
+MODEL_NAME="lunademo"
+LOCAL_AI_URL="http://localhost:8080/v1"
 
 # General
 INTERACTIVE="no"

diff --git a/python/gpt_enterprise/__main__.py b/python/gpt_enterprise/__main__.py
@@ -5,6 +5,7 @@
 import sys
 import asyncio
 import json
+import openai
 from dotenv import load_dotenv
 from gpt_enterprise.enterprise import Enterprise
 
@@ -18,6 +19,18 @@ def main():
     guidelines = os.getenv("CEO_GUIDELINES")
     output_directory = os.getenv("OUTPUT_DIRECTORY")
     manager_retry = int(os.getenv("MANAGER_RETRY"))
+    local_ai_url = os.getenv("LOCAL_AI_URL")
+
+    if local_ai_url:
+        # Change api url to LoacalAi one
+        openai.api_base = local_ai_url
+        openai.api_key = "sx-xxx"
+        OPENAI_API_KEY = "sx-xxx"
+        os.environ['OPENAI_API_KEY'] = OPENAI_API_KEY
+    else:
+        # Initialize openai api_key
+        with open("./openai_key.txt", "r") as file:
+            openai.api_key = (file.read()).strip()
 
     # Interactive mode, change variables with user input
     if interactive:
@@ -49,7 +62,6 @@ def main():
 
     # Create the enterprise
     enterprise = Enterprise(
-        keyfile=keyfile,
         company_name=company_name,
         guidelines=guidelines,
         output_directory=output_directory,

diff --git a/python/gpt_enterprise/enterprise.py b/python/gpt_enterprise/enterprise.py
@@ -4,8 +4,6 @@
 \U0001F3E2
 """
 import os
-import time
-import openai
 
 from gpt_enterprise.team_leader import TeamLeader
 from gpt_enterprise.scrum_master import ScrumMaster
@@ -18,7 +16,6 @@ class Enterprise:
 
     def __init__(
         self,
-        keyfile: str,
         guidelines: str,
         output_directory: str,
         manager_retry: int,
@@ -32,17 +29,13 @@ def __init__(
             - a scrum master that wil be responsible of managing tasks
 
         Args:
-            keyfile (str): Pth to openai key file
             guidelines (str): CEO guidelines
             output_directory (str): Output directory
             manager_retry (int): How many times manager will retry to do the plan
             company_name (str, optional): Enterprise's name. Defaults to "GPTenterprise".
             interactive (bool): Defaults to False
             asynchronous (bool): Defaults to True
         """
-        # Initialize openai api_key
-        with open(keyfile, "r") as file:
-            openai.api_key = (file.read()).strip()
         self.company_name = company_name
         self.employees = {}
         self.tasks_board = []

diff --git a/python/gpt_enterprise/gpt_utils.py b/python/gpt_enterprise/gpt_utils.py
@@ -13,7 +13,7 @@
 
 
 def generate_text(
-    system_prompt: str, user_prompt: str, model: str, temperature: float
+    system_prompt: str, user_prompt: str, temperature: float, model: str = os.getenv("MODEL_NAME", "gpt-3.5-turbo-16k")
 ) -> Generator:
     """
 
@@ -68,7 +68,7 @@ def generate_image(
     # Ask ChatGPT a prompt to generate image with DALL-E
     with open(os.path.join(EMPLOYEE_PROMPTS_PATH, "dall_e_prompter.txt"), "r") as file:
         response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
+            model=os.getenv("MODEL_NAME", "gpt-3.5-turbo-16k"),
             messages=[
                 # Initialize ChatGPT to be a helpful assistant but that it remains the employee
                 {

diff --git a/python/gpt_enterprise/scrum_master.py b/python/gpt_enterprise/scrum_master.py
@@ -216,19 +216,30 @@ def _wait_for_result(self, task_index: int, task: dict, employee: Employee) -> d
         """
         if "no" not in str(task.get("requirements", "no")):
             counter = 0
-            while not self.tasks[int(task["requirements"])].get("result"):
-                time.sleep(0.01)
-                counter += 1
-                if counter >= 2000:
-                    print(
-                        f"Waiting for {self.tasks[int(task['requirements'])]['employee_name']} to finish..."
-                    )
-                    counter = 0
-            # Add the previous employee work to the current task to
-            # be used by the assigned employee.
-            task[
-                "todo"
-            ] += f" Here is the work done by {self.tasks[int(task['requirements'])]['employee_name']} : {self.tasks[int(task['requirements'])]['result']}"
+            # Test if requirements is an int or a list
+            try:
+                task_requirements = list(task["requirements"])
+            except TypeError:
+                task_requirements = [int(task["requirements"])]
+            for task_idx in task_requirements:
+                # Test if it is an int, if not continue
+                try:
+                    task_idx = int(task_idx)
+                except Exception as error:
+                    continue
+                while not self.tasks[task_idx].get("result"):
+                    time.sleep(0.01)
+                    counter += 1
+                    if counter >= 2000:
+                        print(
+                            f"Waiting for {self.tasks[int(task['requirements'])]['employee_name']} to finish..."
+                        )
+                        counter = 0
+                # Add the previous employee work to the current task to
+                # be used by the assigned employee.
+                task[
+                    "todo"
+                ] += f" Here is the work done by {self.tasks[task_idx]['employee_name']} : {self.tasks[task_idx]['result']}"
 
         print(
             f"{self.emoji} {employee.name} is doing task {task_index} : {task['todo']}"

diff --git a/python/gpt_enterprise/team_leader.py b/python/gpt_enterprise/team_leader.py
@@ -91,6 +91,8 @@ def hire_employees(self) -> List[Dict[str, Employee]]:
                 print(
                     f"\n {self.emoji} I've messed up, retrying to find employees... \n Error : \n {err}\n"
                 )
+        else:
+            raise err
         print(f"\n Ok, I've hired them ! Please welcome :\n")
         hired_employees = {}
         if employees_to_hire:

diff --git a/python/tests/conftest.py b/python/tests/conftest.py
@@ -2,6 +2,7 @@
 import json
 import time
 import pytest
+import openai
 
 from gpt_enterprise.employee import Employee
 from gpt_enterprise.scrum_master import ScrumMaster
@@ -13,6 +14,12 @@
 EMPLOYEES_FILES_DIR = os.path.join(TEST_FILES_DIR, "employees")
 
 
+# Change api url to LoacalAi one
+openai.api_base = "local_ai_url"
+openai.api_key = "sx-xxx"
+OPENAI_API_KEY = "sx-xxx"
+os.environ['OPENAI_API_KEY'] = OPENAI_API_KEY
+
 @pytest.fixture
 def scrum_master_test():
     yield ScrumMaster(ceo_guidelines="Test", manager_retry=5, output_directory="Test")

diff --git a/python/tests/test_enterprise.py b/python/tests/test_enterprise.py
@@ -16,9 +16,6 @@
 
 def test_enterprise():
     enterprise = Enterprise(
-        keyfile=os.path.join(
-            os.path.dirname(__file__), "..", "..", "openai_key.txt.template"
-        ),
         guidelines="test",
         output_directory="test",
         manager_retry=1,
@@ -96,9 +93,6 @@ def mock_tl_generate_text(
 
     # Run asynchronously
     enterprise_async = Enterprise(
-        keyfile=os.path.join(
-            os.path.dirname(__file__), "..", "..", "openai_key.txt.template"
-        ),
         guidelines="test",
         output_directory="test",
         manager_retry=1,
@@ -113,9 +107,6 @@ def mock_tl_generate_text(
 
     # Run sequentially
     enterprise_seq = Enterprise(
-        keyfile=os.path.join(
-            os.path.dirname(__file__), "..", "..", "openai_key.txt.template"
-        ),
         guidelines="test",
         output_directory="test",
         manager_retry=1,

diff --git a/python/tests/test_files/tasks/tasks_requirements_list.txt b/python/tests/test_files/tasks/tasks_requirements_list.txt
@@ -0,0 +1,79 @@
+[
+    {
+        "task_name": "Create Product Ideas",
+        "employee_name": "Alice",
+        "todo": "Propose three different innovative and eco-friendly product ideas that can make users' daily lives easier while reducing waste or energy consumption. Provide sketches or computer designs of the products, as well as the materials and the technology used to build them. The prototypes should be functional and visually appealing.",
+        "type": "image",
+        "requirements": "no"
+    },
+    {
+        "task_name": "Create Product Ideas",
+        "employee_name": "Alice",
+        "todo": "Propose three different innovative and eco-friendly product ideas that can make users' daily lives easier while reducing waste or energy consumption. Provide sketches or computer designs of the products, as well as the materials and the technology used to build them. The prototypes should be functional and visually appealing.",
+        "type": "image",
+        "requirements": "no"
+    },
+    {
+        "task_name": "Create Product Ideas",
+        "employee_name": "Alice",
+        "todo": "Propose three different innovative and eco-friendly product ideas that can make users' daily lives easier while reducing waste or energy consumption. Provide sketches or computer designs of the products, as well as the materials and the technology used to build them. The prototypes should be functional and visually appealing.",
+        "type": "image",
+        "requirements": "no"
+    },
+    {
+        "task_name": "Create Product Ideas",
+        "employee_name": "Alice",
+        "todo": "Propose three different innovative and eco-friendly product ideas that can make users' daily lives easier while reducing waste or energy consumption. Provide sketches or computer designs of the products, as well as the materials and the technology used to build them. The prototypes should be functional and visually appealing.",
+        "type": "image",
+        "requirements": "no"
+    },
+    {
+        "task_name": "Choose Products to Develop",
+        "employee_name": "Charlie",
+        "todo": "Evaluate the product ideas proposed and choose the most promising ones. Provide a detailed report of the reasons behind your choice, including the market potential, the feasibility, and the overall appeal of the products. Also, indicate which employee(s) will be responsible for designing, writing, and creating visualization of the chosen product.",
+        "type": "text",
+        "requirements": 0
+    },
+    {
+        "task_name": "Write Product Descriptions",
+        "employee_name": "Bob",
+        "todo": "Write a description of each chosen product in less than 200 words. The articles should be persuasive, informative and engaging, highlighting the benefits and the eco-friendliness of the product in a concise and clear language. Use SEO keywords effectively. ",
+        "type": "text",
+        "requirements": 1
+    },
+    {
+        "task_name": "Design Product Prototypes",
+        "employee_name": "Alice",
+        "todo": "Design and create prototypes of the chosen products following the eco-friendly and sustainable guidelines. Provide sketches and computer designs of the prototypes. Test them and provide a report of the results, including the improvements and the changes needed. Deliver the prototypes to the development team.",
+        "type": "image",
+        "requirements": "1,2"
+    },
+    {
+        "task_name": "Create Product Images",
+        "employee_name": "Charlie",
+        "todo": "Create vector and 3D images of the prototypes in different sizes and formats. The images should convey the functionality and the eco-friendliness of the products in a visually appealing way, consistent with our branding guidelines. Deliver the images to the marketing team.",
+        "type": "image",
+        "requirements": "[1,2,3]"
+    },
+    {
+        "task_name": "Present Product to CEO",
+        "employee_name": "Bob",
+        "todo": "Present the products to the CEO in a visually appealing and informative way, including the prototypes, the product descriptions, and the product images. Explain the benefits and the market potential of the products, as well as the challenges and the improvements needed. ",
+        "type": "text",
+        "requirements": [1,2,3]
+    },
+    {
+        "task_name": "Create Product Article for Website",
+        "employee_name": "Bob",
+        "todo": "Write an article presenting the chosen products that can be used on the company website. The article should be written for the general audience, highlighting the benefits and the eco-friendliness of the products in less than 500 words. The article should include at least two product images and be SEO-optimized. The article must end with the conclusion that the product create the wanted PRODUCT. ",
+        "type": "text",
+        "requirements": 4
+    },
+    {
+        "task_name": "Write Product Conclusion",
+        "employee_name": "Alice",
+        "todo": "Write a final conclusion of the products submitted that make the wanted PRODUCT. Explain how each product helps at making life easier and greener, and how they complete each other. Tell us why these products are innovative and how they help the planet by being eco-friendly. ",
+        "type": "text",
+        "requirements": 6
+    }
+]
diff --git a/python/tests/test_scrum_master.py b/python/tests/test_scrum_master.py
@@ -64,16 +64,23 @@ def test_do_plan(mocker, scrum_master_test, fake_employees):
 
 # TODO: add parametrize with test files
 @pytest.mark.asyncio
-async def test_do_plan_async(mocker, scrum_master_test, fake_employees):
+@pytest.mark.parametrize("test_file", [
+    "tasks.txt",
+    "tasks_requirements_list.txt"
+])
+async def test_do_plan_async(mocker, scrum_master_test, fake_employees, test_file):
     # Mock function and method that requests openai API (to avoid costs)
     mocker.patch(
-        "gpt_enterprise.employee.generate_text", return_value="This is only a test"
+        "gpt_enterprise.employee.generate_text",
+        return_value=mock_open_ai_response_object(
+            mocker=mocker, content="This is only a test"
+        ),
     )
     mocker.patch(
         "gpt_enterprise.employee.generate_image",
         return_value=("Do something", ["./img.jpg"]),
     )
-    with open(os.path.join(TASKS_FILES_DIR, "tasks.txt"), "r") as file:
+    with open(os.path.join(TASKS_FILES_DIR, test_file), "r") as file:
         plan = ast.literal_eval(file.read())
 
     production = await scrum_master_test.do_plan_async(plan, fake_employees)