Skip to content

Commit 1f4d036

Browse files
committed
feat: Added AI extension support and SQL generation APIs
- Added AI extension entry and default port to the README - Added AI extension type and configuration parameters to store.ts - Added AIExtension service and SQL generation related message definitions to server.proto
1 parent e097349 commit 1f4d036

File tree

3 files changed

+113
-1
lines changed

3 files changed

+113
-1
lines changed

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ const ExtensionKindRedis = "atest-store-redis"
5757
const ExtensionKindMongoDB = "atest-store-mongodb"
5858
const ExtensionKindElasticsearch = "atest-store-elasticsearch"
5959
const ExtensionKindOpengeMini = "atest-store-opengemini"
60+
const ExtensionKindAI = "atest-ext-ai"
6061

6162
export const ExtensionKind = {
6263
ExtensionKindGit,
@@ -68,7 +69,8 @@ export const ExtensionKind = {
6869
ExtensionKindRedis,
6970
ExtensionKindMongoDB,
7071
ExtensionKindElasticsearch,
71-
ExtensionKindOpengeMini
72+
ExtensionKindOpengeMini,
73+
ExtensionKindAI
7274
}
7375

7476
const storeExtensions = [
@@ -170,6 +172,24 @@ const storeExtensions = [
170172
name: ExtensionKindOpengeMini,
171173
params: [],
172174
link: 'https://github.com/LinuxSuRen/atest-ext-store-opengemini'
175+
},
176+
{
177+
name: ExtensionKindAI,
178+
params: [{
179+
key: 'model_provider',
180+
defaultValue: 'ollama',
181+
enum: ['ollama', 'openai'],
182+
description: 'AI model provider: ollama or openai'
183+
}, {
184+
key: 'model_name',
185+
defaultValue: 'llama3.2',
186+
description: 'AI model name'
187+
}, {
188+
key: 'api_key',
189+
defaultValue: '',
190+
description: 'API key for external providers (e.g., OpenAI)'
191+
}],
192+
link: 'https://github.com/LinuxSuRen/atest-ext-ai'
173193
}
174194
] as Store[]
175195

extensions/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ Ports in extensions:
1515
| Agent | [collector](https://github.com/LinuxSuRen/atest-ext-collector) | |
1616
| Secret | [Vault](https://github.com/LinuxSuRen/api-testing-vault-extension) | |
1717
| Data | [Swagger](https://github.com/LinuxSuRen/atest-ext-data-swagger) | |
18+
| AI | [ai-extension](https://github.com/LinuxSuRen/atest-ext-ai) | 50051 |
1819

1920
## Contribute a new extension
2021

pkg/server/server.proto

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,18 @@ service ThemeExtension {
330330
}
331331
}
332332

333+
// AIExtension service provides AI-powered capabilities for atest.
334+
service AIExtension {
335+
// Translates a natural language query into an SQL query, providing
336+
// context and explanations for better accuracy and user trust.
337+
rpc GenerateSQLFromNaturalLanguage (GenerateSQLRequest) returns (GenerateSQLResponse) {
338+
option (google.api.http) = {
339+
post: "/api/v1/ai/sql/generate"
340+
body: "*"
341+
};
342+
}
343+
}
344+
333345
message Suites {
334346
map<string, Items> data = 1;
335347
}
@@ -697,3 +709,82 @@ message DataMeta {
697709
string duration = 4;
698710
repeated Pair labels = 5;
699711
}
712+
713+
// AI Extension related messages
714+
715+
// Enum for supported database dialects to ensure type safety.
716+
enum DatabaseType {
717+
// DATABASE_TYPE_UNSPECIFIED indicates a missing or unknown database type.
718+
DATABASE_TYPE_UNSPECIFIED = 0;
719+
MYSQL = 1;
720+
POSTGRESQL = 2;
721+
SQLITE = 3;
722+
}
723+
724+
// Represents an example pair of natural language to SQL for few-shot prompting.
725+
message QueryExample {
726+
// A natural language query example.
727+
string natural_language_prompt = 1;
728+
// The corresponding correct SQL query for the prompt.
729+
string sql_query = 2;
730+
}
731+
732+
// Request message for generating an SQL query from natural language.
733+
message GenerateSQLRequest {
734+
// The user's query in natural language. (e.g., "show me all users from California")
735+
string natural_language_input = 1;
736+
// Target database dialect for SQL generation.
737+
DatabaseType database_type = 2;
738+
// Optional: A list of DDL statements (e.g., CREATE TABLE …) for relevant tables.
739+
// Providing schemas helps the AI generate more accurate queries.
740+
repeated string schemas = 3;
741+
// Optional: A session identifier to maintain context across multiple requests,
742+
// enabling conversational query refinement.
743+
optional string session_id = 4;
744+
// Optional: A list of examples to guide the AI, improving accuracy for
745+
// specific or complex domains.
746+
repeated QueryExample examples = 5;
747+
}
748+
749+
// Response message containing the result of the SQL generation.
750+
message GenerateSQLResponse {
751+
// The response can be one of the following types.
752+
oneof result {
753+
// Contains the successful SQL generation details.
754+
SuccessResponse success = 1;
755+
// Contains details about why the generation failed.
756+
ErrorResponse error = 2;
757+
}
758+
}
759+
760+
// Represents a successful SQL generation.
761+
message SuccessResponse {
762+
// The generated SQL query.
763+
string sql_query = 1;
764+
// An explanation of how the AI interpreted the request and generated the SQL.
765+
// This builds user trust and aids in debugging.
766+
optional string explanation = 2;
767+
// A score between 0.0 and 1.0 indicating the AI's confidence in the generated query.
768+
optional float confidence_score = 3;
769+
}
770+
771+
// Enum for structured error codes.
772+
enum ErrorCode {
773+
ERROR_CODE_UNSPECIFIED = 0;
774+
// The input was invalid (e.g., empty prompt).
775+
INVALID_ARGUMENT = 1;
776+
// The AI model failed to translate the query.
777+
TRANSLATION_FAILED = 2;
778+
// The requested database type is not supported.
779+
UNSUPPORTED_DATABASE = 3;
780+
// An internal error occurred in the service.
781+
INTERNAL_ERROR = 4;
782+
}
783+
784+
// Represents a failed SQL generation attempt.
785+
message ErrorResponse {
786+
// A structured error code for programmatic handling.
787+
ErrorCode code = 1;
788+
// A human-readable message describing the error.
789+
string message = 2;
790+
}

0 commit comments

Comments
 (0)