Add docstrings (sinaptik-ai#20)

Add docstrings

Author: gventuri

pandasai/exceptions.py

@@ -2,20 +2,45 @@
 
 
 class APIKeyNotFoundError(Exception):
-    """Raised when the API key is not defined/declared."""
+    """
+    Raised when the API key is not defined/declared.
+
+    Args:
+        Exception (Exception): APIKeyNotFoundError
+    """
 
 
 class LLMNotFoundError(Exception):
-    """Raised when the LLM is not provided"""
+    """
+    Raised when the LLM is not provided.
+
+    Args:
+        Exception (Exception): LLMNotFoundError
+    """
 
 
 class NoCodeFoundError(Exception):
-    """Raised when no code is found in the response"""
+    """
+    Raised when no code is found in the response.
+
+    Args:
+        Exception (Exception): NoCodeFoundError
+    """
 
 
 class MethodNotImplementedError(Exception):
-    """Raised when a method is not implemented"""
+    """
+    Raised when a method is not implemented.
+
+    Args:
+        Exception (Exception): MethodNotImplementedError
+    """
 
 
 class UnsupportedOpenAIModelError(Exception):
-    """Raised when an unsupported OpenAI model is used"""
+    """
+    Raised when an unsupported OpenAI model is used.
+
+    Args:
+        Exception (Exception): UnsupportedOpenAIModelError
+    """

pandasai/llm/base.py

@@ -17,7 +17,15 @@ class LLM:
 
     @property
     def type(self) -> str:
-        """Return type of llm."""
+        """
+        Return type of LLM.
+
+        Raises:
+            APIKeyNotFoundError: Type has not been implemented
+
+        Returns:
+            str: Type of LLM a string
+        """
         raise APIKeyNotFoundError("Type has not been implemented")
 
     def _remove_imports(self, code: str) -> str:
@@ -32,12 +40,16 @@ def _remove_imports(self, code: str) -> str:
         return astor.to_source(new_tree)
 
     def _polish_code(self, code: str) -> str:
-        """Polish the code:
-        - remove the leading "python" or "py"
-        - remove the imports
-        - remove trailing spaces and new lines
         """
+        Polish the code by removing the leading "python" or "py",  \
+        removing the imports and removing trailing spaces and new lines.
 
+        Args:
+            code (str): Code
+
+        Returns:
+            str: Polished code
+        """
         if re.match(r"^(python|py)", code):
             code = re.sub(r"^(python|py)", "", code)
         if re.match(r"^`.*`$", code):
@@ -54,8 +66,19 @@ def _is_python_code(self, string):
             return False
 
     def _extract_code(self, response: str, separator: str = "```") -> str:
-        """Extract the code from the response"""
+        """
+        Extract the code from the response.
+
+        Args:
+            response (str): Response
+            separator (str, optional): Separator. Defaults to "```".
 
+        Raises:
+            NoCodeFoundError: No code found in the response
+
+        Returns:
+            str: Extracted code from the response
+        """
         code = response
         if len(response.split(separator)) > 1:
             code = response.split(separator)[1]
@@ -67,12 +90,24 @@ def _extract_code(self, response: str, separator: str = "```") -> str:
 
         return code
 
-    def call(self, instruction: str, value: str) -> str:
-        """Execute the llm with the given prompt"""
+    def call(self, instruction: str, value: str) -> None:
+        """
+        Execute the LLM with given prompt.
+
+        Args:
+            instruction (str): Prompt
+            value (str): Value
 
+        Raises:
+            MethodNotImplementedError: Call method has not been implemented
+        """
         raise MethodNotImplementedError("Call method has not been implemented")
 
     def generate_code(self, instruction: str, prompt: str) -> str:
-        """Generate the code based on the instruction and the prompt"""
+        """
+        Generate the code based on the instruction and the given prompt.
 
+        Returns:
+            str: Code
+        """
         return self._extract_code(self.call(instruction, prompt))

pandasai/llm/openai.py

@@ -58,8 +58,15 @@ def _set_params(self, **kwargs):
                 setattr(self, key, value)
 
     def completion(self, prompt: str) -> str:
-        """Query the completion API"""
+        """
+        Query the completation API
 
+        Args:
+            prompt (str): Prompt
+
+        Returns:
+            str: LLM response
+        """
         params = {
             "model": self.model,
             "prompt": prompt,
@@ -78,8 +85,15 @@ def completion(self, prompt: str) -> str:
         return response["choices"][0]["text"]
 
     def chat_completion(self, value: str) -> str:
-        """Query the chat completion API"""
+        """
+        Query the chat completion API
+
+        Args:
+            value (str): Prompt
 
+        Returns:
+            str: LLM response
+        """
         params = {
             "model": self.model,
             "temperature": self.temperature,
@@ -103,7 +117,19 @@ def chat_completion(self, value: str) -> str:
         return response["choices"][0]["message"]["content"]
 
     def call(self, instruction: str, value: str) -> str:
-        """Call the openAI LLM"""
+        """
+        Call the OpenAI LLM.
+
+        Args:
+            instruction (str): Instruction to pass
+            value (str): Value to pass
+
+        Raises:
+            UnsupportedOpenAIModelError: Unsupported model
+
+        Returns:
+            str: Response
+        """
         self.last_prompt = str(instruction) + str(value)
 
         if self.model in self._supported_completion_models: