- adding LLM class which contains multiple connectors including ollama models

- modifying system prompts
- modifying all class to rely on the new LLM class
- Modifying UI to get LLM model from user

Config_files/message_system_config.json

@@ -20,7 +20,8 @@
   "AI_Jargon": ["elevate", "enhance", "ignite", "reignite", "rekindle", "rediscover","passion", "boost", "fuel", "thrill", "revive", "spark", "performing", "fresh", "tone", "enthusiasm", "illuminate"],
   "AI_phrases_singeo": ["your voice deserves more"],
   "header_limit": 30,
-  "message_limit": 110
+  "message_limit": 110,
+  "LLM_models": ["4o-mini", "gpt-4o", "deepseek-r1:1.5b", "gemma3:4b"]
 }
 
 

Messaging_system/CoreConfig.py

@@ -21,9 +21,8 @@ class CoreConfig:
 
         # LLM configs
         self.api_key = None  # will be set by user
-        self.model = "gpt-4o" # will be set by user
+        self.model = "gpt-4o" # default -> will be set by user
         self.temperature = 0.7
-        # self.model = "gpt-4.1-nano" # will be set by user
 
         # will be set by user
         self.CTA = None
@@ -80,6 +79,12 @@ class CoreConfig:
 
         self.message_style = message_style
 
+    # --------------------------------------------------------------
+    # --------------------------------------------------------------
+
+    def set_llm_model(self, model):
+        self.model = model
+
     # --------------------------------------------------------------
     # --------------------------------------------------------------
     def set_involve_recsys_result(self, involve_recsys_result):

Messaging_system/LLM.py

@@ -6,31 +6,58 @@ import json
 import time
 from openai import OpenAI
 import openai
-from tqdm import tqdm
+import ollama
+import torch
+import re
 
 
 
 class LLM:
-    def __init__(self, Core, llm):
+    def __init__(self, Core):
         self.Core = Core
-        self.llm = llm
 
 
+        self.model = None
+        self.model_type = "openai" # valid values -> ["openai", "ollama"]
+        self.client = None
+        self.connect_to_llm()
+
+
+    def get_response(self, prompt, instructions):
+        if self.model_type == "openai":
+            response = self.get_message_openai(prompt, instructions)
+        elif self.model_type == "ollama":
+            response = self.get_message_ollama(prompt, instructions)
+        else:
+            raise f"Invalid model type : {self.model_type}"
+
+        return response
 
     def connect_to_llm(self):
         """
-        connect to selected llm
+        connect to selected llm -> ollama or openai connection
         :return:
         """
+        openai_models = ["4o-mini", "gpt-4o", "gpt-4.1-nano"]
+        ollama_models = ["deepseek-r1:1.5b", "gemma3:4b", "deepseek-r1:7b"]
+
+        if self.Core.model in openai_models:
+            self.model_type = "openai"
+
+        if self.Core.model in ollama_models:
+            self.model_type = "ollama"
+            self.client = ollama.Client()
+
+        self.model = self.Core.model
+
 
 
-    def get_message(self, prompt, max_retries=4):
+    def get_message_openai(self, prompt, instructions, max_retries=4):
         """
-        sending the prompt to the LLM and get back the response
+        sending the prompt to openai LLM and get back the response
         """
 
         openai.api_key = self.Core.api_key
-        instructions = self.llm_instructions()
         client = OpenAI(api_key=self.Core.api_key)
 
         for attempt in range(max_retries):
@@ -59,7 +86,6 @@ class LLM:
                     # Extract JSON code block
 
                     output = json.loads(content)
-                    # output = json.loads(response.choices[0].message.content)
 
                     if 'message' not in output or 'header' not in output:
                         print(f"'message' or 'header' is missing in response on attempt {attempt + 1}. Retrying...")
@@ -92,4 +118,111 @@ class LLM:
                 print(e.response)
 
         print("Max retries exceeded. Returning empty response.")
-        return None
+        return None
+
+    # ======================================================================
+
+    def get_message_ollama(self, prompt, instructions, max_retries=10):
+        """
+        Send the prompt to the LLM and get back the response.
+        Includes handling for GPU memory issues by clearing cache and waiting before retry.
+        """
+        prompt = instructions + "\n \n" + prompt
+        for attempt in range(max_retries):
+            try:
+                # Try generating the response
+                response = self.client.generate(model=self.model, prompt=prompt, format='json')
+            except Exception as e:
+                # This catches errors like the connection being forcibly closed
+                print(f"Error on attempt {attempt + 1}: {e}.")
+                try:
+                    # Clear GPU cache if you're using PyTorch; this may help free up memory
+                    torch.cuda.empty_cache()
+                    print("Cleared GPU cache.")
+                except Exception as cache_err:
+                    print("Failed to clear GPU cache:", cache_err)
+                # Wait a bit before retrying to allow memory to recover
+                time.sleep(2)
+                continue
+
+            try:
+                tokens = {
+                    'prompt_tokens': 0,
+                    'completion_tokens': 0,
+                    'total_tokens': 0
+                }
+
+                try:
+                    output = self.preprocess_and_parse_json(response.response)
+                    if output is None:
+                        continue
+
+                    if 'message' not in output or 'header' not in output:
+                        print(f"'message' or 'header' is missing in response on attempt {attempt + 1}. Retrying...")
+                        continue  # Continue to next attempt
+
+                    else:
+                        if len(output["header"].strip()) > self.Core.config_file["header_limit"] or len(
+                                output["message"].strip()) > self.Core.config_file["message_limit"]:
+                            print(
+                                f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
+                            continue
+
+                except json.JSONDecodeError:
+                    print(f"Invalid JSON from LLM on attempt {attempt + 1}. Retrying...")
+            except Exception as parse_error:
+                print("Error processing output:", parse_error)
+
+        print("Max retries exceeded. Returning empty response.")
+        return [], {}
+
+    # ======================================================================
+
+    # def preprocess_and_parse_json(self, response):
+    #     # Remove any leading/trailing whitespace and newlines
+    #     if response.startswith('```json') and response.endswith('```'):
+    #         response = response[len('```json'):-len('```')].strip()
+    #
+    #     # Parse the cleaned response into a JSON object
+    #     try:
+    #         json_object = json.loads(response)
+    #         return json_object
+    #     except json.JSONDecodeError as e:
+    #         print(f"Failed to parse JSON: {e}")
+    #         return None
+
+    # =====================================================================
+
+    import re
+    import json
+
+    def preprocess_and_parse_json(self, response: str):
+        """
+        Cleans an LLM response by removing <think> tags and extracting JSON
+        from ```json ... ``` fences (or bare text if no fence is found),
+        then returns the parsed object or None on failure.
+        """
+        # 1) Remove all <think>...</think> blocks
+        cleaned = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
+
+        # 2) Look for a ```json ... ``` fenced block
+        fence_pattern = re.compile(r'```json(.*?)```', flags=re.DOTALL)
+        fence_match = fence_pattern.search(cleaned)
+        if fence_match:
+            json_text = fence_match.group(1).strip()
+        else:
+            # No fence; assume whole cleaned text is JSON
+            json_text = cleaned.strip()
+
+        # 3) Attempt to parse
+        try:
+            return json.loads(json_text)
+        except json.JSONDecodeError as e:
+            print(f"Failed to parse JSON: {e}")
+            # Optionally, log the offending text for debugging:
+            # print("Offending text:", json_text)
+            return None
+
+
+
+

Messaging_system/Message_generator.py

@@ -10,12 +10,14 @@ import streamlit as st
 from Messaging_system.MultiMessage import MultiMessage
 from Messaging_system.protection_layer import ProtectionLayer
 import openai
+from Messaging_system.LLM import LLM
 
 
 class MessageGenerator:
 
     def __init__(self, CoreConfig):
         self.Core = CoreConfig
+        self.llm = LLM(CoreConfig)
 
     # --------------------------------------------------------------
     # --------------------------------------------------------------
@@ -36,12 +38,11 @@ class MessageGenerator:
                 progress_callback(progress, total_users)
 
             if row["prompt"] is not None:
-                first_message = self.get_llm_response(row["prompt"])
+                first_message = self.llm.get_response(prompt=row["prompt"], instructions=self.llm_instructions())
 
                 if first_message is not None:
                     # adding protection layer
-                    protect = ProtectionLayer(config_file=self.Core.config_file,
-                                              messaging_mode=self.Core.messaging_mode)
+                    protect = ProtectionLayer(CoreConfig=self.Core)
                     message, total_tokens = protect.criticize(message=first_message, user=row)
 
                     # updating tokens
@@ -171,80 +172,6 @@ class MessageGenerator:
         }
         return output_message
 
-    # --------------------------------------------------------------
-    # --------------------------------------------------------------
-
-    def get_llm_response(self, prompt, max_retries=4):
-        """
-        sending the prompt to the LLM and get back the response
-        """
-
-        openai.api_key = self.Core.api_key
-        instructions = self.llm_instructions()
-        client = OpenAI(api_key=self.Core.api_key)
-
-        for attempt in range(max_retries):
-            try:
-                response = client.chat.completions.create(
-                    model=self.Core.model,
-                    response_format={"type": "json_object"},
-                    messages=[
-                        {"role": "system", "content": instructions},
-                        {"role": "user", "content": prompt}
-                    ],
-                    max_tokens=500,
-                    n=1,
-                    # temperature=self.Core.temperature,
-                    temperature=0.7,
-                )
-
-                tokens = {
-                    'prompt_tokens': response.usage.prompt_tokens,
-                    'completion_tokens': response.usage.completion_tokens,
-                    'total_tokens': response.usage.total_tokens
-                }
-
-                try:
-                    content = response.choices[0].message.content
-
-                    # Extract JSON code block
-
-                    output = json.loads(content)
-                    # output = json.loads(response.choices[0].message.content)
-
-                    if 'message' not in output or 'header' not in output:
-                        print(f"'message' or 'header' is missing in response on attempt {attempt + 1}. Retrying...")
-                        continue  # Continue to next attempt
-
-                    else:
-                        if len(output["header"].strip()) > self.Core.config_file["header_limit"] or len(
-                                output["message"].strip()) > self.Core.config_file["message_limit"]:
-                            print(
-                                f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
-                            continue
-
-                    # validating the JSON
-                    self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    self.Core.temp_token_counter += tokens['total_tokens']
-                    return output
-
-                except json.JSONDecodeError:
-                    print(f"Invalid JSON from LLM on attempt {attempt + 1}. Retrying...")
-
-            except openai.APIConnectionError as e:
-                print("The server could not be reached")
-                print(e.__cause__)  # an underlying Exception, likely raised within httpx.
-            except openai.RateLimitError as e:
-                print("A 429 status code was received; we should back off a bit.")
-            except openai.APIStatusError as e:
-                print("Another non-200-range status code was received")
-                print(e.status_code)
-                print(e.response)
-
-        print("Max retries exceeded. Returning empty response.")
-        return None
-
     # --------------------------------------------------------------
     # --------------------------------------------------------------
     def llm_instructions(self):
@@ -256,7 +183,7 @@ class MessageGenerator:
         jargon_list = "\n".join(f"- {word}" for word in self.Core.config_file["AI_Jargon"])
 
         instructions = f"""
-        You are texting a friend and it should feel like a friendly encouraging nudge. Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like natural everyday speech: friendly, concise, no jargon, and following the instructions.
+        You are a copywriter. Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like natural everyday speech: friendly, concise, no jargon, and following the instructions.
         Write a SUPER CASUAL and NATURAL push notification, as if you are chatting over coffee. Avoid odd phrasings. 
         
         ABSOLUTE RULE – OVERRIDES EVERYTHING ELSE:   

Messaging_system/MultiMessage.py

@@ -3,6 +3,7 @@ import time
 from openai import OpenAI
 from Messaging_system.protection_layer import ProtectionLayer
 import openai
+from Messaging_system.LLM import LLM
 
 class MultiMessage:
     def __init__(self, CoreConfig):
@@ -11,6 +12,7 @@ class MultiMessage:
         for each user, building on previously generated messages.
         """
         self.Core = CoreConfig
+        self.llm = LLM(CoreConfig)
 
     # --------------------------------------------------------------
     def generate_multi_messages(self, user):
@@ -41,8 +43,7 @@ class MultiMessage:
 
         # We'll reuse the same ProtectionLayer
         protect = ProtectionLayer(
-            config_file=self.Core.config_file,
-            messaging_mode=self.Core.messaging_mode
+            CoreConfig=self.Core
         )
 
         # If user requested multiple messages, generate the rest
@@ -107,7 +108,8 @@ class MultiMessage:
         prompt = self.generate_prompt(context, step)
 
         # 2) Call our existing LLM routine
-        response_dict = self.get_llm_response(prompt)
+        response_dict = self.llm.get_response(prompt=prompt, instructions=self.llm_instructions())
+
         return response_dict
 
     # ===============================================================
@@ -292,78 +294,6 @@ Return only JSON of the form:
         return output_message
 
     # --------------------------------------------------------------
-    def get_llm_response(self, prompt, max_retries=4):
-        """
-        Calls the LLM (similar to MessageGenerator) with the prompt, returning a dict
-        with keys like 'header' and 'message' if successful, or None otherwise.
-
-        :param prompt: The text prompt for the LLM.
-        :param max_retries: Number of retries for potential LLM/connection failures.
-        :return: Dictionary with 'header' and 'message', or None if unsuccessful.
-        """
-        openai.api_key = self.Core.api_key
-        instructions = self.llm_instructions()
-        client = OpenAI(api_key=self.Core.api_key)
-
-        for attempt in range(max_retries):
-            try:
-                response = client.chat.completions.create(
-                    model=self.Core.model,
-                    response_format={"type": "json_object"},
-                    messages=[
-                        {"role": "system", "content": instructions},
-                        {"role": "user", "content": prompt}
-                    ],
-                    max_tokens=500,
-                    n=1,
-                    # temperature=self.Core.temperature
-                    temperature=0.7
-                )
-
-                tokens = {
-                    'prompt_tokens': response.usage.prompt_tokens,
-                    'completion_tokens': response.usage.completion_tokens,
-                    'total_tokens': response.usage.total_tokens
-                }
-
-                try:
-                    content = response.choices[0].message.content
-                    output = json.loads(content)
-
-                    # Validate output keys
-                    if 'message' not in output or 'header' not in output:
-                        print(f"'message' or 'header' missing in response (attempt {attempt+1}). Retrying...")
-                        continue
-
-                    # Check character length constraints
-                    if (len(output["header"].strip()) > self.Core.config_file["header_limit"] or
-                            len(output["message"].strip()) > self.Core.config_file["message_limit"]):
-                        print(f"Header or message exceeded character limits (attempt {attempt+1}). Retrying...")
-                        continue
-
-                    # If we're good here, update token usage
-                    self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    self.Core.temp_token_counter += tokens['total_tokens']
-
-                    return output
-
-                except json.JSONDecodeError:
-                    print(f"Invalid JSON from LLM (attempt {attempt+1}). Retrying...")
-
-            except openai.APIConnectionError as e:
-                print("The server could not be reached")
-                print(e.__cause__)
-            except openai.RateLimitError as e:
-                print("Received a 429 status code; backing off might be needed.")
-            except openai.APIStatusError as e:
-                print("A non-200 status code was received")
-                print(e.status_code)
-                print(e.response)
-
-        print("Max retries exceeded. Returning None.")
-        return None
-
     # --------------------------------------------------------------
 
     def llm_instructions(self):
@@ -375,7 +305,7 @@ Return only JSON of the form:
         jargon_list = "\n".join(f"- {word}" for word in self.Core.config_file["AI_Jargon"])
 
         instructions = f"""
-        You are texting a friend and it should feel like a friendly encouraging nudge. Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like natural everyday speech: friendly, concise, no jargon, and following the instructions.
+        You are a copywriter. Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like natural everyday speech: friendly, concise, no jargon, and following the instructions.
         Write a SUPER CASUAL and NATURAL push notification, as if you are chatting over coffee. Avoid odd phrasings. 
         
         ABSOLUTE RULE – OVERRIDES EVERYTHING ELSE:   

Messaging_system/Ollama.py

@@ -1,166 +0,0 @@
-import json
-import time
-
-import torch
-import ollama
-
-class LocalLM:
-
-    def __init__(self, model):
-        # Initialize the Ollama client
-        self.client = ollama.Client()
-        self.model = model
-
-    # def get_llm_response(self, prompt):
-    #
-    #     # Send the query to the model
-    #     response = self.client.generate(model=self.model, prompt=prompt)
-    #     return response.response
-
-    def preprocess_and_parse_json(self, response):
-
-        # Remove any leading/trailing whitespace and newlines
-        if response.startswith('```json') and response.endswith('```'):
-            response = response[len('```json'):-len('```')].strip()
-
-        # Parse the cleaned response into a JSON object
-        try:
-            json_object = json.loads(response)
-            return json_object
-        except json.JSONDecodeError as e:
-            print(f"Failed to parse JSON: {e}")
-            return None
-
-
-    def get_llm_response(self, prompt, max_retries=4):
-        """
-        sending the prompt to the LLM and get back the response
-        """
-
-        openai.api_key = self.Core.api_key
-        instructions = self.llm_instructions()
-        client = OpenAI(api_key=self.Core.api_key)
-
-        for attempt in range(max_retries):
-            try:
-                response = client.chat.completions.create(
-                    model=self.Core.model,
-                    response_format={"type": "json_object"},
-                    messages=[
-                        {"role": "system", "content": instructions},
-                        {"role": "user", "content": prompt}
-                    ],
-                    max_tokens=500,
-                    n=1,
-                    # temperature=self.Core.temperature,
-                    temperature=0.7,
-                )
-
-                tokens = {
-                    'prompt_tokens': response.usage.prompt_tokens,
-                    'completion_tokens': response.usage.completion_tokens,
-                    'total_tokens': response.usage.total_tokens
-                }
-
-                try:
-                    content = response.choices[0].message.content
-
-                    # Extract JSON code block
-
-                    output = json.loads(content)
-                    # output = json.loads(response.choices[0].message.content)
-
-                    if 'message' not in output or 'header' not in output:
-                        print(f"'message' or 'header' is missing in response on attempt {attempt + 1}. Retrying...")
-                        continue  # Continue to next attempt
-
-                    else:
-                        if len(output["header"].strip()) > self.Core.config_file["header_limit"] or len(
-                                output["message"].strip()) > self.Core.config_file["message_limit"]:
-                            print(
-                                f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
-                            continue
-
-                    # validating the JSON
-                    self.Core.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.Core.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    self.Core.temp_token_counter += tokens['total_tokens']
-                    return output
-
-                except json.JSONDecodeError:
-                    print(f"Invalid JSON from LLM on attempt {attempt + 1}. Retrying...")
-
-            except openai.APIConnectionError as e:
-                print("The server could not be reached")
-                print(e.__cause__)  # an underlying Exception, likely raised within httpx.
-            except openai.RateLimitError as e:
-                print("A 429 status code was received; we should back off a bit.")
-            except openai.APIStatusError as e:
-                print("Another non-200-range status code was received")
-                print(e.status_code)
-                print(e.response)
-
-        print("Max retries exceeded. Returning empty response.")
-        return None
-
-    def get_llm_response(self, prompt, mode, max_retries=10):
-        """
-        Send the prompt to the LLM and get back the response.
-        Includes handling for GPU memory issues by clearing cache and waiting before retry.
-        """
-
-        for attempt in range(max_retries):
-            try:
-                # Try generating the response
-                response = self.client.generate(model=self.model, prompt=prompt)
-            except Exception as e:
-                # This catches errors like the connection being forcibly closed
-                print(f"Error on attempt {attempt + 1}: {e}.")
-                try:
-                    # Clear GPU cache if you're using PyTorch; this may help free up memory
-                    torch.cuda.empty_cache()
-                    print("Cleared GPU cache.")
-                except Exception as cache_err:
-                    print("Failed to clear GPU cache:", cache_err)
-                # Wait a bit before retrying to allow memory to recover
-                time.sleep(2)
-                continue
-
-            try:
-                tokens = {
-                    'prompt_tokens': 0,
-                    'completion_tokens': 0,
-                    'total_tokens': 0
-                }
-
-                try:
-                    output = self.preprocess_and_parse_json(response.response)
-                    if output is None:
-                        continue
-
-                    if mode == "rating":
-                        # Check if all keys and values are integers (or convertible to integers)
-                        all_int = True
-                        for k, v in output.items():
-                            try:
-                                int(k)
-                                int(v)
-                            except ValueError:
-                                all_int = False
-                                break
-                        if all_int:
-                            return output, tokens
-                        else:
-                            print(f"Keys and values are not integers on attempt {attempt + 1}. Retrying...")
-                            continue  # Continue to next attempt
-                    else:
-                        print(f"Invalid mode: {mode}")
-                        return None, tokens
-
-                except json.JSONDecodeError:
-                    print(f"Invalid JSON from LLM on attempt {attempt + 1}. Retrying...")
-            except Exception as parse_error:
-                print("Error processing output:", parse_error)
-
-        print("Max retries exceeded. Returning empty response.")
-        return [], {}

Messaging_system/Permes.py

@@ -24,7 +24,7 @@ class Permes:
     def create_personalize_messages(self, session, users, brand, config_file, openai_api_key, CTA, segment_info,
                                     platform="push", number_of_messages=1, instructionset=None, subsequent_examples=None,
                                     message_style=None, selected_input_features=None, selected_source_features=None
-                                    , recsys_contents=None,
+                                    , recsys_contents=None, model=None,
                                     additional_instructions=None, identifier_column="user_id",
                                     sample_example=None, number_of_samples=None, involve_recsys_result=False,
                                     messaging_mode="message", target_column=None, ongoing_df=None,
@@ -83,6 +83,9 @@ class Permes:
         if selected_source_features is not None:
             personalize_message.set_features_to_use(selected_source_features)
 
+        if model is not None:
+            personalize_message.set_llm_model(selected_source_features)
+
         # if involve_recsys_result is not None:
         #     personalize_message.set_messaging_mode("recsys_result")
 

Messaging_system/PromptGenerator.py

@@ -109,7 +109,7 @@ class PromptGenerator:
         """
 
         context = f"""
-        You are texting a friend and it should feel like a friendly encouraging nudge. Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like everyday natural speech: friendly, short, no jargon, and following the instructions.
+        Your task is to write a 'header' and a 'message' as a push notification for a {self.Core.get_instrument()} student that sounds like everyday natural speech: friendly, short, no jargon, and following the instructions.
         """
 
         return context

Messaging_system/protection_layer.py

@@ -7,7 +7,7 @@ import os
 import openai
 from openai import OpenAI
 from dotenv import load_dotenv
-load_dotenv()
+from Messaging_system.LLM import LLM
 
 
 # -----------------------------------------------------------------------
@@ -17,16 +17,11 @@ class ProtectionLayer:
     Protection layer to double check the generated message:
     """
 
-    def __init__(self, config_file, messaging_mode):
+    def __init__(self, CoreConfig):
 
-        self.config_file = config_file
-        self.messaging_mode = messaging_mode
-
-        # LLM configs
-        self.api_key = os.environ.get("OPENAI_API")  # will be set by user
-        self.model = "gpt-4o-mini"  # will be set by user
-        self.temperature = 0
+        self.Core = CoreConfig
 
+        self.llm = LLM(CoreConfig)
         # to trace the number of tokens and estimate the cost if needed
         self.total_tokens = {
             'prompt_tokens': 0,
@@ -39,11 +34,11 @@ class ProtectionLayer:
         Setting instructions for the LLM for the second pass.
         """
 
-        jargon_list = "\n".join(f"- {word}" for word in self.config_file["AI_Jargon"])
+        jargon_list = "\n".join(f"- {word}" for word in self.Core.config_file["AI_Jargon"])
 
         instructions = f"""
 
-        You are texting a friend that should feels like a friendly nudge. Your job is to *either* approve the candidate message or return a corrected version that obeys the style guide.
+        You are friendly copywriter. Your job is to *either* approve the candidate message or return a corrected version that obeys the style guide.
         the 'header' and a 'message' as a push notification should sounds like everyday speech—friendly, short, no jargon, following the instructions.
         
         ABSOLUTE RULE – OVERRIDES EVERYTHING ELSE: 
@@ -108,72 +103,6 @@ class ProtectionLayer:
         return instructions
 
     # --------------------------------------------------------------
-    def get_llm_response(self, prompt, max_retries=3):
-        """
-        sending the prompt to the LLM and get back the response
-        """
-
-        openai.api_key = self.api_key
-        instructions = self.llm_instructions()
-        client = OpenAI(api_key=self.api_key)
-
-        for attempt in range(max_retries):
-            try:
-                response = client.chat.completions.create(
-                    model=self.model,
-                    response_format={"type": "json_object"},
-                    messages=[
-                        {"role": "system", "content": instructions},
-                        {"role": "user", "content": prompt}
-                    ],
-                    max_tokens=500,
-                    n=1,
-                    temperature=self.temperature
-                )
-
-                tokens = {
-                    'prompt_tokens': response.usage.prompt_tokens,
-                    'completion_tokens': response.usage.completion_tokens,
-                    'total_tokens': response.usage.total_tokens
-                }
-
-                try:
-                    content = response.choices[0].message.content
-                    # Extract JSON code block
-
-                    output = json.loads(content)
-                    # output = json.loads(response.choices[0].message.content)
-
-                    if 'message' not in output or 'header' not in output:
-                        print(f"'message' or 'header' is missing in response on attempt {attempt + 1}. Retrying...")
-                        continue  # Continue to next attempt
-
-                    else:
-                        if len(output["header"].strip()) > self.config_file["header_limit"] or len(output["message"].strip()) > self.config_file["message_limit"]:
-                            print(f"'header' or 'message' is more than specified characters in response on attempt {attempt + 1}. Retrying...")
-                            continue
-
-                    # validating the JSON
-                    self.total_tokens['prompt_tokens'] += tokens['prompt_tokens']
-                    self.total_tokens['completion_tokens'] += tokens['completion_tokens']
-                    return output
-
-                except json.JSONDecodeError:
-                    print(f"Invalid JSON from LLM on attempt {attempt + 1}. Retrying...")
-
-            except openai.APIConnectionError as e:
-                print("The server could not be reached")
-                print(e.__cause__)  # an underlying Exception, likely raised within httpx.
-            except openai.RateLimitError as e:
-                print("A 429 status code was received; we should back off a bit.")
-            except openai.APIStatusError as e:
-                print("Another non-200-range status code was received")
-                print(e.status_code)
-                print(e.response)
-
-        print("Max retries exceeded. Returning empty response.")
-        return [], {}
-
     # --------------------------------------------------------------
     def get_context(self):
         """
@@ -196,7 +125,7 @@ class ProtectionLayer:
         :return: new prompt
         """
         # recommended_content = ""
-        # if self.messaging_mode == "recsys_result":
+        # if self.Core.messaging_mode == "recsys_result":
         #     recommended_content = f"""
         # ### ** Recommended Content **
         # {user['recommendation_info']}
@@ -229,7 +158,7 @@ class ProtectionLayer:
         """
 
         prompt = self.generate_prompt(message, user)
-        response = self.get_llm_response(prompt)
+        response = self.llm.get_response(prompt=prompt, instructions=self.llm_instructions())
 
         return response, self.total_tokens
 

app.py

@@ -81,7 +81,7 @@ def init_state() -> None:
         valid_instructions="",
         invalid_instructions="",
         messaging_type="push",
-        generated=False,
+        generated=False, model=None,
         include_recommendation=False,
         data=None, brand=None, recsys_contents=[], csv_output=None,
         users_message=None, messaging_mode=None, target_column=None,
@@ -90,7 +90,7 @@ def init_state() -> None:
         additional_instructions=None, segment_info="", message_style="",
         sample_example="", CTA="", all_features=None, number_of_messages=1,
         instructionset={}, subsequent_examples={}, segment_name="", number_of_samples=10,
-        selected_source_features=[], platform=None, generate_clicked=False,
+        selected_source_features=[], platform=None, generate_clicked=False, config=None,
     )
     for k, v in defaults.items():
         st.session_state.setdefault(k, v)
@@ -163,6 +163,13 @@ with st.sidebar:
             key="brand",
         )
 
+        # ─ Brand
+        st.selectbox(
+            "LLM model *",
+            st.session_state.config["LLM_models"],
+            key="model",
+        )
+
         # ─ Personalisation
         st.text_area("Segment info *", key="segment_info")
         st.text_area("CTA (Call to Action) *", key="CTA")
@@ -258,7 +265,6 @@ with tab2:
                 warehouse=get_credential("snowflake_warehouse"),
                 schema=get_credential("snowflake_schema")
             )
-            config = load_config_("Config_files/message_system_config.json")
             session = Session.builder.configs(conn).create()
 
             # ─ prepare parameters
@@ -295,8 +301,9 @@ with tab2:
                 session=session,
                 users=st.session_state.data,
                 brand=st.session_state.brand,
-                config_file=config,
+                config_file=st.session_state.config,
                 openai_api_key=get_credential("OPENAI_API"),
+                model=st.session_state.model,
                 CTA=st.session_state.CTA,
                 segment_info=st.session_state.segment_info,
                 number_of_samples=st.session_state.number_of_samples,

messaging_main_test.py

@@ -114,13 +114,16 @@ if __name__ == "__main__":
     brand = "singeo"
     identifier_column = "user_id"
 
-    segment_info = """This is a singeo user who didn't practiced for a while."""
+    segment_info = """Student who haven't practiced for a few days"""
 
     # sample inputs
 
-    CTA = """Re-engage the user by reminding them of their goals, and by recommending content that will get them back on track. """
+    CTA = """The goal is to tell them to practice singing"""
 
-    # sample_example = """we have crafted a perfect set of courses just for you! come and check it out!"""
+    sample_example = """
+    Header: Sing Your Heart Out 
+    Message: It’s been a few days. Take a lesson today and start practicing!
+"""
 
     # additional_instructions = """Include weeks_since _last_interaction in the message if you can create a better message to re-engage the user."""
     additional_instructions = None
@@ -133,14 +136,14 @@ if __name__ == "__main__":
     # number of messages to generate
     number_of_messages = 3
     instructionset = {
-        1: "Be highly motivational positive and kind",
-        2: "Be highly motivational positive and kind",
-        3: "Be highly motivational positive and kind",
+        1: "Talk like a singing coach motivating your student. Don't say things a singer wouldn't say. Make the message quick, concise, and in casual language. Tell them to practice, take a lesson, or warm up today. ",
+        2: "Talk like a singing coach motivating your student. Don't say things a singer wouldn't say. Make the message quick, concise, and in casual language. Tell them to practice, take a lesson, or warm up today. ",
+        3: "Talk like a singing coach motivating your student. Don't say things a singer wouldn't say. Make the message quick, concise, and in casual language. Tell them to practice, take a lesson, or warm up today. ",
     }
 
     subsequent_examples = {
-        1: "header: Don’t Stop Believin’"
-           "message: You're closer to your goals than you think! Start singing now.",
+        1: "Header: Sing Your Heart Out!"
+           "Message: It’s been a few days. Take a lesson today and start practicing!",
         2: "header: Here Comes The Sun"
            "message: A quick practice session will light up your day. Let’s get right back at it. ",
         3: "header: Ain’t No Mountain High Enough"
@@ -152,21 +155,10 @@ if __name__ == "__main__":
 
     # messaging_mode = "recommend_playlist"
 
-    sample_example = """
-    Below are sample messages from us. make the generated message close to our sound in terms of style, tune, and the way we write messages. 
-
-
-    Example 1
-    header: Your voice Misses You, [first_name]
-    message: It’s been a while. Jump back in with this quick lesson! 
-    
-    """
-
-    # sample_example = None
 
     platform = "push"
 
-    selected_source_features = ["first_name", "weeks_since_last_interaction"]
+    selected_source_features = None
     selected_input_features = None
 
     segment_name = "no_recent_activity"