Merge pull request #169 from Cloud-Code-AI/167-updating-the-readme-for-the-launch

feat: Added custom model selection

Author: sauravpanda

README.md

@@ -74,7 +74,7 @@ PYTHONPATH=. poetry run python examples/basic/execute.py
 Or using the default pytest module:
 
 ```
-pytest -v .kaizen/tests/
+pytest -v .kaizen/ui-tests/
 ```
 
 Kaizen will execute the generated tests and provide detailed reports.

config.json

@@ -2,10 +2,24 @@
     "language_model": {
         "provider": "litellm",
         "enable_observability_logging": true,
-        "model": {
-            "name": "azure/gpt4-o",
-            "input_token_cost": 0.000005,
-            "output_token_cost": 0.000015
+        "model_config": {
+            "model": "gpt-3.5-turbo-16k",
+            "input_cost_per_token": 0.0000005,
+            "output_cost_per_token": 0.0000015
+        },
+        "models": {
+            "best": {
+                "model": "gpt-4o",
+                "input_cost_per_token": 0.000005,
+                "output_cost_per_token": 0.000015,
+                "type": "best"
+            },
+            "default": {
+                "model": "gpt-3.5-turbo-16k",
+                "input_cost_per_token": 0.0000005,
+                "output_cost_per_token": 0.0000015,
+                "type": "default"
+            }
         }
     },
     "github_app": {

kaizen/generator/ui.py

@@ -14,6 +14,11 @@ class UITestGenerator:
     def __init__(self):
         self.logger = logging.getLogger(__name__)
         self.provider = LLMProvider(system_prompt=UI_TESTS_SYSTEM_PROMPT)
+        self.custom_model = None
+        if self.provider.models and "best" in self.provider.models:
+            self.custom_model = self.provider.models["best"]
+            if "type" in self.custom_model:
+                del self.custom_model["type"]
 
     def generate_ui_tests(
         self,
@@ -45,7 +50,7 @@ def identify_modules(self, web_content: str, user: Optional[str] = None):
         This method identifies the different UI modules from a webpage.
         """
         prompt = UI_MODULES_PROMPT.format(WEB_CONTENT=web_content)
-        resp, usage = self.provider.chat_completion(prompt, user=user)
+        resp, usage = self.provider.chat_completion(prompt, user=user, custom_model=self.custom_model)
         modules = parser.extract_multi_json(resp)
         return {"modules": modules, "usage": usage}
 
@@ -63,7 +68,7 @@ def generate_playwright_code(
             WEB_CONTENT=web_content, TEST_DESCRIPTION=test_description, URL=web_url
         )
 
-        resp, usage = self.provider.chat_completion(prompt, user=user)
+        resp, usage = self.provider.chat_completion(prompt, user=user, custom_model=self.custom_model)
 
         return {"code": resp, "usage": usage}
 
@@ -92,7 +97,7 @@ def store_tests_files(self, json_tests: list, folder_path: str = ""):
         if not folder_path:
             folder_path = output.get_parent_folder()
 
-        folder_path = os.path.join(folder_path, ".kaizen/tests")
+        folder_path = os.path.join(folder_path, ".kaizen/ui-tests")
         output.create_folder(folder_path)
         output.create_test_files(json_tests, folder_path)
 

kaizen/llms/provider.py

@@ -17,6 +17,10 @@ def __init__(
         self.model_config = model_config
         if "default_model_config" in self.config.get("language_model", {}):
             self.model_config = self.config["language_model"]["default_model_config"]
+        
+        if "models" in self.config.get("language_model"):
+            self.models = self.config["language_model"]["models"]
+
         self.model = self.model_config["model"]
         if self.config.get("language_model", {}).get(
             "enable_observability_logging", False
@@ -25,13 +29,15 @@ def __init__(
             litellm.success_callback = ["supabase"]
             litellm.failure_callback = ["supabase"]
 
-    def chat_completion(self, prompt, user: str = None):
+    def chat_completion(self, prompt, user: str = None, custom_model=None):
         messages = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": prompt},
         ]
+        if not custom_model:
+            custom_model = self.model_config
 
-        response = litellm.completion(messages=messages, user=user, **self.model_config)
+        response = litellm.completion(messages=messages, user=user, **custom_model)
         return response["choices"][0]["message"]["content"], response["usage"]
 
     def is_inside_token_limit(self, PROMPT, percentage=0.7):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kaizen-cloudcode"
-version = "0.1.17"
+version = "0.2.0"
 description = "An intelligent coding companion that accelerates your development workflow by providing efficient assistance, enabling you to craft high-quality code more rapidly."
 authors = ["Saurav Panda <[email protected]>"]
 license = "Apache2.0"