feat: Added AI extension support and SQL generation APIs

console/atest-ui/src/views/store.ts

@@ -57,6 +57,7 @@ const ExtensionKindRedis = "atest-store-redis"
 const ExtensionKindMongoDB = "atest-store-mongodb"
 const ExtensionKindElasticsearch = "atest-store-elasticsearch"
 const ExtensionKindOpengeMini = "atest-store-opengemini"
+const ExtensionKindAI = "atest-ext-ai"
 
 export const ExtensionKind = {
     ExtensionKindGit,
@@ -68,7 +69,8 @@ export const ExtensionKind = {
     ExtensionKindRedis,
     ExtensionKindMongoDB,
     ExtensionKindElasticsearch,
-    ExtensionKindOpengeMini
+    ExtensionKindOpengeMini,
+    ExtensionKindAI
 }
 
 const storeExtensions = [
@@ -170,6 +172,36 @@ const storeExtensions = [
         name: ExtensionKindOpengeMini,
         params: [],
         link: 'https://github.com/LinuxSuRen/atest-ext-store-opengemini'
+    },
+    {
+        name: ExtensionKindAI,
+        params: [{
+            key: 'model_provider',
+            defaultValue: 'ollama',
+            enum: ['ollama', 'openai'],
+            description: 'AI model provider used in GenerateContentRequest.parameters for content generation'
+        }, {
+            key: 'model_name',
+            defaultValue: 'llama3.2',
+            description: 'AI model name passed to GenerateContentRequest.parameters["model_name"] for specific model selection'
+        }, {
+            key: 'api_key',
+            defaultValue: '',
+            description: 'API key for external providers, used in GenerateContentRequest.parameters["api_key"] for authentication'
+        }, {
+            key: 'base_url',
+            defaultValue: 'http://localhost:11434',
+            description: 'Base URL for AI service, used in GenerateContentRequest.parameters["base_url"] to specify service endpoint'
+        }, {
+            key: 'temperature',
+            defaultValue: '0.7',
+            description: 'Controls randomness in AI responses (0.0-1.0), passed via GenerateContentRequest.parameters["temperature"]'
+        }, {
+            key: 'max_tokens',
+            defaultValue: '2048',
+            description: 'Maximum tokens in AI response, used in GenerateContentRequest.parameters["max_tokens"] to limit output length'
+        }],
+        link: 'https://github.com/LinuxSuRen/atest-ext-ai'
     }
 ] as Store[]
 

extensions/README.md

@@ -15,6 +15,68 @@
 | Agent | [collector](https://github.com/LinuxSuRen/atest-ext-collector)           |  |
 | Secret | [Vault](https://github.com/LinuxSuRen/api-testing-vault-extension)       | |
 | Data | [Swagger](https://github.com/LinuxSuRen/atest-ext-data-swagger) | |
+| AI | [ai-extension](https://github.com/LinuxSuRen/atest-ext-ai) | 50051 |
+
+## AI Extension Interface
+
+The AI extension provides intelligent testing capabilities through gRPC interface on port 50051. It offers a unified AI interface that can handle various content generation tasks.
+
+### Available Services
+
+#### GenerateContent (Recommended)
+A general-purpose AI interface that can handle multiple content types:
+- **SQL Generation**: Convert natural language to SQL queries
+- **Test Case Generation**: Automatically generate test cases based on API specifications
+- **Mock Service Creation**: Generate mock services and responses
+- **Test Data Generation**: Create realistic test data using AI models
+- **Result Analysis**: Intelligent analysis of test results and failure patterns
+
+**Request Parameters:**
+- `prompt`: Natural language description of what to generate
+- `contentType`: Type of content ("sql", "testcase", "mock", "analysis")
+- `context`: Additional context information (schemas, API specs, etc.)
+- `sessionId`: Optional session ID for conversational context
+- `parameters`: AI model configuration (model_provider, model_name, api_key, etc.)
+
+#### GenerateSQLFromNaturalLanguage (Legacy)
+Specific interface for SQL generation, maintained for backward compatibility.
+
+### Configuration Parameters
+
+These parameters are configured in the store extension settings and passed to the AI service:
+
+- `model_provider`: AI provider ("ollama", "openai") - used in GenerateContentRequest.parameters
+- `model_name`: Specific model name (e.g., "llama3.2", "gpt-4") - passed via parameters["model_name"]
+- `api_key`: Authentication key for external providers - used in parameters["api_key"]
+- `base_url`: Service endpoint URL - specified in parameters["base_url"]
+- `temperature`: Response randomness control (0.0-1.0) - set via parameters["temperature"]
+- `max_tokens`: Maximum response length - controlled by parameters["max_tokens"]
+
+### Usage Examples
+
+```protobuf
+// Generate SQL query
+GenerateContentRequest {
+  prompt: "Show me all users from California"
+  contentType: "sql"
+  context: {
+    "schema": "CREATE TABLE users (id INT, name VARCHAR(100), state VARCHAR(50))"
+  }
+  parameters: {
+    "model_provider": "ollama"
+    "model_name": "llama3.2"
+  }
+}
+
+// Generate test case
+GenerateContentRequest {
+  prompt: "Create a test case for user registration API"
+  contentType: "testcase"
+  context: {
+    "api_spec": "POST /api/users {name, email, password}"
+  }
+}
+```
 
 ## Contribute a new extension
 

pkg/server/server.proto

@@ -330,6 +330,27 @@ service ThemeExtension {
     }
 }
 
+// AIExtension service provides AI-powered capabilities for atest.
+service AIExtension {
+    // Generates content based on natural language prompts with context.
+    // This is a general-purpose AI interface that can handle SQL generation,
+    // test case writing, mock service creation, and other AI-powered tasks.
+    rpc GenerateContent (GenerateContentRequest) returns (GenerateContentResponse) {
+        option (google.api.http) = {
+            post: "/api/v1/ai/generate"
+            body: "*"
+        };
+    }
+
+    // Legacy SQL generation endpoint for backward compatibility
+    rpc GenerateSQLFromNaturalLanguage (GenerateSQLRequest) returns (GenerateSQLResponse) {
+        option (google.api.http) = {
+            post: "/api/v1/ai/sql/generate"
+            body: "*"
+        };
+    }
+}
+
 message Suites {
   map<string, Items> data = 1;
 }
@@ -697,3 +718,122 @@ message DataMeta {
     string duration = 4;
     repeated Pair labels = 5;
 }
+
+// AI Extension related messages
+
+// Enum for supported database dialects to ensure type safety.
+enum DatabaseType {
+  // DATABASE_TYPE_UNSPECIFIED indicates a missing or unknown database type.
+  DATABASE_TYPE_UNSPECIFIED = 0;
+  MYSQL = 1;
+  POSTGRESQL = 2;
+  SQLITE = 3;
+}
+
+// Represents an example pair of natural language to SQL for few-shot prompting.
+message QueryExample {
+  // A natural language query example.
+  string naturalLanguagePrompt = 1;
+  // The corresponding correct SQL query for the prompt.
+  string sqlQuery = 2;
+}
+
+// Request message for generating content based on natural language prompts.
+message GenerateContentRequest {
+  // The user's prompt in natural language.
+  string prompt = 1;
+  // The type of content to generate (e.g., "sql", "testcase", "mock").
+  string contentType = 2;
+  // Context information to help with generation.
+  map<string, string> context = 3;
+  // Optional: A session identifier to maintain context across multiple requests.
+  optional string sessionId = 4;
+  // Optional: Additional parameters specific to the content type.
+  map<string, string> parameters = 5;
+}
+
+// Request message for generating an SQL query from natural language.
+message GenerateSQLRequest {
+  // The user's query in natural language. (e.g., "show me all users from California")
+  string naturalLanguageInput = 1;
+  // Target database dialect for SQL generation.
+  DatabaseType databaseType = 2;
+  // Optional: A list of DDL statements (e.g., CREATE TABLE …) for relevant tables.
+  // Providing schemas helps the AI generate more accurate queries.
+  repeated string schemas = 3;
+  // Optional: A session identifier to maintain context across multiple requests,
+  // enabling conversational query refinement.
+  optional string sessionId = 4;
+  // Optional: A list of examples to guide the AI, improving accuracy for
+  // specific or complex domains.
+  repeated QueryExample examples = 5;
+}
+
+// Response message containing the result of content generation.
+message GenerateContentResponse {
+  // The response can be one of the following types.
+  oneof result {
+    // Contains the successful content generation details.
+    ContentSuccessResponse success = 1;
+    // Contains details about why the generation failed.
+    ErrorResponse error = 2;
+  }
+}
+
+// Response message containing the result of the SQL generation.
+message GenerateSQLResponse {
+  // The response can be one of the following types.
+  oneof result {
+    // Contains the successful SQL generation details.
+    SuccessResponse success = 1;
+    // Contains details about why the generation failed.
+    ErrorResponse error = 2;
+  }
+}
+
+// Represents a successful content generation.
+message ContentSuccessResponse {
+  // The generated content.
+  string content = 1;
+  // The type of content that was generated.
+  string contentType = 2;
+  // An explanation of how the AI interpreted the request and generated the content.
+  // This builds user trust and aids in debugging.
+  optional string explanation = 3;
+  // A score between 0.0 and 1.0 indicating the AI's confidence in the generated content.
+  optional float confidenceScore = 4;
+  // Additional metadata about the generation process.
+  map<string, string> metadata = 5;
+}
+
+// Represents a successful SQL generation.
+message SuccessResponse {
+  // The generated SQL query.
+  string sqlQuery = 1;
+  // An explanation of how the AI interpreted the request and generated the SQL.
+  // This builds user trust and aids in debugging.
+  optional string explanation = 2;
+  // A score between 0.0 and 1.0 indicating the AI's confidence in the generated query.
+  optional float confidenceScore = 3;
+}
+
+// Enum for structured error codes.
+enum ErrorCode {
+  ERROR_CODE_UNSPECIFIED = 0;
+  // The input was invalid (e.g., empty prompt).
+  INVALID_ARGUMENT = 1;
+  // The AI model failed to translate the query.
+  TRANSLATION_FAILED = 2;
+  // The requested database type is not supported.
+  UNSUPPORTED_DATABASE = 3;
+  // An internal error occurred in the service.
+  INTERNAL_ERROR = 4;
+}
+
+// Represents a failed SQL generation attempt.
+message ErrorResponse {
+  // A structured error code for programmatic handling.
+  ErrorCode code = 1;
+  // A human-readable message describing the error.
+  string message = 2;
+}