sys_prompt

Author: sanika-n

mesa_llm/llm_agent.py

@@ -2,10 +2,14 @@
 
 
 class LLMAgent:
-    def __init__(self, api_key: str, model: str = "openai/gpt-4o"):
+    def __init__(self, api_key: str, model: str = "openai/gpt-4o", system_prompt: str|None = None):
         """Initialize the LLMAgent with a ModuleLLM instance."""
-        self.llm = ModuleLLM(api_key=api_key, model=model)
+        self.llm = ModuleLLM(api_key=api_key, model=model, system_prompt=system_prompt)
 
-    def set_llm(self, api_key: str, model: str = "openai/gpt-4o") -> None:
-        """Replace the current LLM with a new configuration."""
-        self.llm = ModuleLLM(api_key=api_key, model=model)
+    def set_model(self, api_key: str, model: str = "openai/gpt-4o") -> None:
+        """Set the model of the LLM."""
+        self.llm = ModuleLLM(api_key=api_key, model=model, system_prompt=self.llm.system_prompt)
+        
+    def set_system_prompt(self, system_prompt: str) -> None:
+        """Set the system prompt for the LLM."""
+        self.llm= ModuleLLM(api_key=self.llm.api_key, model=self.llm.model, system_prompt=system_prompt)

mesa_llm/module_llm.py

@@ -10,32 +10,30 @@ class ModuleLLM:
     Note : Currently supports OpenAI, Anthropic, xAI, Huggingface, Ollama, OpenRouter, NovitaAI
     """
 
-    def __init__(self, api_key: str, model: str):
+    def __init__(self, api_key: str, model: str, system_prompt: str | None = None):
         """
         Initialize the LLM module
 
         Args:
             api_key: The API key for the LLM provider
-            model: The model to use for the LLM
+            model: The model to use for the LLM in the format of {provider}/{model}
+            system_prompt: The system prompt to use for the LLM
         """
         self.api_key = api_key
-        provider = model.split("/")[0].upper()
-
+        self.model=model
+        self.system_prompt = system_prompt
+        provider = self.model.split("/")[0].upper()
         os.environ[f"{provider}_API_KEY"] = self.api_key
 
-    def generate(self, prompt: str, system_prompt: str | None = None) -> str:
-        """
-        Generate a response from the LLM
-        """
-        if system_prompt:
+    def generate(self, prompt: str) -> str:
+        if self.system_prompt:
             messages = [
-                {"role": "system", "content": system_prompt},
+                {"role": "system", "content": self.system_prompt},
                 {"role": "user", "content": prompt},
             ]
         else:
             messages = [{"role": "user", "content": prompt}]
-
-        response = completion(model="openai/gpt-4o", messages=messages)
+        response = completion(model=self.model, messages=messages)
         return response