Merge pull request #27 from mindsdb/feature/evaluate_kb

Added the Syntax for EVALUATE KNOWLEDGE_BASE

Author: ea-rus

mindsdb_sql_parser/ast/mindsdb/__init__.py

@@ -15,7 +15,8 @@
 from .drop_job import DropJob
 from .chatbot import CreateChatBot, UpdateChatBot, DropChatBot
 from .trigger import CreateTrigger, DropTrigger
-from .knowledge_base import CreateKnowledgeBase, DropKnowledgeBase, CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex
+from .knowledge_base import CreateKnowledgeBase, DropKnowledgeBase, CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex \
+    , EvaluateKnowledgeBase
 from .skills import CreateSkill, DropSkill, UpdateSkill
 
 # remove it in next release

mindsdb_sql_parser/ast/mindsdb/knowledge_base.py

@@ -1,4 +1,5 @@
 from mindsdb_sql_parser.ast.base import ASTNode
+from mindsdb_sql_parser.ast import Identifier
 from mindsdb_sql_parser.utils import indent
 
 
@@ -106,6 +107,7 @@ 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
 
+
 class CreateKnowledgeBaseIndex(ASTNode):
     """
     Create a new index in the knowledge base
@@ -127,6 +129,7 @@ def get_string(self, *args, **kwargs):
         out_str = f'CREATE INDEX ON KNOWLEDGE_BASE {self.name.to_string()}'
         return out_str
 
+
 class DropKnowledgeBaseIndex(ASTNode):
     """
     Delete an index in the knowledge base
@@ -146,4 +149,61 @@ def to_tree(self, *args, level=0, **kwargs):
 
     def get_string(self, *args, **kwargs):
         out_str = f'DROP INDEX ON KNOWLEDGE_BASE {self.name.to_string()}'
-        return out_str
+        return out_str
+
+
+class EvaluateKnowledgeBase(ASTNode):
+    """
+    Evaluate a knowledge base.
+    """
+    def __init__(self, name: Identifier, params: dict = None, *args, **kwargs):
+        """
+        Args:
+            name: Identifier -- name of the knowledge base to evaluate.
+            params: dict -- parameters to pass to the evaluation.
+        """
+        super().__init__(*args, **kwargs)
+        self.name = name
+        self.params = params if params is not None else {}
+
+    def to_tree(self, *args, level=0, **kwargs):
+        ind = indent(level)
+
+        param_str = ""
+        if self.params:
+            param_items = []
+            for k, v in self.params.items():
+                if isinstance(v, Identifier):
+                    param_items.append(f"{ind}    {k}={v.to_string()}")
+                else:
+                    param_items.append(f"{ind}    {k}={repr(v)}")
+            param_str = ",\n".join(param_items)
+
+        output_str = (
+            f"{ind}EvaluateKnowledgeBase(\n"
+            f"{ind}    name={self.name.to_string()},\n"
+            f"{param_str}\n"
+            f"{ind})"
+        )
+
+        return output_str
+
+    def get_string(self, *args, **kwargs):
+        using_str = ""
+
+        if self.params:
+            using_args = []
+            for k, v in self.params.items():
+                if isinstance(v, Identifier):
+                    using_args.append(f"{k}={v.to_string()}")
+                else:
+                    using_args.append(f"{k}={repr(v)}")
+
+            using_str = "USING " + ", ".join(using_args)
+
+        output_str = (
+            f"EVALUATE KNOWLEDGE_BASE {self.name.to_string()} "
+            f"{using_str}"
+        )
+
+        return output_str.strip()

mindsdb_sql_parser/parser.py

@@ -17,7 +17,7 @@
 from mindsdb_sql_parser.ast.mindsdb.latest import Latest
 from mindsdb_sql_parser.ast.mindsdb.evaluate import Evaluate
 from mindsdb_sql_parser.ast.mindsdb.knowledge_base import CreateKnowledgeBase, DropKnowledgeBase, \
-    CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex
+    CreateKnowledgeBaseIndex, DropKnowledgeBaseIndex, EvaluateKnowledgeBase
 from mindsdb_sql_parser.ast.mindsdb.skills import CreateSkill, DropSkill, UpdateSkill
 from mindsdb_sql_parser.exceptions import ParsingException
 from mindsdb_sql_parser.ast.mindsdb.retrain_predictor import RetrainPredictor
@@ -89,6 +89,7 @@ class MindsDBParser(Parser):
        'drop_trigger',
        'create_kb',
        'drop_kb',
+       'evaluate_kb',
        'create_skill',
        'drop_skill',
        'update_skill',
@@ -149,6 +150,13 @@ def  drop_index(self, p):
     def drop_kb(self, p):
         return DropKnowledgeBase(name=p.identifier, if_exists=p.if_exists_or_empty)
 
+    @_('EVALUATE KNOWLEDGE_BASE identifier USING kw_parameter_list')
+    def evaluate_kb(self, p):
+        return EvaluateKnowledgeBase(
+            name=p.identifier,
+            params={k.lower(): v for k, v in p.kw_parameter_list.items()}
+        )
+
     # -- Skills --
     @_('CREATE SKILL if_not_exists_or_empty identifier USING kw_parameter_list')
     def create_skill(self, p):

tests/test_mindsdb/test_knowledgebase.py

@@ -2,6 +2,7 @@
 from mindsdb_sql_parser.ast.mindsdb.knowledge_base import (
     CreateKnowledgeBase,
     DropKnowledgeBase,
+    EvaluateKnowledgeBase
 )
 from mindsdb_sql_parser.ast import (
     Select,
@@ -370,3 +371,38 @@ def test_insert_into_knowledge_base(self):
             ),
         )
         assert ast == expected_ast
+
+    def test_evaluate_knowledge_base(self):
+        sql = """
+            EVALUATE KNOWLEDGE_BASE my_knowledge_base
+            USING
+                TEST_TABLE = my_database.some_table_1,
+                SAVE_TO = my_database.some_table_2,
+                LLM = {
+                    "provider": "openai",
+                    "model": "gpt-3.5-turbo",
+                    "api_key": "my_api_key"
+                },
+                GENERATE_DATA = {
+                    "from_sql": "SELECT content FROM my_database.some_table",
+                    "count": 100
+                }
+        """
+        ast = parse_sql(sql)
+        expected_ast = EvaluateKnowledgeBase(
+            name=Identifier("my_knowledge_base"),
+            params={
+                "test_table": Identifier(parts=["my_database", "some_table_1"]),
+                "save_to": Identifier(parts=["my_database", "some_table_2"]),
+                "llm": {
+                    "provider": "openai",
+                    "model": "gpt-3.5-turbo",
+                    "api_key": "my_api_key"
+                },
+                "generate_data": {
+                    "from_sql": "SELECT content FROM my_database.some_table",
+                    "count": 100
+                }
+            }
+        )
+        assert ast == expected_ast