feat(SemanticLayerSchema): Adding a SemanticLayerSchema class with logical and business validation

extensions/connectors/sql/tests/test_sql.py

@@ -1,13 +1,14 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
 
 # Assuming the functions are in a module called db_loader
 from pandasai_sql import (
+    load_from_cockroachdb,
     load_from_mysql,
     load_from_postgres,
     load_from_sqlite,
-    load_from_cockroachdb,
 )
 
 

extensions/connectors/yfinance/tests/test_yahoo_finance.py

@@ -1,5 +1,6 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
 
 # Assuming the functions are in a module called yahoo_finance

extensions/ee/connectors/bigquery/pandasai_bigquery/__init__.py

@@ -1,5 +1,5 @@
-from google.cloud import bigquery
 import pandas as pd
+from google.cloud import bigquery
 
 
 def load_from_bigquery(connection_info, query):

extensions/ee/connectors/bigquery/tests/test_bigquery.py

@@ -1,6 +1,7 @@
-import pytest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
+import pytest
 from pandasai_bigquery import load_from_bigquery
 
 

extensions/ee/connectors/databricks/tests/test_databricks.py

@@ -1,5 +1,6 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 from pandasai_databricks import (
     load_from_databricks,
 )

extensions/ee/connectors/oracle/pandasai_oracle/__init__.py

@@ -1,5 +1,5 @@
-import pandas as pd
 import cx_Oracle
+import pandas as pd
 
 
 def load_from_oracle(connection_info, query):

extensions/ee/connectors/oracle/tests/test_oracle.py

@@ -1,5 +1,6 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
 from pandasai_oracle import load_from_oracle
 

extensions/ee/connectors/snowflake/tests/test_snowflake.py

@@ -1,5 +1,6 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
 from pandasai_snowflake import load_from_snowflake
 

extensions/ee/vectorstores/chromadb/tests/test_chromadb.py

@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
 
 from extensions.ee.vectorstores.chromadb.pandasai_chromadb import ChromaDB
 

extensions/ee/vectorstores/pinecone/pandasai_pinecone/pinecone.py

@@ -1,7 +1,8 @@
+import uuid
 from typing import Any, Callable, Iterable, List, Optional, Union
 
 import pinecone
-import uuid
+
 from pandasai.helpers.logger import Logger
 from pandasai.vectorstores.vectorstore import VectorStore
 

extensions/ee/vectorstores/qdrant/tests/test_qdrant.py

@@ -1,12 +1,12 @@
 import unittest
 import uuid
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
 
 from qdrant_client import models
 
 from extensions.ee.vectorstores.qdrant.pandasai_qdrant.qdrant import (
-    Qdrant,
     UUID_NAMESPACE,
+    Qdrant,
 )
 
 

extensions/llms/bedrock/pandasai_bedrock/__init__.py

@@ -1,4 +1,3 @@
 from .claude import BedrockClaude
 
-
 __all__ = ["BedrockClaude"]

extensions/llms/bedrock/pandasai_bedrock/base.py

@@ -4,14 +4,12 @@
 from typing import TYPE_CHECKING, Optional
 
 from pandasai.core.prompts.base import BasePrompt
-from pandasai.helpers.memory import Memory
-
 from pandasai.exceptions import (
     MethodNotImplementedError,
 )
+from pandasai.helpers.memory import Memory
 from pandasai.llm.base import LLM
 
-
 if TYPE_CHECKING:
     from pandasai.agent.state import AgentState
 

extensions/llms/bedrock/tests/test_bedrock_claude.py

@@ -6,9 +6,9 @@
 
 import pytest
 
+from extensions.llms.bedrock.pandasai_bedrock.claude import BedrockClaude
 from pandasai.core.prompts.base import BasePrompt
 from pandasai.exceptions import APIKeyNotFoundError, UnsupportedModelError
-from extensions.llms.bedrock.pandasai_bedrock.claude import BedrockClaude
 
 
 class MockBedrockRuntimeClient:

extensions/llms/google/pandasai_google/__init__.py

@@ -1,4 +1,4 @@
 from .google_gemini import GoogleGemini
 from .google_vertexai import GoogleVertexAI
 
-__all__ = ["GoogleGemini", "GoogleVertexAI"]
+__all__ = ["GoogleGemini", "GoogleVertexAI"]

extensions/llms/google/pandasai_google/base.py

@@ -4,14 +4,12 @@
 from typing import TYPE_CHECKING, Optional
 
 from pandasai.core.prompts.base import BasePrompt
-from pandasai.helpers.memory import Memory
-
 from pandasai.exceptions import (
     MethodNotImplementedError,
 )
+from pandasai.helpers.memory import Memory
 from pandasai.llm.base import LLM
 
-
 if TYPE_CHECKING:
     from pandasai.agent.state import AgentState
 

extensions/llms/google/pandasai_google/google_gemini.py

@@ -1,8 +1,9 @@
 from typing import Any, Optional
 
+from pandasai.core.code_execution.environment import import_dependency
 from pandasai.exceptions import APIKeyNotFoundError
 from pandasai.helpers.memory import Memory
-from pandasai.core.code_execution.environment import import_dependency
+
 from .base import BaseGoogle
 
 

extensions/llms/google/pandasai_google/google_vertexai.py

@@ -1,9 +1,9 @@
 from typing import Optional
 
+from pandasai.core.code_execution.environment import import_dependency
+from pandasai.exceptions import UnsupportedModelError
 from pandasai.helpers.memory import Memory
 
-from pandasai.exceptions import UnsupportedModelError
-from pandasai.core.code_execution.environment import import_dependency
 from .base import BaseGoogle
 
 

extensions/llms/google/tests/test_google_vertexai.py

@@ -1,9 +1,9 @@
 """Unit tests for the openai LLM class"""
 
 import pytest
+from pandasai_google.google_vertexai import GoogleVertexAI
 
 from pandasai.exceptions import UnsupportedModelError
-from pandasai_google.google_vertexai import GoogleVertexAI
 
 
 class MockedCompletion:

extensions/llms/huggingface/pandasai_huggingface/huggingface_text_gen.py

@@ -2,10 +2,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
-
 from pandasai.core.prompts.base import BasePrompt
 from pandasai.helpers import load_dotenv
-
 from pandasai.llm.base import LLM
 
 if TYPE_CHECKING:

extensions/llms/huggingface/tests/test_huggingface_text_gen.py

@@ -3,6 +3,7 @@
 from pandasai_huggingface.huggingface_text_gen import (
     HuggingFaceTextGen,
 )
+
 from pandasai.core.prompts.base import BasePrompt
 
 

extensions/llms/ibm/pandasai_ibm/ibm_watsonx.py

@@ -7,7 +7,6 @@
 from pandasai.core.prompts.base import BasePrompt
 from pandasai.exceptions import APIKeyNotFoundError
 from pandasai.helpers import load_dotenv
-
 from pandasai.llm.base import LLM
 
 if TYPE_CHECKING:

extensions/llms/ibm/tests/test_ibm_watsonx.py

@@ -3,8 +3,8 @@
 
 import pytest
 
-from pandasai.exceptions import APIKeyNotFoundError
 from extensions.llms.ibm.pandasai_ibm.ibm_watsonx import IBMwatsonx
+from pandasai.exceptions import APIKeyNotFoundError
 
 
 @pytest.mark.skipif(sys.version_info < (3, 10), reason="requires python3.10 or higher")

extensions/llms/langchain/pandasai_langchain/langchain.py

@@ -14,7 +14,6 @@
 
 from typing import TYPE_CHECKING
 
-
 from pandasai.llm.base import LLM
 
 if TYPE_CHECKING:

extensions/llms/langchain/tests/test_langchain_llm.py

@@ -9,10 +9,10 @@
     GenerationChunk,
     LLMResult,
 )
-
-from pandasai.llm.base import LLM
 from pandasai_langchain.langchain import LangchainLLM
+
 from pandasai.core.prompts.base import BasePrompt
+from pandasai.llm.base import LLM
 
 
 class TestLangchainLLM:
@@ -69,9 +69,10 @@ def test_langchain_chat_call(self, langchain_chat_llm, prompt):
         )
 
     def test_agent_integration(self):
-        from pandasai.agent import Agent
         from unittest.mock import MagicMock, PropertyMock
 
+        from pandasai.agent import Agent
+
         class FakeChatOpenAI(LLM):
             openai_api_key: str = "fake_key"
 

extensions/llms/local/pandasai_local/local_llm.py

@@ -6,7 +6,6 @@
 
 from pandasai.core.prompts.base import BasePrompt
 from pandasai.helpers.memory import Memory
-
 from pandasai.llm.base import LLM
 
 if TYPE_CHECKING:

extensions/llms/local/tests/test_local_llm.py

@@ -3,10 +3,10 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-
-from pandasai.helpers.memory import Memory
 from pandasai_local.local_llm import LocalLLM
+
 from pandasai.core.prompts.base import BasePrompt
+from pandasai.helpers.memory import Memory
 
 
 @pytest.fixture

extensions/llms/openai/pandasai_openai/__init__.py

@@ -1,4 +1,4 @@
-from .openai import OpenAI
 from .azure_openai import AzureOpenAI
+from .openai import OpenAI
 
 __all__ = ["OpenAI", "AzureOpenAI"]

extensions/llms/openai/pandasai_openai/azure_openai.py

@@ -5,6 +5,7 @@
 
 from pandasai.exceptions import APIKeyNotFoundError, MissingModelError
 from pandasai.helpers import load_dotenv
+
 from .base import BaseOpenAI
 
 load_dotenv()

extensions/llms/openai/pandasai_openai/base.py

@@ -4,7 +4,6 @@
 
 from pandasai.core.prompts.base import BasePrompt
 from pandasai.helpers.memory import Memory
-
 from pandasai.llm.base import LLM
 
 if TYPE_CHECKING:

extensions/llms/openai/pandasai_openai/openai.py

@@ -5,6 +5,7 @@
 
 from pandasai.exceptions import APIKeyNotFoundError, UnsupportedModelError
 from pandasai.helpers import load_dotenv
+
 from .base import BaseOpenAI
 
 load_dotenv()

extensions/llms/openai/tests/test_azure_openai.py

@@ -1,9 +1,9 @@
 """Unit tests for the openai LLM class"""
 import openai
 import pytest
+from pandasai_openai import AzureOpenAI
 
 from pandasai.exceptions import APIKeyNotFoundError, MissingModelError
-from pandasai_openai import AzureOpenAI
 
 
 class OpenAIObject:

extensions/llms/openai/tests/test_openai.py

@@ -1,13 +1,14 @@
 """Unit tests for the openai LLM class"""
 
+import os
+from unittest import mock
+
 import openai
 import pytest
-from unittest import mock
-import os
 
-from pandasai.exceptions import APIKeyNotFoundError, UnsupportedModelError
 from extensions.llms.openai.pandasai_openai import OpenAI
 from pandasai.core.prompts.base import BasePrompt
+from pandasai.exceptions import APIKeyNotFoundError, UnsupportedModelError
 
 
 class OpenAIObject:

pandasai/constants.py

@@ -47,3 +47,7 @@
     "databricks",
     "oracle",
 ]
+
+VALID_COLUMN_TYPES = ["string", "integer", "float", "datetime", "boolean"]
+
+VALID_TRANSFORMATION_TYPES = ["anonymize", "convert_timezone"]