doc : add comments to methods in module_llm.py

Author: colinfrisch

mesa_llm/module_llm.py

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