@@ -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+
333345message 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