chore: return model list (#1206)

This commit is contained in:
Nathan.fooo 2025-02-01 12:45:20 +08:00 committed by GitHub
parent fe6d11eb25
commit 18b1386bc2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 73 additions and 4 deletions

View File

@ -1,6 +1,6 @@
use crate::dto::{
AIModel, CalculateSimilarityParams, ChatAnswer, ChatQuestion, CompleteTextParams,
CreateChatContext, Document, LocalAIConfig, MessageData, QuestionMetadata,
CreateChatContext, Document, LocalAIConfig, MessageData, ModelList, QuestionMetadata,
RepeatedLocalAIPackage, RepeatedRelatedQuestion, ResponseFormat, SearchDocumentsRequest,
SimilarityResponse, SummarizeRowResponse, TranslateRowData, TranslateRowResponse,
};
@ -284,6 +284,15 @@ impl AppFlowyAIClient {
.into_data()
}
pub async fn get_model_list(&self) -> Result<ModelList, AIError> {
let url = format!("{}/model/list", self.url);
let resp = self.async_http_client(Method::GET, &url)?.send().await?;
AIResponse::<ModelList>::from_reqwest_response(resp)
.await?
.into_data()
}
pub async fn calculate_similarity(
&self,
params: CalculateSimilarityParams,

View File

@ -414,6 +414,11 @@ pub struct LocalAIConfig {
pub plugin: AppFlowyOfflineAI,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelList {
pub models: Vec<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CreateChatContext {
pub chat_id: String,

View File

@ -1,3 +1,4 @@
mod completion_test;
mod context_test;
mod model_config_test;
mod qa_test;

View File

@ -0,0 +1,8 @@
use crate::appflowy_ai_client;
#[tokio::test]
async fn get_model_list_test() {
let client = appflowy_ai_client();
let models = client.get_model_list().await.unwrap().models;
assert!(models.len() >= 5, "models.len() = {}", models.len());
}

View File

@ -4,8 +4,8 @@ use bytes::Bytes;
use futures_core::Stream;
use reqwest::Method;
use shared_entity::dto::ai_dto::{
CompleteTextParams, LocalAIConfig, SummarizeRowParams, SummarizeRowResponse, TranslateRowParams,
TranslateRowResponse,
CompleteTextParams, LocalAIConfig, ModelList, SummarizeRowParams, SummarizeRowResponse,
TranslateRowParams, TranslateRowResponse,
};
use shared_entity::response::{AppResponse, AppResponseError};
use std::time::Duration;
@ -96,4 +96,18 @@ impl Client {
.await?
.into_data()
}
#[instrument(level = "debug", skip_all, err)]
pub async fn get_model_list(&self, workspace_id: &str) -> Result<ModelList, AppResponseError> {
let url = format!("{}/api/ai/{workspace_id}/model/list", self.base_url);
let resp = self
.http_client_with_auth(Method::GET, &url)
.await?
.send()
.await?;
log_request_id(&resp);
AppResponse::<ModelList>::from_response(resp)
.await?
.into_data()
}
}

View File

@ -5,7 +5,7 @@ use actix_web::web::{Data, Json};
use actix_web::{web, HttpRequest, HttpResponse, Scope};
use app_error::AppError;
use appflowy_ai_client::dto::{
CalculateSimilarityParams, LocalAIConfig, SimilarityResponse, TranslateRowParams,
CalculateSimilarityParams, LocalAIConfig, ModelList, SimilarityResponse, TranslateRowParams,
TranslateRowResponse,
};
@ -28,6 +28,7 @@ pub fn ai_completion_scope() -> Scope {
.service(
web::resource("/calculate_similarity").route(web::post().to(calculate_similarity_handler)),
)
.service(web::resource("/model/list").route(web::get().to(model_list_handler)))
}
async fn stream_complete_text_handler(
@ -126,6 +127,7 @@ struct ConfigQuery {
platform: String,
app_version: Option<String>,
}
#[instrument(level = "debug", skip_all, err)]
async fn local_ai_config_handler(
state: web::Data<AppState>,
@ -165,3 +167,15 @@ async fn calculate_similarity_handler(
.map_err(|err| AppError::AIServiceUnavailable(err.to_string()))?;
Ok(AppResponse::Ok().with_data(response).into())
}
#[instrument(level = "debug", skip_all, err)]
async fn model_list_handler(
state: web::Data<AppState>,
) -> actix_web::Result<Json<AppResponse<ModelList>>> {
let model_list = state
.ai_client
.get_model_list()
.await
.map_err(|err| AppError::AIServiceUnavailable(err.to_string()))?;
Ok(AppResponse::Ok().with_data(model_list).into())
}

View File

@ -484,6 +484,24 @@ async fn get_question_message_test() {
assert_eq!(find_question.reply_message_id.unwrap(), answer.message_id);
}
#[tokio::test]
async fn get_model_list_test() {
if !ai_test_enabled() {
return;
}
let test_client = TestClient::new_user().await;
let workspace_id = test_client.workspace_id().await;
let models = test_client
.api_client
.get_model_list(&workspace_id)
.await
.unwrap()
.models;
assert!(!models.is_empty());
assert!(models.len() >= 5, "models.len() = {}", models.len());
println!("models: {:?}", models);
}
async fn collect_answer(mut stream: QuestionStream) -> String {
let mut answer = String::new();
while let Some(value) = stream.next().await {