#[cfg(feature = "postgres")]
use langchain_rust::{
chain::{options::ChainCallOptions, Chain, SQLDatabaseChainBuilder},
llm::openai::OpenAI,
tools::{postgres::PostgreSQLEngine, SQLDatabaseBuilder},
};
#[cfg(feature = "postgres")]
use std::io::{self, Write};
#[cfg(feature = "postgres")]
#[tokio::main]
async fn main() {
let options = ChainCallOptions::default();
let llm = OpenAI::default();
let db = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
let engine = PostgreSQLEngine::new(&db).await.unwrap();
let db = SQLDatabaseBuilder::new(engine).build().await.unwrap();
let chain = SQLDatabaseChainBuilder::new()
.llm(llm)
.top_k(4)
.database(db)
.options(options)
.build()
.expect("Failed to build LLMChain");
print!("Please enter a question: ");
io::stdout().flush().unwrap();
let mut input = String::new();
io::stdin().read_line(&mut input).unwrap();
let input = input.trim();
let input_variables = chain.prompt_builder().query(input).build();
match chain.invoke(input_variables).await {
Ok(result) => {
println!("Result: {:?}", result);
}
Err(e) => panic!("Error invoking LLMChain: {:?}", e),
}
}
#[cfg(not(feature = "postgres"))]
fn main() {
println!("This example requires the 'postgres' feature to be enabled.");
println!("Please run the command as follows:");
println!("cargo run --example sql_chain --features postgres");
}