@@ -374,7 +374,7 @@ def test_insert_into_knowledge_base(self):
374
374
375
375
def test_evaluate_knowledge_base (self ):
376
376
sql = """
377
- EVALUATE my_knowledge_base
377
+ EVALUATE KNOWLEDGE_BASE my_knowledge_base
378
378
USING
379
379
TEST_TABLE = my_database.some_table_1,
380
380
SAVE_TO = my_database.some_table_2,
@@ -383,24 +383,26 @@ def test_evaluate_knowledge_base(self):
383
383
"model": "gpt-3.5-turbo",
384
384
"api_key": "my_api_key"
385
385
},
386
- generate_data = {
386
+ GENERATE_DATA = {
387
387
"from_sql": "SELECT content FROM my_database.some_table",
388
388
"count": 100
389
389
}
390
390
"""
391
391
ast = parse_sql (sql )
392
392
expected_ast = EvaluateKnowledgeBase (
393
393
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
+ }
404
406
}
405
407
)
406
408
assert ast == expected_ast
0 commit comments