HamaWhiteGG/langchain-java

[Feature] Support Spark SQL Agent

Closed this issue · 1 comments

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);
}