test(elasticsearch): add integration tests for semantic vector search

.ci/integration.cloudbuild.yaml

@@ -1159,7 +1159,7 @@ steps:
     env:
       - "GOPATH=/gopath"
       - "SERVICE_ACCOUNT_EMAIL=$SERVICE_ACCOUNT_EMAIL"
-    secretEnv: ["CLIENT_ID", "ELASTICSEARCH_USER", "ELASTICSEARCH_PASS", "ELASTICSEARCH_HOST"]
+    secretEnv: ["CLIENT_ID", "ELASTICSEARCH_USER", "ELASTICSEARCH_PASS", "ELASTICSEARCH_HOST", "API_KEY"]
     volumes:
       - name: "go"
         path: "/gopath"
@@ -1173,7 +1173,9 @@ steps:
           .ci/test_with_coverage.sh \
             "Elasticsearch" \
             elasticsearch \
-            elasticsearch
+            elasticsearch \
+            "" \
+            "API_KEY"
         else
           echo "No relevant changes for Elasticsearch. Skipping shard."
           exit 0

tests/elasticsearch/elasticsearch_integration_test.go

@@ -74,6 +74,8 @@
 
 	toolsConfig := getElasticsearchToolsConfig(sourceConfig, ElasticsearchToolType, paramToolStatement, idParamToolStatement, nameParamToolStatement, arrayParamToolStatement, authToolStatement)
 
+	toolsConfig = addSemanticSearchConfig(t, toolsConfig, index)
+
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsConfig, args...)
 	if err != nil {
 		t.Fatalf("failed to start cmd: %v", err)
@@ -113,12 +115,44 @@
 									"email": "%s"
 								}`, tests.ServiceAccountEmail)
 
+	// Create index with mapping for vector search
+	mapping := `{
+		"mappings": {
+			"properties": {
+				"embedding": {
+					"type": "dense_vector",
+					"dims": 768,
+					"index": true,
+					"similarity": "cosine"
+				}
+			}
+		}
+	}`
+	_, err = esapi.IndicesCreateRequest{
+		Index: index,
+		Body:  strings.NewReader(mapping),
+	}.Do(ctx, esClient)
+	if err != nil {
+		t.Fatalf("error creating index: %s", err)
+	}
+
+	vectorSize := 768
+	var sb strings.Builder
+	sb.WriteString("[")
+	if vectorSize > 0 {
+		sb.WriteString("0.1")
+		sb.WriteString(strings.Repeat(", 0.1", vectorSize-1))
+	}
+	sb.WriteString("]")
+	semanticDoc := fmt.Sprintf(`{"id": 5, "name": "Semantic", "embedding": %s}`, sb.String())
+
 	// Index sample documents
 	sampleDocs := []string{
 		alice,
 		`{"id": 2, "name": "Jane", "email": "janedoe@gmail.com"}`,
 		`{"id": 3, "name": "Sid"}`,
 		`{"id": 4, "name": "null"}`,
+		semanticDoc,
 	}
 	for _, doc := range sampleDocs {
 		res, err := esapi.IndexRequest{
@@ -146,6 +180,10 @@
 		tests.WithNullWant(wants.Null),
 	)
 	tests.RunMCPToolCallMethod(t, wants.McpMyFailTool, wants.McpSelect1, tests.WithMcpMyToolId3NameAliceWant(wants.McpMyToolId3NameAlice))
+
+	// Semantic search tests
+	semanticSearchWant := `[{"id":5,"name":"Semantic","name.keyword":"Semantic"}]`
+	tests.RunSemanticSearchToolInvokeTest(t, "", "", semanticSearchWant)
 }
 
 func getElasticsearchQueries(index string) (string, string, string, string, string) {
@@ -307,3 +345,58 @@
 	}
 	return toolsFile
 }
+
+func addSemanticSearchConfig(t *testing.T, config map[string]any, index string) map[string]any {
+	config["embeddingModels"] = map[string]any{
+		"gemini_model": map[string]any{
+			"kind":      "gemini",
+			"model":     "gemini-embedding-001",
+			"apiKey":    tests.ApiKey,
+			"dimension": 768,
+		},
+	}
+
+	tools, ok := config["tools"].(map[string]any)
+	if !ok {
+		t.Fatalf("unable to get tools from config")
+	}
+
+	tools["insert_docs"] = map[string]any{
+		"kind":        ElasticsearchToolType,
+		"source":      "my-instance",
+		"description": "Stores content and its vector embedding into the documents table.",
+		"query":   fmt.Sprintf("FROM %s | WHERE name == ?content OR name == ?text_to_embed | LIMIT 0", index),
+		"parameters": []any{
+			map[string]any{
+				"name":        "content",
+				"type":        "string",
+				"description": "The text content associated with the vector.",
+			},
+			map[string]any{
+				"name":           "text_to_embed",
+				"type":           "string",
+				"description":    "The text content used to generate the vector.",
+				"embeddedBy":     "gemini_model",
+				"valueFromParam": "content",
+			},
+		},
+	}
+
+	tools["search_docs"] = map[string]any{
+		"kind":        ElasticsearchToolType,
+		"source":      "my-instance",
+		"description": "Finds the most semantically similar document to the query vector.",
+		"query":   fmt.Sprintf("FROM %s | WHERE embedding IS NOT NULL | EVAL score = COSINE_SIMILARITY(embedding, ?query) | SORT score DESC | LIMIT 1 | KEEP id, name", index),
+		"parameters": []any{
+			map[string]any{
+				"name":        "query",
+				"type":        "string",
+				"description": "The text content to search for.",
+				"embeddedBy":  "gemini_model",
+			},
+		},
+	}
+
+	config["tools"] = tools
+	return config
+}

tests/embedding.go

@@ -33,7 +33,7 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 )
 
-var apiKey = os.Getenv("API_KEY")
+var ApiKey = os.Getenv("API_KEY")
 
 // AddSemanticSearchConfig adds embedding models and semantic search tools to the config
 // with configurable tool kind and SQL statements.
@@ -42,7 +42,7 @@ func AddSemanticSearchConfig(t *testing.T, config map[string]any, toolKind, inse
 		"gemini_model": map[string]any{
 			"kind":      "gemini",
 			"model":     "gemini-embedding-001",
-			"apiKey":    apiKey,
+			"apiKey":    ApiKey,
 			"dimension": 768,
 		},
 	}