From 8fb7727fb8157e58c877b5ad116b0949f5d39cc6 Mon Sep 17 00:00:00 2001
From: ChaoLiu <chaoliu719@gmail.com>
Date: Mon, 18 Aug 2025 17:14:24 +0800
Subject: [PATCH] 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>
---
 proto/api/v1/workspace_service.proto      | 30 +++++++++
 server/router/api/v1/workspace_service.go | 82 +++++++++++++++++++++++
 2 files changed, 112 insertions(+)

diff --git a/proto/api/v1/workspace_service.proto b/proto/api/v1/workspace_service.proto
index 75737ae12..36afe175e 100644
--- a/proto/api/v1/workspace_service.proto
+++ b/proto/api/v1/workspace_service.proto
@@ -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;
+}
diff --git a/server/router/api/v1/workspace_service.go b/server/router/api/v1/workspace_service.go
index 10bd688b2..e2f6b87f7 100644
--- a/server/router/api/v1/workspace_service.go
+++ b/server/router/api/v1/workspace_service.go
@@ -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
+}