gmakstutis/ragflow
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

go through upload, create kb, add doc to kb

Browse source
This commit is contained in:
kevinhu 2023-12-19 19:25:43 +08:00
parent 2d1ffc5f68
commit 0fa9243a13
16 changed files with 262 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Cargo.toml
View file

@ -24,6 +24,7 @@ listenfd = "1.0.1"
chrono = "0.4.31"
migration = { path = "./migration" }
futures-util = "0.3.29"
actix-multipart-extract = "0.1.5"
[[bin]]
name = "doc_gpt"

40
migration/src/m20220101_000001_create_table.rs
View file

@ -47,8 +47,8 @@ impl MigrationTrait for Migration {
.col(ColumnDef::new(TagInfo::Uid).big_integer().not_null())
.col(ColumnDef::new(TagInfo::TagName).string().not_null())
.col(ColumnDef::new(TagInfo::Regx).string())
.col(ColumnDef::new(TagInfo::Color).big_integer().default(1))
.col(ColumnDef::new(TagInfo::Icon).big_integer().default(1))
.col(ColumnDef::new(TagInfo::Color).tiny_unsigned().default(1))
.col(ColumnDef::new(TagInfo::Icon).tiny_unsigned().default(1))
.col(ColumnDef::new(TagInfo::Dir).string())
.col(ColumnDef::new(TagInfo::CreatedAt).date().not_null())
.col(ColumnDef::new(TagInfo::UpdatedAt).date().not_null())
@ -62,6 +62,13 @@ impl MigrationTrait for Migration {
Table::create()
.table(Tag2Doc::Table)
.if_not_exists()
.col(
ColumnDef::new(Tag2Doc::Id)
.big_integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(Tag2Doc::TagId).big_integer())
.col(ColumnDef::new(Tag2Doc::Did).big_integer())
.to_owned(),
@ -73,6 +80,13 @@ impl MigrationTrait for Migration {
Table::create()
.table(Kb2Doc::Table)
.if_not_exists()
.col(
ColumnDef::new(Kb2Doc::Id)
.big_integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(Kb2Doc::KbId).big_integer())
.col(ColumnDef::new(Kb2Doc::Did).big_integer())
.to_owned(),
@ -84,6 +98,13 @@ impl MigrationTrait for Migration {
Table::create()
.table(Dialog2Kb::Table)
.if_not_exists()
.col(
ColumnDef::new(Dialog2Kb::Id)
.big_integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(Dialog2Kb::DialogId).big_integer())
.col(ColumnDef::new(Dialog2Kb::KbId).big_integer())
.to_owned(),
@ -95,6 +116,13 @@ impl MigrationTrait for Migration {
Table::create()
.table(Doc2Doc::Table)
.if_not_exists()
.col(
ColumnDef::new(Doc2Doc::Id)
.big_integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(Doc2Doc::ParentId).big_integer())
.col(ColumnDef::new(Doc2Doc::Did).big_integer())
.to_owned(),
@ -112,7 +140,7 @@ impl MigrationTrait for Migration {
.primary_key())
.col(ColumnDef::new(KbInfo::Uid).big_integer().not_null())
.col(ColumnDef::new(KbInfo::KbName).string().not_null())
.col(ColumnDef::new(KbInfo::Icon).big_integer().default(1))
.col(ColumnDef::new(KbInfo::Icon).tiny_unsigned().default(1))
.col(ColumnDef::new(KbInfo::CreatedAt).date().not_null())
.col(ColumnDef::new(KbInfo::UpdatedAt).date().not_null())
.col(ColumnDef::new(KbInfo::IsDeleted).boolean().default(false))
@ -149,7 +177,7 @@ impl MigrationTrait for Migration {
.table(DialogInfo::Table)
.if_not_exists()
.col(ColumnDef::new(DialogInfo::DialogId)
.big_integer()
.big_integer()
.not_null()
.auto_increment()
.primary_key())
@ -241,6 +269,7 @@ enum TagInfo {
#[derive(DeriveIden)]
enum Tag2Doc {
Table,
Id,
TagId,
Did,
}
@ -248,6 +277,7 @@ enum Tag2Doc {
#[derive(DeriveIden)]
enum Kb2Doc {
Table,
Id,
KbId,
Did,
}
@ -255,6 +285,7 @@ enum Kb2Doc {
#[derive(DeriveIden)]
enum Dialog2Kb {
Table,
Id,
DialogId,
KbId,
}
@ -262,6 +293,7 @@ enum Dialog2Kb {
#[derive(DeriveIden)]
enum Doc2Doc {
Table,
Id,
ParentId,
Did,
}

73
src/api/doc_info.rs
View file

@ -1,15 +1,20 @@
use std::collections::HashMap;
use actix_multipart::Multipart;
use std::io::Write;
use std::slice::Chunks;
//use actix_multipart::{Multipart, MultipartError, Field};
use actix_multipart_extract::{File, Multipart, MultipartForm};
use actix_web::{get, HttpResponse, post, web};
use actix_web::web::Bytes;
use chrono::Local;
use futures_util::StreamExt;
use serde::Deserialize;
use std::io::Write;
use sea_orm::DbConn;
use crate::api::JsonResponse;
use crate::AppState;
use crate::entity::doc_info::Model;
use crate::errors::AppError;
use crate::service::doc_info::{Mutation, Query};
use serde::Deserialize;
#[derive(Debug, Deserialize)]
pub struct Params {
@ -53,42 +58,54 @@ async fn list(params: web::Json<Params>, data: web::Data<AppState>) -> Result<Ht
.body(serde_json::to_string(&json_response)?))
}
#[derive(Deserialize, MultipartForm, Debug)]
pub struct UploadForm {
#[multipart(max_size = 512MB)]
file_field: File,
uid: i64,
did: i64
}
#[post("/v1.0/upload")]
async fn upload(mut payload: Multipart, filename: web::Data<String>, did: web::Data<i64>, uid: web::Data<i64>, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
let mut size = 0;
while let Some(item) = payload.next().await {
let mut field = item.unwrap();
let filepath = format!("./uploads/{}", filename.as_str());
let mut file = web::block(|| std::fs::File::create(filepath))
.await
.unwrap()?;
while let Some(chunk) = field.next().await {
let data = chunk.unwrap();
size += data.len() as u64;
file = web::block(move || file.write_all(&data).map(|_| file))
.await
.unwrap()?;
async fn upload(payload: Multipart<UploadForm>, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
let uid = payload.uid;
async fn add_number_to_filename(file_name: String, conn:&DbConn, uid:i64) -> String {
let mut i = 0;
let mut new_file_name = file_name.to_string();
let arr: Vec<&str> = file_name.split(".").collect();
let suffix = String::from(arr[arr.len()-1]);
let preffix = arr[..arr.len()-1].join(".");
let mut docs = Query::find_doc_infos_by_name(conn, uid, new_file_name.clone()).await.unwrap();
while docs.len()>0 {
i += 1;
new_file_name = format!("{}_{}.{}", preffix, i, suffix);
docs = Query::find_doc_infos_by_name(conn, uid, new_file_name.clone()).await.unwrap();
}
new_file_name
}
let fnm = add_number_to_filename(payload.file_field.name.clone(), &data.conn, uid).await;
let _ = Mutation::create_doc_info(&data.conn, Model {
did: *did.into_inner(),
uid: *uid.into_inner(),
doc_name: filename.to_string(),
size,
std::fs::create_dir_all(format!("./upload/{}/", uid));
let filepath = format!("./upload/{}/{}-{}", payload.uid, payload.did, fnm.clone());
let mut f =std::fs::File::create(&filepath)?;
f.write(&payload.file_field.bytes)?;
let doc = Mutation::create_doc_info(&data.conn, Model {
did:Default::default(),
uid: uid,
doc_name: fnm,
size: payload.file_field.bytes.len() as i64,
kb_infos: Vec::new(),
kb_progress: 0.0,
kb_progress_msg: "".to_string(),
location: "".to_string(),
r#type: "".to_string(),
location: filepath,
r#type: "doc".to_string(),
created_at: Local::now().date_naive(),
updated_at: Local::now().date_naive(),
}).await?;
let _ = Mutation::place_doc(&data.conn, payload.did, doc.did.unwrap()).await?;
Ok(HttpResponse::Ok().body("File uploaded successfully"))
}

43
src/api/kb_info.rs
View file

@ -1,23 +1,58 @@
use std::collections::HashMap;
use actix_web::{get, HttpResponse, post, web};
use serde::Serialize;
use crate::api::JsonResponse;
use crate::AppState;
use crate::entity::kb_info;
use crate::errors::AppError;
use crate::service::kb_info::Mutation;
use crate::service::kb_info::Query;
use serde::Deserialize;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AddDocs2KbParams {
pub uid: i64,
pub dids: Vec<i64>,
pub kb_id: i64,
}
#[post("/v1.0/create_kb")]
async fn create(model: web::Json<kb_info::Model>, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
let model = Mutation::create_kb_info(&data.conn, model.into_inner()).await?;
let mut docs = Query::find_kb_infos_by_name(&data.conn, model.kb_name.to_owned()).await.unwrap();
if docs.len() >0 {
let json_response = JsonResponse {
code: 201,
err: "Duplicated name.".to_owned(),
data: ()
};
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response)?))
}else{
let model = Mutation::create_kb_info(&data.conn, model.into_inner()).await?;
let mut result = HashMap::new();
result.insert("kb_id", model.kb_id.unwrap());
let mut result = HashMap::new();
result.insert("kb_id", model.kb_id.unwrap());
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),
data: result,
};
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response)?))
}
}
#[post("/v1.0/add_docs_to_kb")]
async fn add_docs_to_kb(param: web::Json<AddDocs2KbParams>, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
let _ = Mutation::add_docs(&data.conn, param.kb_id, param.dids.to_owned()).await?;
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),
data: result,
data: (),
};
Ok(HttpResponse::Ok()

58
src/api/tag.rs Normal file
View file

@ -0,0 +1,58 @@
use std::collections::HashMap;
use actix_web::{get, HttpResponse, post, web};
use actix_web::http::Error;
use crate::api::JsonResponse;
use crate::AppState;
use crate::entity::tag_info;
use crate::service::tag_info::{Mutation, Query};
#[post("/v1.0/create_tag")]
async fn create(model: web::Json<tag_info::Model>, data: web::Data<AppState>) -> Result<HttpResponse, Error> {
let model = Mutation::create_tag(&data.conn, model.into_inner()).await.unwrap();
let mut result = HashMap::new();
result.insert("tid", model.tid.unwrap());
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),
data: result,
};
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response).unwrap()))
}
#[post("/v1.0/delete_tag")]
async fn delete(model: web::Json<tag_info::Model>, data: web::Data<AppState>) -> Result<HttpResponse, Error> {
let _ = Mutation::delete_tag(&data.conn, model.tid).await.unwrap();
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),
data: (),
};
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response).unwrap()))
}
#[get("/v1.0/tags")]
async fn list(data: web::Data<AppState>) -> Result<HttpResponse, Error> {
let tags = Query::find_tag_infos(&data.conn).await.unwrap();
let mut result = HashMap::new();
result.insert("tags", tags);
let json_response = JsonResponse {
code: 200,
err: "".to_owned(),
data: result,
};
Ok(HttpResponse::Ok()
.content_type("application/json")
.body(serde_json::to_string(&json_response).unwrap()))
}

5
src/entity/dialog2_kb.rs
View file

@ -4,10 +4,11 @@ use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Deserialize, Serialize)]
#[sea_orm(table_name = "dialog2_kb")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(primary_key, auto_increment = true)]
pub id: i64,
#[sea_orm(index)]
pub dialog_id: i64,
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(index)]
pub kb_id: i64,
}

5
src/entity/doc2_doc.rs
View file

@ -4,10 +4,11 @@ use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Deserialize, Serialize)]
#[sea_orm(table_name = "doc2_doc")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(primary_key, auto_increment = true)]
pub id: i64,
#[sea_orm(index)]
pub parent_id: i64,
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(index)]
pub did: i64,
}

4
src/entity/doc_info.rs
View file

@ -10,10 +10,10 @@ pub struct Model {
#[sea_orm(index)]
pub uid: i64,
pub doc_name: String,
pub size: u64,
pub size: i64,
#[sea_orm(column_name = "type")]
pub r#type: String,
pub kb_progress: f64,
pub kb_progress: f32,
pub kb_progress_msg: String,
pub location: String,
#[sea_orm(ignore)]

11
src/entity/kb2_doc.rs
View file

@ -4,11 +4,12 @@ use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Deserialize, Serialize)]
#[sea_orm(table_name = "kb2_doc")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(primary_key, auto_increment = true)]
pub id: i64,
#[sea_orm(index)]
pub kb_id: i64,
#[sea_orm(primary_key, auto_increment = false)]
pub uid: i64,
#[sea_orm(index)]
pub did: i64,
}
#[derive(Debug, Clone, Copy, EnumIter)]
@ -21,8 +22,8 @@ impl RelationTrait for Relation {
fn def(&self) -> RelationDef {
match self {
Self::DocInfo => Entity::belongs_to(super::doc_info::Entity)
.from(Column::Uid)
.to(super::doc_info::Column::Uid)
.from(Column::Did)
.to(super::doc_info::Column::Did)
.into(),
Self::KbInfo => Entity::belongs_to(super::kb_info::Entity)
.from(Column::KbId)

4
src/entity/kb_info.rs
View file

@ -8,8 +8,8 @@ pub struct Model {
pub kb_id: i64,
#[sea_orm(index)]
pub uid: i64,
pub kn_name: String,
pub icon: i64,
pub kb_name: String,
pub icon: i16,
#[serde(skip_deserializing)]
pub created_at: Date,

6
src/entity/mod.rs
View file

@ -1,8 +1,8 @@
pub(crate) mod user_info;
pub(crate) mod tag_info;
mod tag2_doc;
mod kb2_doc;
mod dialog2_kb;
pub(crate) mod tag2_doc;
pub(crate) mod kb2_doc;
pub(crate) mod dialog2_kb;
pub(crate) mod doc2_doc;
pub(crate) mod kb_info;
pub(crate) mod doc_info;

5
src/entity/tag2_doc.rs
View file

@ -4,10 +4,11 @@ use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Deserialize, Serialize)]
#[sea_orm(table_name = "tag2_doc")]
pub struct Model {
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(primary_key, auto_increment = true)]
pub id: i64,
#[sea_orm(index)]
pub tag_id: i64,
#[sea_orm(primary_key, auto_increment = false)]
#[sea_orm(index)]
pub uid: i64,
}

4
src/entity/tag_info.rs
View file

@ -10,8 +10,8 @@ pub struct Model {
pub uid: i64,
pub tag_name: String,
pub regx: Option<String>,
pub color: i64,
pub icon: i64,
pub color: u16,
pub icon: u16,
pub dir: Option<String>,
#[serde(skip_deserializing)]

1
src/main.rs
View file

@ -97,6 +97,7 @@ fn init(cfg: &mut web::ServiceConfig) {
cfg.service(api::kb_info::create);
cfg.service(api::kb_info::delete);
cfg.service(api::kb_info::list);
cfg.service(api::kb_info::add_docs_to_kb);
cfg.service(api::doc_info::list);
cfg.service(api::doc_info::delete);

36
src/service/doc_info.rs
View file

@ -1,5 +1,5 @@
use chrono::Local;
use sea_orm::{ActiveModelTrait, ColumnTrait, DbConn, DbErr, DeleteResult, EntityTrait, PaginatorTrait, QueryOrder};
use sea_orm::{ActiveModelTrait, ColumnTrait, DbConn, DbErr, DeleteResult, EntityTrait, PaginatorTrait, QueryOrder, Unset, Unchanged, ConditionalStatement};
use sea_orm::ActiveValue::Set;
use sea_orm::QueryFilter;
use crate::api::doc_info::Params;
@ -24,6 +24,14 @@ impl Query {
.await
}
pub async fn find_doc_infos_by_name(db: &DbConn, uid: i64, name: String) -> Result<Vec<doc_info::Model>, DbErr> {
Entity::find()
.filter(doc_info::Column::DocName.eq(name))
.filter(doc_info::Column::Uid.eq(uid))
.all(db)
.await
}
pub async fn find_doc_infos_by_params(db: &DbConn, params: Params) -> Result<Vec<doc_info::Model>, DbErr> {
// Setup paginator
let paginator = Entity::find();
@ -80,18 +88,34 @@ impl Mutation {
dids: &[i64]
) -> Result<(), DbErr> {
for did in dids {
let d = doc2_doc::Entity::find().filter(doc2_doc::Column::Did.eq(did.to_owned())).all(db).await?;
let _ = doc2_doc::ActiveModel {
parent_id: Set(dest_did),
did: Set(*did),
id: Set(d[0].id),
did: Set(did.to_owned()),
parent_id: Set(dest_did)
}
.save(db)
.await
.unwrap();
.update(db)
.await?;
}
Ok(())
}
pub async fn place_doc(
db: &DbConn,
dest_did: i64,
did: i64
) -> Result<doc2_doc::ActiveModel, DbErr> {
doc2_doc::ActiveModel {
id: Default::default(),
parent_id: Set(dest_did),
did: Set(did),
}
.save(db)
.await
}
pub async fn create_doc_info(
db: &DbConn,
form_data: doc_info::Model,

30
src/service/kb_info.rs
View file

@ -2,6 +2,7 @@ use chrono::Local;
use sea_orm::{ActiveModelTrait, ColumnTrait, DbConn, DbErr, DeleteResult, EntityTrait, PaginatorTrait, QueryFilter, QueryOrder};
use sea_orm::ActiveValue::Set;
use crate::entity::kb_info;
use crate::entity::kb2_doc;
use crate::entity::kb_info::Entity;
pub struct Query;
@ -21,6 +22,13 @@ impl Query {
.all(db)
.await
}
pub async fn find_kb_infos_by_name(db: &DbConn, name: String) -> Result<Vec<kb_info::Model>, DbErr> {
Entity::find()
.filter(kb_info::Column::KbName.eq(name))
.all(db)
.await
}
pub async fn find_kb_infos_in_page(
db: &DbConn,
@ -48,7 +56,7 @@ impl Mutation {
kb_info::ActiveModel {
kb_id: Default::default(),
uid: Set(form_data.uid.to_owned()),
kn_name: Set(form_data.kn_name.to_owned()),
kb_name: Set(form_data.kb_name.to_owned()),
icon: Set(form_data.icon.to_owned()),
created_at: Set(Local::now().date_naive()),
updated_at: Set(Local::now().date_naive()),
@ -57,6 +65,24 @@ impl Mutation {
.await
}
pub async fn add_docs(
db: &DbConn,
kb_id: i64,
doc_ids: Vec<i64>
)-> Result<(), DbErr> {
for did in doc_ids{
let _ = kb2_doc::ActiveModel {
id: Default::default(),
kb_id: Set(kb_id),
did: Set(did),
}
.save(db)
.await?;
}
Ok(())
}
pub async fn update_kb_info_by_id(
db: &DbConn,
id: i64,
@ -71,7 +97,7 @@ impl Mutation {
kb_info::ActiveModel {
kb_id: kb_info.kb_id,
uid: kb_info.uid,
kn_name: Set(form_data.kn_name.to_owned()),
kb_name: Set(form_data.kb_name.to_owned()),
icon: Set(form_data.icon.to_owned()),
created_at: Default::default(),
updated_at: Set(Local::now().date_naive()),