Add MindsDB SQL Support for CREATE, UPDATE, and DROP Agent

mindsdb_sql/parser/dialects/mindsdb/__init__.py

@@ -1,3 +1,4 @@
+from .agents import CreateAgent, DropAgent, UpdateAgent
 from .create_view import CreateView
 from .create_database import CreateDatabase
 from .create_predictor import CreatePredictor, CreateAnomalyDetectionModel

mindsdb_sql/parser/dialects/mindsdb/agents.py

@@ -0,0 +1,94 @@
+from mindsdb_sql.parser.ast.base import ASTNode
+from mindsdb_sql.parser.utils import indent
+
+
+class CreateAgent(ASTNode):
+    """
+    Node for creating a new agent
+    """
+
+    def __init__(self, name, model, params, if_not_exists=False, *args, **kwargs):
+        """
+        Parameters:
+            name (Identifier): name of the agent to create
+            model (str): name of the underlying model to use with the agent
+            params (dict): USING parameters to create the agent with
+            if_not_exists (bool): if True, do not raise an error if the agent exists 
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.model = model
+        self.params = params
+        self.if_not_exists = if_not_exists
+
+    def to_tree(self, level=0, *args, **kwargs):
+        ind = indent(level)
+        out_str = f'{ind}CreateAgent(' \
+                  f'if_not_exists={self.if_not_exists}' \
+                  f'name={self.name.to_string()}, ' \
+                  f'model={self.model}, ' \
+                  f'params={self.params})'
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        using_ar = [f'model={repr(self.model)}']
+        using_ar += [f'{k}={repr(v)}' for k, v in self.params.items()]
+        using_str = ', '.join(using_ar)
+
+        out_str = f'CREATE AGENT {"IF NOT EXISTS " if self.if_not_exists else ""}{self.name.to_string()} USING {using_str}'
+        return out_str
+
+
+class UpdateAgent(ASTNode):
+    """
+    Node for updating an agent
+    """
+
+    def __init__(self, name, updated_params, *args, **kwargs):
+        """
+        Parameters:
+            name (Identifier): name of the agent to update
+            updated_params (dict): new SET parameters of the agent to update
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.params = updated_params
+
+    def to_tree(self, level=0, *args, **kwargs):
+        ind = indent(level)
+        out_str = f'{ind}UpdateAgent(' \
+                  f'name={self.name.to_string()}, ' \
+                  f'updated_params={self.params})'
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        set_ar = [f'{k}={repr(v)}' for k, v in self.params.items()]
+        set_str = ', '.join(set_ar)
+
+        out_str = f'UPDATE AGENT {self.name.to_string()} SET {set_str}'
+        return out_str
+
+
+class DropAgent(ASTNode):
+    """
+    Node for dropping an agent
+    """
+
+    def __init__(self, name, if_exists=False, *args, **kwargs):
+        """
+        Parameters:
+            name (Identifier): name of the agent to drop
+            if_exists (bool): if True, do not raise an error if the agent does not exist
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.if_exists = if_exists
+
+    def to_tree(self, level=0, *args, **kwargs):
+        ind = indent(level)
+        out_str = f'{ind}DropAgent(if_exists={self.if_exists}, name={self.name.to_string()})'
+        return out_str
+
+    def get_string(self, *args, **kwargs):
+        out_str = f'DROP AGENT {"IF EXISTS " if self.if_exists else ""}{str(self.name.to_string())}'
+        return out_str

mindsdb_sql/parser/dialects/mindsdb/lexer.py

@@ -32,6 +32,7 @@ class MindsDBLexer(Lexer):
         ANOMALY, DETECTION,
         KNOWLEDGE_BASE, KNOWLEDGE_BASES,
         SKILL,
+        AGENT,
 
         # SHOW/DDL Keywords
 
@@ -121,6 +122,7 @@ class MindsDBLexer(Lexer):
     KNOWLEDGE_BASE = r'\bKNOWLEDGE[_|\s]BASE\b'
     KNOWLEDGE_BASES = r'\bKNOWLEDGE[_|\s]BASES\b'
     SKILL = r'\bSKILL\b'
+    AGENT = r'\bAGENT\b'
 
     # Misc
     SET = r'\bSET\b'

mindsdb_sql/parser/dialects/mindsdb/parser.py

@@ -1,6 +1,7 @@
 from sly import Parser
 from mindsdb_sql.parser.ast import *
 from mindsdb_sql.parser.ast.drop import DropDatabase, DropView
+from mindsdb_sql.parser.dialects.mindsdb.agents import CreateAgent, DropAgent, UpdateAgent
 from mindsdb_sql.parser.dialects.mindsdb.drop_datasource import DropDatasource
 from mindsdb_sql.parser.dialects.mindsdb.drop_predictor import DropPredictor
 from mindsdb_sql.parser.dialects.mindsdb.drop_dataset import DropDataset
@@ -88,7 +89,10 @@ class MindsDBParser(Parser):
        'drop_kb',
        'create_skill',
        'drop_skill',
-       'update_skill'
+       'update_skill',
+       'create_agent',
+       'drop_agent',
+       'update_agent'
        )
     def query(self, p):
         return p[0]
@@ -154,6 +158,26 @@ def drop_skill(self, p):
     def update_skill(self, p):
         return UpdateSkill(name=p.identifier, updated_params=p.kw_parameter_list)
 
+    # -- Agent --
+    @_('CREATE AGENT if_not_exists_or_empty identifier USING kw_parameter_list')
+    def create_agent(self, p):
+        params = p.kw_parameter_list
+
+        return CreateAgent(
+            name=p.identifier,
+            model=params.pop('model'),
+            params=params,
+            if_not_exists=p.if_not_exists_or_empty
+        )
+
+    @_('DROP AGENT if_exists_or_empty identifier')
+    def drop_agent(self, p):
+        return DropAgent(name=p.identifier, if_exists=p.if_exists_or_empty)
+
+    @_('UPDATE AGENT identifier SET kw_parameter_list')
+    def update_agent(self, p):
+        return UpdateAgent(name=p.identifier, updated_params=p.kw_parameter_list)
+
     # -- ChatBot --
     @_('CREATE CHATBOT identifier USING kw_parameter_list')
     def create_chat_bot(self, p):
@@ -1633,6 +1657,7 @@ def parameter(self, p):
        'WARNINGS',
        'MODEL',
        'MODELS',
+       'AGENT'
        )
     def id(self, p):
         return p[0]

tests/test_parser/test_mindsdb/test_agents.py

@@ -0,0 +1,64 @@
+from mindsdb_sql import parse_sql
+from mindsdb_sql.parser.dialects.mindsdb import *
+from mindsdb_sql.parser.ast import *
+
+
+class TestAgents:
+    def test_create_agent(self):
+        sql = '''
+            create agent if not exists my_agent
+            using 
+            model = 'my_model', 
+            skills = ['skill1', 'skill2']
+        '''
+        ast = parse_sql(sql, dialect='mindsdb')
+        expected_ast = CreateAgent(
+            name=Identifier('my_agent'),
+            model='my_model',
+            params={'skills': ['skill1', 'skill2']},
+            if_not_exists=True
+        )
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()
+
+        # Parse again after rendering to catch problems with rendering.
+        ast = parse_sql(str(ast), dialect='mindsdb')
+        assert str(ast) == str(expected_ast)
+
+    def test_update_agent(self):
+        sql = '''
+            update agent my_agent
+            set
+            model = 'new_model',
+            skills = ['new_skill1', 'new_skill2']
+        '''
+        ast = parse_sql(sql, dialect='mindsdb')
+        expected_params = {
+            'model': 'new_model',
+            'skills': ['new_skill1', 'new_skill2']
+        }
+        expected_ast = UpdateAgent(
+            name=Identifier('my_agent'),
+            updated_params=expected_params
+        )
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()
+
+        # Parse again after rendering to catch problems with rendering.
+        ast = parse_sql(str(ast), dialect='mindsdb')
+        assert str(ast) == str(expected_ast)
+
+    def test_drop_agent(self):
+        sql = '''
+            drop agent if exists my_agent
+        '''
+        ast = parse_sql(sql, dialect='mindsdb')
+        expected_ast = DropAgent(
+            name=Identifier('my_agent'), if_exists=True
+        )
+        assert str(ast) == str(expected_ast)
+        assert ast.to_tree() == expected_ast.to_tree()
+
+        # Parse again after rendering to catch problems with rendering.
+        ast = parse_sql(str(ast), dialect='mindsdb')
+        assert str(ast) == str(expected_ast)