Update private_gpt/components/llm/llm_component.py

private_gpt/components/llm/llm_component.py

@@ -70,39 +70,39 @@ class LLMComponent:
                 self.llm = OpenAI(model="gpt-3.5-turbo", api_key=os.environ.get("OPENAI_API_KEY"))
                 
                 
-                def switch_model(new_model: str) -> None:
-                    nonlocal self 
-                    from llama_index.llms import LlamaCPP
-                    openai_settings = settings.openai.api_key
-                    
-                    if new_model == "gpt-3.5-turbo":
-                        self.llm = OpenAI(model="gpt-3.5-turbo", api_key=os.environ.get("OPENAI_API_KEY"))
-                        logger.info("Initializing the LLM Model in=%s", "gpt-3.5-turbo")
-                        
-                    elif new_model == "gpt-4":
-                        # Initialize with the new model 
-                        self.llm = OpenAI(model="gpt-4", api_key=os.environ.get("OPENAI_API_KEY"))
-                        logger.info("Initializing the LLM Model in=%s", "gpt-4")
+    def switch_model(new_model: str) -> None:
+        nonlocal self 
+        from llama_index.llms import LlamaCPP
+        openai_settings = settings.openai.api_key
+        
+        if new_model == "gpt-3.5-turbo":
+            self.llm = OpenAI(model="gpt-3.5-turbo", api_key=os.environ.get("OPENAI_API_KEY"))
+            logger.info("Initializing the LLM Model in=%s", "gpt-3.5-turbo")
             
-                    
-                    elif new_model == "mistral-7B":
-                        prompt_style_cls = get_prompt_style(settings.local.prompt_style)
-                        prompt_style = prompt_style_cls(
-                            default_system_prompt=settings.local.default_system_prompt
-                        )
-                        self.llm = LlamaCPP(
-                            #model_path=str(models_path / settings.local.llm_hf_model_file),
-                            model_url= "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q4_K_M.gguf?download=true",
-                            temperature=0.1,
-                            max_new_tokens=settings.llm.max_new_tokens,
-                            context_window=3900,
-                            generate_kwargs={},
-                            model_kwargs={"n_gpu_layers": -1},
-                            messages_to_prompt=prompt_style.messages_to_prompt,
-                            completion_to_prompt=prompt_style.completion_to_prompt,
-                            verbose=True,
-                        )
-                        logger.info("Initializing the LLM Model in=%s", "Mistral-7B")
+        elif new_model == "gpt-4":
+            # Initialize with the new model 
+            self.llm = OpenAI(model="gpt-4", api_key=os.environ.get("OPENAI_API_KEY"))
+            logger.info("Initializing the LLM Model in=%s", "gpt-4")
+
+        
+        elif new_model == "mistral-7B":
+            prompt_style_cls = get_prompt_style(settings.local.prompt_style)
+            prompt_style = prompt_style_cls(
+                default_system_prompt=settings.local.default_system_prompt
+            )
+            self.llm = LlamaCPP(
+                #model_path=str(models_path / settings.local.llm_hf_model_file),
+                model_url= "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF/resolve/main/mistral-7b-instruct-v0.1.Q4_K_M.gguf?download=true",
+                temperature=0.1,
+                max_new_tokens=settings.llm.max_new_tokens,
+                context_window=3900,
+                generate_kwargs={},
+                model_kwargs={"n_gpu_layers": -1},
+                messages_to_prompt=prompt_style.messages_to_prompt,
+                completion_to_prompt=prompt_style.completion_to_prompt,
+                verbose=True,
+            )
+            logger.info("Initializing the LLM Model in=%s", "Mistral-7B")
                         
         def switch_to_model(self, new_model: str):
             if self.llm_mode == "dynamic":