Merge pull request #1831 from sinaptik-ai/move_skills_to_ee

refactor move skills to ee folder

Author: ArslanSaleem

pandasai/__init__.py

@@ -2,29 +2,31 @@
 """
 PandasAI is a wrapper around a LLM to make dataframes conversational
 """
+from __future__ import annotations
 
 import os
 from io import BytesIO
 from typing import Hashable, List, Optional, Union
 
 import pandas as pd
 
-from pandasai.config import APIKeyManager, ConfigManager, SkillsManager
+from pandasai.config import APIKeyManager, ConfigManager
 from pandasai.data_loader.semantic_layer_schema import (
     Column,
     Relation,
     SemanticLayerSchema,
     Source,
     Transformation,
 )
+from pandasai.ee.skills import skill
+from pandasai.ee.skills.manager import SkillsManager
 from pandasai.exceptions import DatasetNotFound, InvalidConfigError
 from pandasai.helpers.path import (
     find_project_root,
     get_validated_dataset_path,
     transform_dash_to_underscore,
 )
 from pandasai.sandbox.sandbox import Sandbox
-from pandasai.skills import skill
 
 from .agent import Agent
 from .data_loader.loader import DatasetLoader

pandasai/agent/base.py

@@ -65,6 +65,12 @@ def __init__(
                 stacklevel=2,
             )
 
+        # Transition pd dataframe to pandasai dataframe
+        if isinstance(dfs, list):
+            dfs = [DataFrame(df) if self.is_pd_dataframe(df) else df for df in dfs]
+        elif self.is_pd_dataframe(dfs):
+            dfs = DataFrame(dfs)
+
         if isinstance(dfs, list):
             sources = [df.schema.source or df._loader.source for df in dfs]
             if not BaseQueryBuilder.check_compatible_sources(sources):
@@ -80,6 +86,9 @@ def __init__(
         self._response_parser = ResponseParser()
         self._sandbox = sandbox
 
+    def is_pd_dataframe(self, df: Union[DataFrame, VirtualDataFrame]) -> bool:
+        return not isinstance(df, DataFrame) and isinstance(df, pd.DataFrame)
+
     def chat(self, query: str, output_type: Optional[str] = None):
         """
         Start a new chat interaction with the assistant on Dataframe.

pandasai/agent/state.py

@@ -5,9 +5,10 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from pandasai.config import Config, ConfigManager, SkillsManager
+from pandasai.config import Config, ConfigManager
 from pandasai.constants import DEFAULT_CHART_DIRECTORY
 from pandasai.data_loader.semantic_layer_schema import is_schema_source_same
+from pandasai.ee.skills.manager import SkillsManager
 from pandasai.exceptions import InvalidConfigError
 from pandasai.helpers.folder import Folder
 from pandasai.helpers.logger import Logger

pandasai/config.py

@@ -1,12 +1,10 @@
 import os
-from importlib.util import find_spec
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, ConfigDict
 
 from pandasai.helpers.filemanager import DefaultFileManager, FileManager
 from pandasai.llm.base import LLM
-from pandasai.skills import Skill
 
 
 class Config(BaseModel):
@@ -59,88 +57,3 @@ def set(cls, api_key: str):
     @classmethod
     def get(cls) -> Optional[str]:
         return cls._api_key
-
-
-class SkillsManager:
-    """
-    A singleton class to manage the global skills list.
-    """
-
-    _skills: List[Skill] = []
-
-    @classmethod
-    def add_skills(cls, *skills: Skill):
-        """
-        Add skills to the global list of skills. If a skill with the same name
-             already exists, raise an error.
-
-        Args:
-            *skills: Variable number of skill objects to add.
-        """
-        for skill in skills:
-            if any(existing_skill.name == skill.name for existing_skill in cls._skills):
-                raise ValueError(f"Skill with name '{skill.name}' already exists.")
-
-        cls._skills.extend(skills)
-
-    @classmethod
-    def skill_exists(cls, name: str):
-        """
-        Check if a skill with the given name exists in the global list of skills.
-
-        Args:
-            name (str): The name of the skill to check.
-
-        Returns:
-            bool: True if a skill with the given name exists, False otherwise.
-        """
-        return any(skill.name == name for skill in cls._skills)
-
-    @classmethod
-    def has_skills(cls):
-        """
-        Check if there are any skills in the global list of skills.
-
-        Returns:
-            bool: True if there are skills, False otherwise.
-        """
-        return len(cls._skills) > 0
-
-    @classmethod
-    def get_skill_by_func_name(cls, name: str):
-        """
-        Get a skill by its name from the global list.
-
-        Args:
-            name (str): The name of the skill to retrieve.
-
-        Returns:
-            Skill or None: The skill with the given name, or None if not found.
-        """
-        return next((skill for skill in cls._skills if skill.name == name), None)
-
-    @classmethod
-    def get_skills(cls) -> List[Skill]:
-        """
-        Get the global list of skills.
-
-        Returns:
-            List[Skill]: The list of all skills.
-        """
-        return cls._skills.copy()
-
-    @classmethod
-    def clear_skills(cls):
-        """
-        Clear all skills from the global list.
-        """
-        cls._skills.clear()
-
-    @classmethod
-    def __str__(cls) -> str:
-        """
-        Present all skills
-        Returns:
-            str: String representation of all skills
-        """
-        return "\n".join(str(skill) for skill in cls._skills)

pandasai/ee/LICENSE

@@ -0,0 +1,36 @@
+The PandasAI Enterprise license (the “Enterprise License”)
+Copyright (c) 2024 Sinaptik GmbH
+
+With regard to the PandasAI Software:
+
+This software and associated documentation files (the "Software") may only be
+used in production, if you (and any entity that you represent) have agreed to,
+and are in compliance with, the PandasAI Subscription Terms of Service, available
+at https://pandas-ai.com/terms (the “Enterprise Terms”), or other
+agreement governing the use of the Software, as agreed by you and PandasAI,
+and otherwise have a valid PandasAI Enterprise license for the
+correct number of user seats. Subject to the foregoing sentence, you are free to
+modify this Software and publish patches to the Software. You agree that PandasAI
+and/or its licensors (as applicable) retain all right, title and interest in and
+to all such modifications and/or patches, and all such modifications and/or
+patches may only be used, copied, modified, displayed, distributed, or otherwise
+exploited with a valid PandasAI Enterprise license for the correct
+number of user seats. Notwithstanding the foregoing, you may copy and modify
+the Software for development and testing purposes, without requiring a
+subscription. You agree that PandasAI and/or its licensors (as applicable) retain
+all right, title and interest in and to all such modifications. You are not
+granted any other rights beyond what is expressly stated herein. Subject to the
+foregoing, it is forbidden to copy, merge, publish, distribute, sublicense,
+and/or sell the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+For all third party components incorporated into the PandasAI Software, those
+components are licensed under the original license provided by the owner of the
+applicable component.

pandasai/skills/__init__.py renamed to pandasai/ee/skills/__init__.py

@@ -4,7 +4,7 @@
 from pydantic import BaseModel, PrivateAttr
 
 
-class Skill(BaseModel):
+class SkillType(BaseModel):
     """Skill that takes a function usable by pandasai"""
 
     func: Callable[..., Any]
@@ -41,7 +41,7 @@ def __init__(
             )
         signature = f"def {name}{inspect.signature(func)}:"
 
-        super(Skill, self).__init__(
+        super(SkillType, self).__init__(
             func=func, description=description, name=name, **kwargs
         )
         self._signature = signature
@@ -51,7 +51,7 @@ def __call__(self, *args, **kwargs) -> Any:
         return self.func(*args, **kwargs)
 
     @classmethod
-    def from_function(cls, func: Callable, **kwargs: Any) -> "Skill":
+    def from_function(cls, func: Callable, **kwargs: Any) -> "SkillType":
         """
         Creates a skill object from a function
 
@@ -95,8 +95,8 @@ def compute_flight_prices(offers: pd.Dataframe) -> List[float]:
     """
 
     def _make_skill_with_name(skill_name: str) -> Callable:
-        def _make_skill(skill_fn: Callable) -> Skill:
-            skill_obj = Skill(
+        def _make_skill(skill_fn: Callable) -> SkillType:
+            skill_obj = SkillType(
                 name=skill_name,  # func.__name__ if None
                 # when this decorator is used, the function MUST have a docstring
                 description=skill_fn.__doc__,
@@ -105,7 +105,7 @@ def _make_skill(skill_fn: Callable) -> Skill:
 
             # Automatically add the skill to the global skills manager
             try:
-                from pandasai.config import SkillsManager
+                from pandasai.ee.skills.manager import SkillsManager
 
                 SkillsManager.add_skills(skill_obj)
             except ImportError:
@@ -125,11 +125,14 @@ def _make_skill(skill_fn: Callable) -> Skill:
     elif not args:
         # Covers the case in which a function is decorated with "@skill()"
         # with the intended behavior of "@skill"
-        def _func_wrapper(fn: Callable) -> Skill:
+        def _func_wrapper(fn: Callable) -> SkillType:
             return _make_skill_with_name(fn.__name__)(fn)
 
         return _func_wrapper
     else:
         raise ValueError(
             f"Too many arguments for skill decorator. Received: {len(args)}"
         )
+
+
+__all__ = ["skill", "SkillType"]

pandasai/ee/skills/manager.py

@@ -0,0 +1,88 @@
+from typing import List
+
+from pandasai.ee.skills import SkillType
+
+
+class SkillsManager:
+    """
+    A singleton class to manage the global skills list.
+    """
+
+    _skills: List[SkillType] = []
+
+    @classmethod
+    def add_skills(cls, *skills: SkillType):
+        """
+        Add skills to the global list of skills. If a skill with the same name
+             already exists, raise an error.
+
+        Args:
+            *skills: Variable number of skill objects to add.
+        """
+        for skill in skills:
+            if any(existing_skill.name == skill.name for existing_skill in cls._skills):
+                raise ValueError(f"Skill with name '{skill.name}' already exists.")
+
+        cls._skills.extend(skills)
+
+    @classmethod
+    def skill_exists(cls, name: str):
+        """
+        Check if a skill with the given name exists in the global list of skills.
+
+        Args:
+            name (str): The name of the skill to check.
+
+        Returns:
+            bool: True if a skill with the given name exists, False otherwise.
+        """
+        return any(skill.name == name for skill in cls._skills)
+
+    @classmethod
+    def has_skills(cls):
+        """
+        Check if there are any skills in the global list of skills.
+
+        Returns:
+            bool: True if there are skills, False otherwise.
+        """
+        return len(cls._skills) > 0
+
+    @classmethod
+    def get_skill_by_func_name(cls, name: str):
+        """
+        Get a skill by its name from the global list.
+
+        Args:
+            name (str): The name of the skill to retrieve.
+
+        Returns:
+            Skill or None: The skill with the given name, or None if not found.
+        """
+        return next((skill for skill in cls._skills if skill.name == name), None)
+
+    @classmethod
+    def get_skills(cls) -> List[SkillType]:
+        """
+        Get the global list of skills.
+
+        Returns:
+            List[SkillType]: The list of all skills.
+        """
+        return cls._skills.copy()
+
+    @classmethod
+    def clear_skills(cls):
+        """
+        Clear all skills from the global list.
+        """
+        cls._skills.clear()
+
+    @classmethod
+    def __str__(cls) -> str:
+        """
+        Present all skills
+        Returns:
+            str: String representation of all skills
+        """
+        return "\n".join(str(skill) for skill in cls._skills)

pandasai/smart_datalake/__init__.py

@@ -34,7 +34,11 @@ def load_dfs(self, dfs: List[pd.DataFrame]):
         load_dfs = []
         for df in dfs:
             if isinstance(df, pd.DataFrame):
-                load_dfs.append(DataFrame(df))
+                load_dfs.append(
+                    DataFrame(df)
+                    if not isinstance(df, DataFrame) and isinstance(df, pd.DataFrame)
+                    else df
+                )
             else:
                 raise ValueError(
                     "Invalid input data. We cannot convert it to a dataframe."

tests/unit_tests/skills/test_shared_template.py

@@ -8,8 +8,8 @@
 import pytest
 from jinja2 import Environment, FileSystemLoader
 
-from pandasai.config import SkillsManager
-from pandasai.skills import Skill, skill
+from pandasai.ee.skills import skill
+from pandasai.ee.skills.manager import SkillsManager
 
 
 class TestSharedTemplate: