refactor(skills): move skills to ee

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,14 @@ def __init__(
                 stacklevel=2,
             )
 
+        # Transition pd dataframe to pandasai dataframe
+        if isinstance(dfs, list):
+            for df in dfs:
+                if isinstance(df, pd.DataFrame):
+                    df = DataFrame(df)
+        elif isinstance(dfs, pd.DataFrame):
+            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):

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/skills/__init__.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: