61 lines
2.0 KiB
Rust
61 lines
2.0 KiB
Rust
use swiftide::{
|
|
indexing::{
|
|
self, EmbeddedField,
|
|
loaders::FileLoader,
|
|
transformers::{ChunkMarkdown, Embed, MetadataQAText, metadata_qa_text},
|
|
},
|
|
integrations::{dashscope::DashscopeBuilder, lancedb::LanceDB},
|
|
query::{
|
|
self,
|
|
answers::{self},
|
|
query_transformers::{self},
|
|
response_transformers,
|
|
},
|
|
};
|
|
use temp_dir::TempDir;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
tracing_subscriber::fmt::init();
|
|
|
|
let client = DashscopeBuilder::default()
|
|
.default_embed_model("text-embedding-v2")
|
|
.default_prompt_model("qwen-long")
|
|
.build()?;
|
|
let tempdir = TempDir::new().unwrap();
|
|
let lancedb = LanceDB::builder()
|
|
.uri(tempdir.child("lancedb").to_str().unwrap())
|
|
.vector_size(1536)
|
|
.with_vector(EmbeddedField::Combined)
|
|
.with_metadata(metadata_qa_text::NAME)
|
|
.table_name("swiftide_test")
|
|
.build()
|
|
.unwrap();
|
|
|
|
indexing::Pipeline::from_loader(FileLoader::new(".").with_extensions(&["md"]))
|
|
.with_default_llm_client(client.clone())
|
|
.then_chunk(ChunkMarkdown::from_chunk_range(10..2048))
|
|
.then(MetadataQAText::new(client.clone()))
|
|
.then_in_batch(Embed::new(client.clone()).with_batch_size(10))
|
|
.then_store_with(lancedb.clone())
|
|
.run()
|
|
.await?;
|
|
|
|
let pipeline = query::Pipeline::default()
|
|
.then_transform_query(query_transformers::GenerateSubquestions::from_client(
|
|
client.clone(),
|
|
))
|
|
.then_transform_query(query_transformers::Embed::from_client(client.clone()))
|
|
.then_retrieve(lancedb.clone())
|
|
.then_transform_response(response_transformers::Summary::from_client(client.clone()))
|
|
.then_answer(answers::Simple::from_client(client.clone()));
|
|
|
|
let result = pipeline
|
|
.query("What is swiftide? Please provide an elaborate explanation")
|
|
.await?;
|
|
|
|
println!("====");
|
|
println!("{:?}", result.answer());
|
|
Ok(())
|
|
}
|