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) {
|
||||
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.
|
||||
|
|
@ -233,3 +241,25 @@ message GetDefaultTagRecommendationPromptResponse {
|
|||
// The default system prompt for tag recommendation.
|
||||
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(),
|
||||
}, 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