Merge pull request #299 from mindsdb/feature/knowledge-base-sql

Add sql syntax for knowledge base

mindsdb_sql/parser/dialects/mindsdb/__init__.py

@@ -16,6 +16,8 @@
 from .drop_job import DropJob
 from .chatbot import CreateChatBot, UpdateChatBot, DropChatBot
 from .trigger import CreateTrigger, DropTrigger
+from .knowledge_base import CreateKnowledgeBase, DropKnowledgeBase
 
 # remove it in next release
 CreateDatasource = CreateDatabase
+

mindsdb_sql/parser/dialects/mindsdb/knowledge_base.py

@@ -0,0 +1,98 @@
+from mindsdb_sql.parser.ast.base import ASTNode
+from mindsdb_sql.parser.utils import indent
+
+
+class CreateKnowledgeBase(ASTNode):
+    """
+    Create a new knowledge base
+    """
+    def __init__(
+        self,
+        name,
+        model,
+        storage,
+        from_select=None,
+        params=None,
+        if_not_exists=False,
+        *args,
+        **kwargs,
+    ):
+        """
+        Args:
+            name: Identifier -- name of the knowledge base
+            model: Identifier -- name of the model to use
+            storage: Identifier -- name of the storage to use
+            from_select: SelectStatement -- select statement to use as the source of the knowledge base
+            params: dict -- additional parameters to pass to the knowledge base. E.g., chunking strategy, etc.
+            if_not_exists: bool -- if True, do not raise an error if the knowledge base already exists
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.model = model
+        self.storage = storage
+        self.params = params
+        self.if_not_exists = if_not_exists
+        self.from_query = from_select
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = f"""
+        {ind}CreateKnowledgeBase(
+        {ind}    if_not_exists={self.if_not_exists},
+        {ind}    name={self.name.to_string()},
+        {ind}    from_query={self.from_query.to_tree(level=level+1) if self.from_query else None},
+        {ind}    model={self.model.to_string()},
+        {ind}    storage={self.storage.to_string()},
+        {ind}    params={self.params}
+        {ind})
+        """
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        params = self.params.copy()
+        using_ar = [f"{k}={repr(v)}" for k, v in params.items()]
+        using_str = ", ".join(using_ar)
+        from_query_str = (
+            f"FROM ({self.from_query.get_string()})" if self.from_query else ""
+        )
+
+        out_str = (
+            f"CREATE KNOWLEDGE_BASE {'IF NOT EXISTS' if self.if_not_exists else ''}{self.name.to_string()} "
+            f"{from_query_str} "
+            f"USING {using_str},"
+            f"  MODEL = {self.model.to_string()}, "
+            f"  STORAGE {self.storage.to_string()} "
+        )
+
+        return out_str
+
+    def __repr__(self) -> str:
+        return self.to_tree()
+
+
+class DropKnowledgeBase(ASTNode):
+    """
+    Delete a knowledge base
+    """
+    def __init__(self, name, if_exists=False, *args, **kwargs):
+        """
+        Args:
+            name: Identifier -- name of the knowledge base
+            if_exists: bool -- if True, do not raise an error if the knowledge base does not exist
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.if_exists = if_exists
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+        out_str = (
+            f"{ind}DropKnowledgeBase("
+            f"{ind}    if_exists={self.if_exists},"
+            f"name={self.name.to_string()})"
+        )
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        out_str = f'DROP KNOWLEDGE_BASE {"IF EXISTS" if self.if_exists else ""}{self.name.to_string()}'
+        return out_str

mindsdb_sql/parser/dialects/mindsdb/lexer.py

@@ -30,6 +30,7 @@ class MindsDBLexer(Lexer):
         LATEST, HORIZON, USING,
         ENGINE, TRAIN, PREDICT, PARAMETERS, JOB, CHATBOT, EVERY,PROJECT,
         ANOMALY, DETECTION,
+        KNOWLEDGE_BASE, KNOWLEDGE_BASES,
 
         # SHOW/DDL Keywords
 
@@ -115,6 +116,9 @@ class MindsDBLexer(Lexer):
     ANOMALY = r'\bANOMALY\b'
     DETECTION = r'\bDETECTION\b'
 
+    KNOWLEDGE_BASE = r'\bKNOWLEDGE[_|\s]BASE\b'
+    KNOWLEDGE_BASES = r'\bKNOWLEDGE[_|\s]BASES\b'
+
     # Misc
     SET = r'\bSET\b'
     START = r'\bSTART\b'