[Feature] Support Spark SQL Agent
Closed this issue · 1 comments
HamaWhiteGG commented
I have implemented a local version that allows you to view SparkSqlToolkitTest.
@Test
void testRunningQuery() {
var sparkSql = SparkSql.builder()
.spark(spark)
.schema(SCHEMA)
.build().init();
var llm = ChatOpenAI.builder()
.temperature(0)
.build().init();
var toolkit = new SparkSqlToolkit(sparkSql, llm);
var agentExecutor = createSparkSqlAgent(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.45.'
assertEquals("5.45", actual);
// SELECT Name FROM titanic WHERE Survived = 1 ORDER BY Age DESC LIMIT 1
actual = agentExecutor.run("What's the name of the oldest survived passenger?");
assertEquals("Barkworth, Mr. Algernon Henry Wilson", actual);
}