chore: pass chat metadata (#789)

* chore: pass chat metadata

* chore: clippy
This commit is contained in:
Nathan.fooo 2024-09-04 17:19:39 +08:00 committed by GitHub
parent 2b57fb7125
commit 4e384b982e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 56 additions and 38 deletions

View file

@ -1,22 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "\n SELECT content\n FROM af_chat_messages\n WHERE message_id = $1\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "content",
"type_info": "Text"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
false
]
},
"hash": "5fd78d55ed9c4b866f1ce883c5cc89d4df0d2b5daf485a9957e71112d2682f9a"
}

View file

@ -0,0 +1,28 @@
{
"db_name": "PostgreSQL",
"query": "\n SELECT content,meta_data\n FROM af_chat_messages\n WHERE message_id = $1\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "content",
"type_info": "Text"
},
{
"ordinal": 1,
"name": "meta_data",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
false,
false
]
},
"hash": "6f5d6d79587d7f7a52c920acccfe338a8c001ea30b722d3a6a1a60259d47913c"
}

View file

@ -201,11 +201,13 @@ impl AppFlowyAIClient {
chat_id: &str,
content: &str,
model: &AIModel,
metadata: Option<serde_json::Value>,
) -> Result<ChatAnswer, AIError> {
let json = ChatQuestion {
chat_id: chat_id.to_string(),
data: MessageData {
content: content.to_string(),
metadata,
},
};
let url = format!("{}/chat/message", self.url);
@ -224,12 +226,14 @@ impl AppFlowyAIClient {
&self,
chat_id: &str,
content: &str,
metadata: Option<serde_json::Value>,
model: &AIModel,
) -> Result<impl Stream<Item = Result<Bytes, AIError>>, AIError> {
let json = ChatQuestion {
chat_id: chat_id.to_string(),
data: MessageData {
content: content.to_string(),
metadata,
},
};
let url = format!("{}/chat/message/stream", self.url);
@ -247,12 +251,14 @@ impl AppFlowyAIClient {
&self,
chat_id: &str,
content: &str,
metadata: Option<serde_json::Value>,
model: &AIModel,
) -> Result<impl Stream<Item = Result<Bytes, AIError>>, AIError> {
let json = ChatQuestion {
chat_id: chat_id.to_string(),
data: MessageData {
content: content.to_string(),
metadata,
},
};
let url = format!("{}/v2/chat/message/stream", self.url);

View file

@ -20,6 +20,8 @@ pub struct ChatQuestion {
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MessageData {
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]

View file

@ -14,7 +14,7 @@ async fn create_chat_context_test() {
};
client.create_chat_text_context(context).await.unwrap();
let resp = client
.send_question(&chat_id, "Where I live?", &AIModel::GPT35)
.send_question(&chat_id, "Where I live?", &AIModel::GPT35, None)
.await
.unwrap();
// response will be something like:

View file

@ -10,7 +10,7 @@ async fn qa_test() {
client.health_check().await.unwrap();
let chat_id = uuid::Uuid::new_v4().to_string();
let resp = client
.send_question(&chat_id, "I feel hungry", &AIModel::GPT35)
.send_question(&chat_id, "I feel hungry", &AIModel::GPT35, None)
.await
.unwrap();
assert!(!resp.content.is_empty());
@ -29,7 +29,7 @@ async fn stop_stream_test() {
client.health_check().await.unwrap();
let chat_id = uuid::Uuid::new_v4().to_string();
let mut stream = client
.stream_question(&chat_id, "I feel hungry", &AIModel::GPT35)
.stream_question(&chat_id, "I feel hungry", None, &AIModel::GPT35)
.await
.unwrap();
@ -51,7 +51,7 @@ async fn stream_test() {
client.health_check().await.unwrap();
let chat_id = uuid::Uuid::new_v4().to_string();
let stream = client
.stream_question_v2(&chat_id, "I feel hungry", &AIModel::GPT35)
.stream_question_v2(&chat_id, "I feel hungry", None, &AIModel::GPT35)
.await
.unwrap();
let json_stream = JsonStream::<serde_json::Value>::new(stream);

View file

@ -608,10 +608,10 @@ pub async fn update_chat_message_meta(
pub async fn select_chat_message_content<'a, E: Executor<'a, Database = Postgres>>(
executor: E,
message_id: i64,
) -> Result<String, AppError> {
) -> Result<(String, serde_json::Value), AppError> {
let row = sqlx::query!(
r#"
SELECT content
SELECT content,meta_data
FROM af_chat_messages
WHERE message_id = $1
"#,
@ -619,5 +619,5 @@ pub async fn select_chat_message_content<'a, E: Executor<'a, Database = Postgres
)
.fetch_one(executor)
.await?;
Ok(row.content)
Ok((row.content, row.meta_data))
}

View file

@ -258,11 +258,12 @@ async fn answer_stream_handler(
req: HttpRequest,
) -> actix_web::Result<HttpResponse> {
let (_workspace_id, chat_id, question_id) = path.into_inner();
let content = chat::chat_ops::select_chat_message_content(&state.pg_pool, question_id).await?;
let (content, metadata) =
chat::chat_ops::select_chat_message_content(&state.pg_pool, question_id).await?;
let ai_model = ai_model_from_header(&req);
match state
.ai_client
.stream_question(&chat_id, &content, &ai_model)
.stream_question(&chat_id, &content, Some(metadata), &ai_model)
.await
{
Ok(answer_stream) => {
@ -290,11 +291,12 @@ async fn answer_stream_v2_handler(
req: HttpRequest,
) -> actix_web::Result<HttpResponse> {
let (_workspace_id, chat_id, question_id) = path.into_inner();
let content = chat::chat_ops::select_chat_message_content(&state.pg_pool, question_id).await?;
let (content, metadata) =
chat::chat_ops::select_chat_message_content(&state.pg_pool, question_id).await?;
let ai_model = ai_model_from_header(&req);
match state
.ai_client
.stream_question_v2(&chat_id, &content, &ai_model)
.stream_question_v2(&chat_id, &content, Some(metadata), &ai_model)
.await
{
Ok(answer_stream) => {

View file

@ -60,8 +60,9 @@ pub async fn update_chat_message(
))
})?;
// TODO(nathan): query the metadata from the database
let new_answer = ai_client
.send_question(&params.chat_id, &params.content, &ai_model)
.send_question(&params.chat_id, &params.content, &ai_model, None)
.await?;
let _answer = insert_answer_message(
pg_pool,
@ -83,9 +84,10 @@ pub async fn generate_chat_message_answer(
chat_id: &str,
ai_model: AIModel,
) -> Result<ChatMessage, AppError> {
let content = chat::chat_ops::select_chat_message_content(pg_pool, question_message_id).await?;
let (content, metadata) =
chat::chat_ops::select_chat_message_content(pg_pool, question_message_id).await?;
let new_answer = ai_client
.send_question(chat_id, &content, &ai_model)
.send_question(chat_id, &content, &ai_model, Some(metadata))
.await?;
info!("new_answer: {:?}", new_answer);
@ -210,7 +212,7 @@ pub async fn create_chat_message_stream(
ChatAuthor::new(uid, ChatAuthorType::Human),
&chat_id,
params.content.clone(),
params.metadata,
params.metadata.clone(),
).await {
Ok(question) => question,
Err(err) => {
@ -236,7 +238,7 @@ pub async fn create_chat_message_stream(
match params.message_type {
ChatMessageType::System => {}
ChatMessageType::User => {
let answer = match ai_client.send_question(&chat_id, &params.content, &ai_model).await {
let answer = match ai_client.send_question(&chat_id, &params.content, &ai_model, params.metadata).await {
Ok(response) => response,
Err(err) => {
error!("Failed to send question to AI: {}", err);