Merge pull request #69 from BohuTANG/dev-llmchain

feat: Upgrade to llmchain.rs crate for enhanced performance and power

Author: bohutang

Cargo.lock

@@ -1,4 +1,5 @@
 [workspace]
 members = [
     "app",
-]
+]
+

Cargo.toml

@@ -1,49 +1,16 @@
-# AskBend: SQL-based Knowledge Base Search and Completion using Databend
-## Demo https://ask.databend.rs/
-
-AskBend is a Rust project that utilizes the power of Databend and OpenAI to create a SQL-based knowledge base from Markdown files.
-
-Databend is a cloud-native data warehouse adept at storing and performing vector computations, making it suitable for this use case.
-
-[Databend Cloud](https://databend.com) seamlessly integrates with OpenAI's capabilities, such as embedding generation, cosine distance calculation, and text completion. This integration means you don't need to interact with OpenAI directly; Databend Cloud manages everything.
-
-The project automatically generates document embeddings from the content, enabling users to search and retrieve the most relevant information to their queries using SQL.
-
-SQL-Based means you don't need any OpenAI API knowledge. With the Databend Cloud platform, you can perform these tasks using SQL. Some SQL AI functions of Databend Cloud include:
-
-- [ai_embedding_vector](https://databend.rs/doc/sql-functions/ai-functions/ai-embedding-vector): Get the vector from OpenAI API
-- [ai_text_completion](https://databend.rs/doc/sql-functions/ai-functions/ai-text-completion): Get the completion of a text
-- [cosine_distance](https://databend.rs/doc/sql-functions/ai-functions/ai-cosine-distance): Calculate the distance between embedding vectors
+# AskBend: Empower developers to explore & implement your knowledge base
 
-## Overview
-
-The project follows this general process:
+## Demo https://ask.databend.rs/
 
-1. Read and parse Markdown files from a directory.
-2. Extract the content and store it in the askbend.doc table.
-3. Compute embeddings for the content using Databend Cloud's built-in AI capabilities, including OpenAI's embedding generation, all through SQL.
-4. When a user queries, generate the query embedding using Databend Cloud's SQL-based `ai_embedding_vector` function.
-5. Perform a vector calculation to find the most relevant doc.content using Databend Cloud's SQL-based `cosine_distance` function.
-6. Concatenate the retrieved content and use OpenAI's completion capabilities with Databend Cloud's SQL-based `ai_text_completion` function.
-7. Output the completion result in Markdown format.
+AskBend is a project built in Rust that leverages the [llmchain.rs](https://github.com/shafishlabs/llmchain.rs) library to create a SQL-based knowledge base from Markdown files.
 
 ## Setup
 
 ### 1. Download 
 
 https://github.com/datafuselabs/askbend/releases
 
-### 2. Create a table in your Databend Cloud
-
-[table](schema/table.sql):
-```
-CREATE DATABASE askbend;
-USE askbend;
-
-CREATE TABLE doc (path VARCHAR, content VARCHAR, embedding ARRAY(FLOAT32));
-```
-
-### 3. Modify the configuration file [conf/askbend.toml](conf/askbend.toml)
+### 2. Modify the configuration file [conf/askbend.toml](conf/askbend.toml)
 
 ```
 # Usage:
@@ -66,20 +33,11 @@ port = 8081
 
 [query]
 top = 3
-product = "YourProductName"
-prompt = '''
-<your prompt including {{product}}> ... 
-Documentation sections:
-{{context}}
-
-Question:
-{{query}}
-'''
 ```
 
-### 4. Prepare your Markdown files by copying them to the `data/` directory
+### 3. Prepare your Markdown files by copying them to the `data/` directory
 
-### 5. Parse the Markdown files and build embeddings
+### 4. Parse the Markdown files and build embeddings
 
 ```
 ./target/release/askbend -c conf/askbend.toml --rebuild
@@ -96,13 +54,13 @@ Question:
 The `--rebuild` flag rebuilds all the embeddings for the data directory. This process may take a few minutes, depending on the number of Markdown files.
 
 
-### 6. Start the API server
+### 5. Start the API server
 
 ```
 ./target/release/askbend -c conf/askbend.toml
 ```
 
-### 7. Query your Markdown knowledge base using the API
+### 6. Query your Markdown knowledge base using the API
 ```
 curl -X POST -H "Content-Type: application/json" -d '{"query": "tell me how to do copy"}' http://localhost:8081/query
 ```

README.md

@@ -14,24 +14,27 @@ doctest = false
 test = false
 
 [dependencies]
+actix-cors = "0.6.4"
+actix-web = "4.3.1"
 clap = { version = "4.1.7", features = ["derive", "env"] }
 chrono = "0.4.24"
-comrak = "0.18.0"
-databend-driver = "0.1.4"
+databend-driver = "0.2.23"
 log = "0.4.0"
 env_logger = "0.9.0"
 regex = "1.7.3"
 serde = { version = "1.0.159", features = ["derive"] }
 serfig = "0.0.3"
-actix-web = "4.3.1"
-actix-cors = "0.6.4"
 anyhow = "1.0.58"
 syn = "1.0"
 proc-macro2 = "1.0"
 quote = "1.0"
 markdown = "0.3.0"
-tokio = { version = "1.19.2", features = ["full"] }
+tokio = { version = "1.28", features = ["full"] }
 tokio-stream = "0.1.12"
 walkdir = "2.3.3"
 
+llmchain = { git = "https://github.com/shafishlabs/llmchain.rs" , rev = "f204230"}
+
 [dev-dependencies]
+
+

app/Cargo.toml

@@ -12,16 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
 use anyhow::Result;
 use askbend::APIHandler;
 use askbend::Config;
-use askbend::DatabendDriver;
-use askbend::FileOperator;
-use askbend::Markdown;
-use askbend::Parse;
-use askbend::RustCode;
 use env_logger::Builder;
 use env_logger::Env;
+use llmchain::DatabendEmbedding;
+use llmchain::DatabendVectorStore;
+use llmchain::DocumentLoader;
+use llmchain::DocumentPath;
+use llmchain::DocumentSplitter;
+use llmchain::MarkdownLoader;
+use llmchain::MarkdownSplitter;
+use llmchain::VectorStore;
 use log::info;
 use tokio::time::Instant;
 
@@ -46,58 +51,44 @@ async fn main() -> Result<()> {
 
 /// Rebuild all embeddings.
 async fn rebuild_embedding(conf: &Config) -> Result<()> {
-    let ext = conf.data.file_ext.clone();
-    let ignore_dirs = conf.data.ignore_dirs.clone();
-    info!("Step-1: begin parser all {} files", ext);
-    let file_opt = FileOperator::create(&conf.data.path, &conf.data.file_ext, &ignore_dirs);
-    let files = file_opt.list()?;
-
-    let snippet_files = match conf.data.file_ext.as_str() {
-        "md" => Markdown::parse_multiple(
-            &files
-                .iter()
-                .map(|v| v.full_path.clone())
-                .collect::<Vec<String>>(),
-        ),
+    let local_disk = llmchain::LocalDisk::create()?;
+    let markdown_loader = MarkdownLoader::create(local_disk.clone());
+    let directory_loader =
+        llmchain::DirectoryLoader::create(local_disk).with_loader("**/*.md", markdown_loader);
+    let documents = directory_loader
+        .load(DocumentPath::Str(conf.data.path.clone()))
+        .await?;
+    info!("Step-1: parser all files:{}", documents.len());
 
-        "rs" => RustCode::parse_multiple(
-            &files
-                .iter()
-                .map(|v| v.full_path.clone())
-                .collect::<Vec<String>>(),
-        ),
-        _ => Err(anyhow::Error::msg(format!(
-            "Only support file ext: md or rs, got:{}",
-            conf.data.file_ext
-        ))),
-    }?;
+    let documents = MarkdownSplitter::create().split_documents(&documents)?;
+    info!("Step-2: split all files to:{}", documents.len());
 
+    let now = Instant::now();
     info!(
-        "Step-1: finish parser all {} files:{}, sections:{}, tokens:{}",
-        ext,
-        files.len(),
-        snippet_files.all_snippets(),
-        snippet_files.all_tokens()
+        "Step-3: begin embedding to table:{}.{}",
+        conf.database.database, conf.database.table
     );
+    let dsn = conf.database.dsn.clone();
+    let databend_embedding = Arc::new(DatabendEmbedding::create(&dsn));
+    let databend_vector_store = DatabendVectorStore::create(&dsn, databend_embedding)
+        .with_database(&conf.database.database)
+        .with_table(&conf.database.table);
+    databend_vector_store.init().await?;
 
-    let dal = DatabendDriver::connect(conf)?;
-
-    info!("Step-2: begin insert to table");
-    dal.insert(&snippet_files).await?;
-    info!("Step-2: finish insert to table");
-
-    info!("Step-3: begin generate embedding, may take some minutes");
-    dal.embedding().await?;
-    info!("Step-3: finish generate embedding");
-
+    let _ = databend_vector_store.add_documents(documents).await?;
+    info!(
+        "Step-3: finish embedding to table:{}.{}, cost {}",
+        conf.database.database,
+        conf.database.table,
+        now.elapsed().as_secs()
+    );
     Ok(())
 }
 
 /// Start the api server.
 async fn start_api_server(conf: &Config) -> Result<()> {
     info!("Start api server {}:{}", conf.server.host, conf.server.port);
-    let dal = DatabendDriver::connect(conf)?;
-    let handler = APIHandler::create(&conf.server, dal.clone());
+    let handler = APIHandler::create(conf);
     handler.start().await?;
     Ok(())
 }

app/bin/ask.rs

@@ -21,27 +21,21 @@ use anyhow::Result;
 
 use crate::api::query::query_handler;
 use crate::api::status::status_handler;
-use crate::configs::ServerConfig;
-use crate::DatabendDriver;
+use crate::Config;
 
 pub struct APIHandler {
-    pub conf: ServerConfig,
-    pub db: DatabendDriver,
+    pub conf: Config,
 }
 
 impl APIHandler {
-    pub fn create(conf: &ServerConfig, db: DatabendDriver) -> Self {
-        APIHandler {
-            conf: conf.clone(),
-            db,
-        }
+    pub fn create(conf: &Config) -> Self {
+        APIHandler { conf: conf.clone() }
     }
 
     pub async fn start(self) -> Result<()> {
         let conf = self.conf.clone();
-        let host = conf.host.clone();
-        let port = conf.port;
-        let data = self.db.clone();
+        let host = conf.server.host.clone();
+        let port = conf.server.port;
 
         HttpServer::new(move || {
             let mut cors = Cors::default()
@@ -52,12 +46,12 @@ impl APIHandler {
                     http::header::CONTENT_TYPE,
                 ])
                 .max_age(3600);
-            for origin in &conf.cors {
+            for origin in &conf.server.cors {
                 cors = cors.allowed_origin(origin);
             }
             App::new()
                 .wrap(cors)
-                .app_data(web::Data::new(data.clone()))
+                .app_data(web::Data::new(conf.clone()))
                 .route("/status", web::get().to(status_handler))
                 .route("/query", web::post().to(query_handler))
         })

app/src/api/http.rs

@@ -18,7 +18,8 @@ use actix_web::HttpResponse;
 use actix_web::Responder;
 use log::error;
 
-use crate::DatabendDriver;
+use crate::llm::BendLLM;
+use crate::Config;
 
 #[derive(serde::Deserialize)]
 pub struct Query {
@@ -31,17 +32,13 @@ struct Response {
 }
 
 /// curl -X POST -H "Content-Type: application/json" -d '{"query": "whats the fast way to load data to databend"}' http://localhost:8081/query
-pub async fn query_handler(
-    query: web::Json<Query>,
-    db: web::Data<DatabendDriver>,
-) -> impl Responder {
-    let result = db.query(&query.query).await;
+pub async fn query_handler(query: web::Json<Query>, conf: web::Data<Config>) -> impl Responder {
+    let llm = BendLLM::create(&conf);
+    let result = llm.query(&query.query).await;
     match result {
         Ok(result) => {
             let response = if !result.is_empty() {
-                Response {
-                    result: result[0].to_string(),
-                }
+                Response { result }
             } else {
                 Response {
                     result: "Sorry, I dont know how to help with that.".to_string(),

app/src/api/query.rs

@@ -12,35 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use regex::Regex;
-
 pub fn escape_sql_string(input: &str) -> String {
     input
         .replace('\\', "\\\\")
         .replace('\'', "''")
         .replace('\n', " ")
         .replace('\r', "\\r")
 }
-
-pub fn remove_markdown_links(input: &str) -> String {
-    let link_regex = Regex::new(r"\[(?P<text>[^\]]+)\]\((?P<url>[^\)]+)\)").unwrap();
-    let result = link_regex.replace_all(input, "$text");
-    result.to_string()
-}
-
-pub fn replace_multiple_spaces(input: &str) -> String {
-    let re = Regex::new(r" +").unwrap();
-    re.replace_all(input, " ").to_string()
-}
-
-pub trait LengthWithoutSymbols {
-    fn length_without_symbols(&self) -> usize;
-}
-
-impl LengthWithoutSymbols for String {
-    fn length_without_symbols(&self) -> usize {
-        self.chars()
-            .filter(|c| c.is_alphanumeric() || c.is_whitespace())
-            .count()
-    }
-}