mirror of https://github.com/usememos/memos.git
feat: add AI connection test API and protobuf definitions
- Add TestAiConnection RPC endpoint with request/response messages - Implement connection testing with OpenAI API validation - Add proper authentication and error handling for test endpoint Signed-off-by: ChaoLiu <chaoliu719@gmail.com>
This commit is contained in:
parent
fcd85d9651
commit
8fb7727fb8
|
|
@ -35,6 +35,14 @@ service WorkspaceService {
|
||||||
rpc GetDefaultTagRecommendationPrompt(GetDefaultTagRecommendationPromptRequest) returns (GetDefaultTagRecommendationPromptResponse) {
|
rpc GetDefaultTagRecommendationPrompt(GetDefaultTagRecommendationPromptRequest) returns (GetDefaultTagRecommendationPromptResponse) {
|
||||||
option (google.api.http) = {get: "/api/v1/workspace/ai/tag-recommendation/default-prompt"};
|
option (google.api.http) = {get: "/api/v1/workspace/ai/tag-recommendation/default-prompt"};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Tests AI API connection and configuration.
|
||||||
|
rpc TestAiConnection(TestAiConnectionRequest) returns (TestAiConnectionResponse) {
|
||||||
|
option (google.api.http) = {
|
||||||
|
post: "/api/v1/workspace/ai/test-connection"
|
||||||
|
body: "*"
|
||||||
|
};
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Workspace profile message containing basic workspace information.
|
// Workspace profile message containing basic workspace information.
|
||||||
|
|
@ -233,3 +241,25 @@ message GetDefaultTagRecommendationPromptResponse {
|
||||||
// The default system prompt for tag recommendation.
|
// The default system prompt for tag recommendation.
|
||||||
string system_prompt = 1;
|
string system_prompt = 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Request message for TestAiConnection method.
|
||||||
|
message TestAiConnectionRequest {
|
||||||
|
// base_url is the base URL for AI API.
|
||||||
|
string base_url = 1;
|
||||||
|
// api_key is the API key for AI service.
|
||||||
|
string api_key = 2;
|
||||||
|
// model is the AI model to use.
|
||||||
|
string model = 3;
|
||||||
|
// timeout_seconds is the timeout for AI requests in seconds.
|
||||||
|
int32 timeout_seconds = 4;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Response message for TestAiConnection method.
|
||||||
|
message TestAiConnectionResponse {
|
||||||
|
// success indicates whether the connection test was successful.
|
||||||
|
bool success = 1;
|
||||||
|
// message provides additional information about the test result.
|
||||||
|
string message = 2;
|
||||||
|
// model_info contains information about the tested model (if successful).
|
||||||
|
string model_info = 3;
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -378,3 +378,85 @@ func (s *APIV1Service) GetDefaultTagRecommendationPrompt(ctx context.Context, _
|
||||||
SystemPrompt: ai.GetDefaultSystemPrompt(),
|
SystemPrompt: ai.GetDefaultSystemPrompt(),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestAiConnection tests the AI API connection and configuration.
|
||||||
|
func (s *APIV1Service) TestAiConnection(ctx context.Context, request *v1pb.TestAiConnectionRequest) (*v1pb.TestAiConnectionResponse, error) {
|
||||||
|
// Check permissions - only host can test AI connection
|
||||||
|
user, err := s.GetCurrentUser(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return nil, status.Errorf(codes.Internal, "failed to get current user: %v", err)
|
||||||
|
}
|
||||||
|
if user == nil || user.Role != store.RoleHost {
|
||||||
|
return nil, status.Errorf(codes.PermissionDenied, "permission denied")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate request
|
||||||
|
if request.BaseUrl == "" {
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: false,
|
||||||
|
Message: "Base URL is required",
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
if request.ApiKey == "" {
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: false,
|
||||||
|
Message: "API Key is required",
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
if request.Model == "" {
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: false,
|
||||||
|
Message: "Model is required",
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create AI config for testing
|
||||||
|
config := &ai.Config{
|
||||||
|
Enabled: true,
|
||||||
|
BaseURL: request.BaseUrl,
|
||||||
|
APIKey: request.ApiKey,
|
||||||
|
Model: request.Model,
|
||||||
|
TimeoutSeconds: int(request.TimeoutSeconds),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Default timeout if not specified
|
||||||
|
if config.TimeoutSeconds <= 0 {
|
||||||
|
config.TimeoutSeconds = 10
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create AI client
|
||||||
|
client, err := ai.NewClient(config)
|
||||||
|
if err != nil {
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: false,
|
||||||
|
Message: fmt.Sprintf("Failed to create AI client: %v", err),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test with a simple chat request
|
||||||
|
chatRequest := &ai.ChatRequest{
|
||||||
|
Messages: []ai.Message{
|
||||||
|
{
|
||||||
|
Role: "user",
|
||||||
|
Content: "Hello, please respond with 'AI connection test successful'",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
MaxTokens: 50,
|
||||||
|
Temperature: 0.1,
|
||||||
|
}
|
||||||
|
|
||||||
|
response, err := client.Chat(ctx, chatRequest)
|
||||||
|
if err != nil {
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: false,
|
||||||
|
Message: fmt.Sprintf("AI API test failed: %v", err),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test successful
|
||||||
|
return &v1pb.TestAiConnectionResponse{
|
||||||
|
Success: true,
|
||||||
|
Message: "AI connection test successful",
|
||||||
|
ModelInfo: fmt.Sprintf("Model: %s, Response: %s", request.Model, response.Content),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue