@@ -374,7 +374,7 @@ def test_insert_into_knowledge_base(self):
374374
375375 def test_evaluate_knowledge_base (self ):
376376 sql = """
377- EVALUATE my_knowledge_base
377+ EVALUATE KNOWLEDGE_BASE my_knowledge_base
378378 USING
379379 TEST_TABLE = my_database.some_table_1,
380380 SAVE_TO = my_database.some_table_2,
@@ -383,24 +383,26 @@ def test_evaluate_knowledge_base(self):
383383 "model": "gpt-3.5-turbo",
384384 "api_key": "my_api_key"
385385 },
386- generate_data = {
386+ GENERATE_DATA = {
387387 "from_sql": "SELECT content FROM my_database.some_table",
388388 "count": 100
389389 }
390390 """
391391 ast = parse_sql (sql )
392392 expected_ast = EvaluateKnowledgeBase (
393393 name = Identifier ("my_knowledge_base" ),
394- test_table = Identifier (parts = ["my_database" , "some_table_1" ]),
395- save_to = Identifier (parts = ["my_database" , "some_table_2" ]),
396- llm = {
397- "provider" : "openai" ,
398- "model" : "gpt-3.5-turbo" ,
399- "api_key" : "my_api_key"
400- },
401- generate_data = {
402- "from_sql" : "SELECT content FROM my_database.some_table" ,
403- "count" : 100
394+ params = {
395+ "test_table" : Identifier (parts = ["my_database" , "some_table_1" ]),
396+ "save_to" : Identifier (parts = ["my_database" , "some_table_2" ]),
397+ "llm" : {
398+ "provider" : "openai" ,
399+ "model" : "gpt-3.5-turbo" ,
400+ "api_key" : "my_api_key"
401+ },
402+ "generate_data" : {
403+ "from_sql" : "SELECT content FROM my_database.some_table" ,
404+ "count" : 100
405+ }
404406 }
405407 )
406408 assert ast == expected_ast
0 commit comments