Closed HamaWhiteGG closed 1 year ago
I have implemented a local version FlinkSqlToolkitTest.
@Test
void testRunQuery() {
var flinkSql = FlinkSql.builder()
.tableEnv(tableEnv)
.build()
.init();
var llm = ChatOpenAI.builder()
.temperature(0)
.build().init();
var toolkit = new FlinkSqlToolkit(flinkSql, llm);
var agentExecutor = createFlinkSqlAgent(llm, toolkit);
// SELECT SQRT(AVG(Age)) FROM titanic
var actual = agentExecutor.run("whats the square root of the average age?");
// sometimes it's 'The square root of the average age is approximately 5.07.'
assertEquals("5.07", actual);
// TODO: It should be DESC here, not ASC.
// SELECT Name FROM titanic WHERE Survived = 1 ORDER BY Age ASC LIMIT 1
actual = agentExecutor.run("What's the name of the oldest survived passenger?");
assertEquals("Sandstrom, Miss. Marguerite Rut", actual);
}
Similar to Spark SQL Agent, visit https://python.langchain.com/docs/modules/agents/toolkits/spark_sql for more details.