Skip to content

Commit c480404

Browse files
fzowlfzowl
committed
Add VoyageAI integration to Elasticsearch
- Add text, multimodal, and contextual embeddings support - Add rerank functionality - All code in services/voyageai directory - Includes comprehensive test coverage
1 parent cd926f1 commit c480404

30 files changed

+2166
-78
lines changed

‎test-voyageai-e2e.sh‎

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
#!/bin/bash
2+
3+
# VoyageAI End-to-End Test Script
4+
# This script tests the VoyageAI integration with real API calls
5+
6+
set -e # Exit on error
7+
8+
# Color codes for output
9+
RED='\033[0;31m'
10+
GREEN='\033[0;32m'
11+
YELLOW='\033[1;33m'
12+
NC='\033[0m' # No Color
13+
14+
# Configuration
15+
ES_URL="${ES_URL:-http://localhost:9200}"
16+
ES_USER="${ES_USER:-elastic}"
17+
ES_PASSWORD="${ES_PASSWORD:-changeme}"
18+
VOYAGE_API_KEY="${VOYAGE_API_KEY}"
19+
20+
# Check prerequisites
21+
if [ -z "$VOYAGE_API_KEY" ]; then
22+
echo -e "${RED}Error: VOYAGE_API_KEY environment variable is not set${NC}"
23+
echo "Usage: export VOYAGE_API_KEY='your-api-key' && ./test-voyageai-e2e.sh"
24+
exit 1
25+
fi
26+
27+
echo -e "${YELLOW}=== VoyageAI End-to-End Test Suite ===${NC}"
28+
echo "ES URL: $ES_URL"
29+
echo "ES User: $ES_USER"
30+
echo ""
31+
32+
# Test counter
33+
TESTS_RUN=0
34+
TESTS_PASSED=0
35+
TESTS_FAILED=0
36+
37+
# Function to run a test
38+
run_test() {
39+
local test_name="$1"
40+
local test_command="$2"
41+
local expected_pattern="$3"
42+
43+
TESTS_RUN=$((TESTS_RUN + 1))
44+
echo -e "${YELLOW}[TEST $TESTS_RUN] $test_name${NC}"
45+
46+
if output=$(eval "$test_command" 2>&1); then
47+
if [ -n "$expected_pattern" ]; then
48+
if echo "$output" | grep -q "$expected_pattern"; then
49+
echo -e "${GREEN}✓ PASSED${NC}"
50+
TESTS_PASSED=$((TESTS_PASSED + 1))
51+
return 0
52+
else
53+
echo -e "${RED}✗ FAILED - Expected pattern not found: $expected_pattern${NC}"
54+
echo "Output: $output"
55+
TESTS_FAILED=$((TESTS_FAILED + 1))
56+
return 1
57+
fi
58+
else
59+
echo -e "${GREEN}✓ PASSED${NC}"
60+
TESTS_PASSED=$((TESTS_PASSED + 1))
61+
return 0
62+
fi
63+
else
64+
echo -e "${RED}✗ FAILED${NC}"
65+
echo "Error: $output"
66+
TESTS_FAILED=$((TESTS_FAILED + 1))
67+
return 1
68+
fi
69+
}
70+
71+
# Cleanup function
72+
cleanup() {
73+
echo -e "\n${YELLOW}Cleaning up test endpoints...${NC}"
74+
curl -s -X DELETE "${ES_URL}/_inference/text_embedding/voyage-text-test" -u "${ES_USER}:${ES_PASSWORD}" > /dev/null 2>&1 || true
75+
curl -s -X DELETE "${ES_URL}/_inference/text_embedding/voyage-multimodal-test" -u "${ES_USER}:${ES_PASSWORD}" > /dev/null 2>&1 || true
76+
curl -s -X DELETE "${ES_URL}/_inference/text_embedding/voyage-contextual-test" -u "${ES_USER}:${ES_PASSWORD}" > /dev/null 2>&1 || true
77+
curl -s -X DELETE "${ES_URL}/_inference/rerank/voyage-rerank-test" -u "${ES_USER}:${ES_PASSWORD}" > /dev/null 2>&1 || true
78+
curl -s -X DELETE "${ES_URL}/_inference/text_embedding/voyage-35-test" -u "${ES_USER}:${ES_PASSWORD}" > /dev/null 2>&1 || true
79+
}
80+
81+
# Trap to ensure cleanup on exit
82+
trap cleanup EXIT
83+
84+
# Initial cleanup
85+
cleanup
86+
87+
echo -e "\n${YELLOW}=== Test Suite 1: Text Embeddings (voyage-3) ===${NC}\n"
88+
89+
run_test "Create text embeddings endpoint" \
90+
"curl -s -X PUT '${ES_URL}/_inference/text_embedding/voyage-text-test' \
91+
-u '${ES_USER}:${ES_PASSWORD}' \
92+
-H 'Content-Type: application/json' \
93+
-d '{
94+
\"service\": \"voyageai\",
95+
\"service_settings\": {
96+
\"api_key\": \"${VOYAGE_API_KEY}\",
97+
\"model_id\": \"voyage-3\"
98+
}
99+
}'" \
100+
'"model_id":"voyage-3"'
101+
102+
run_test "Get text embeddings endpoint configuration" \
103+
"curl -s -X GET '${ES_URL}/_inference/text_embedding/voyage-text-test' \
104+
-u '${ES_USER}:${ES_PASSWORD}'" \
105+
'"model_id":"voyage-3"'
106+
107+
run_test "Inference with text embeddings" \
108+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-text-test' \
109+
-u '${ES_USER}:${ES_PASSWORD}' \
110+
-H 'Content-Type: application/json' \
111+
-d '{
112+
\"input\": [\"Hello world\", \"This is a test\"]
113+
}'" \
114+
'"text_embedding"'
115+
116+
run_test "Inference with INGEST input type" \
117+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-text-test' \
118+
-u '${ES_USER}:${ES_PASSWORD}' \
119+
-H 'Content-Type: application/json' \
120+
-d '{
121+
\"input\": [\"Document for indexing\"],
122+
\"task_settings\": {
123+
\"input_type\": \"ingest\"
124+
}
125+
}'" \
126+
'"text_embedding"'
127+
128+
run_test "Inference with SEARCH input type" \
129+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-text-test' \
130+
-u '${ES_USER}:${ES_PASSWORD}' \
131+
-H 'Content-Type: application/json' \
132+
-d '{
133+
\"input\": [\"Search query\"],
134+
\"task_settings\": {
135+
\"input_type\": \"search\"
136+
}
137+
}'" \
138+
'"text_embedding"'
139+
140+
echo -e "\n${YELLOW}=== Test Suite 2: Multimodal Embeddings (voyage-multimodal-3) ===${NC}\n"
141+
142+
run_test "Create multimodal embeddings endpoint" \
143+
"curl -s -X PUT '${ES_URL}/_inference/text_embedding/voyage-multimodal-test' \
144+
-u '${ES_USER}:${ES_PASSWORD}' \
145+
-H 'Content-Type: application/json' \
146+
-d '{
147+
\"service\": \"voyageai\",
148+
\"service_settings\": {
149+
\"api_key\": \"${VOYAGE_API_KEY}\",
150+
\"model_id\": \"voyage-multimodal-3\"
151+
}
152+
}'" \
153+
'"model_id":"voyage-multimodal-3"'
154+
155+
run_test "Get multimodal embeddings configuration" \
156+
"curl -s -X GET '${ES_URL}/_inference/text_embedding/voyage-multimodal-test' \
157+
-u '${ES_USER}:${ES_PASSWORD}'" \
158+
'"model_id":"voyage-multimodal-3"'
159+
160+
run_test "Inference with multimodal embeddings" \
161+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-multimodal-test' \
162+
-u '${ES_USER}:${ES_PASSWORD}' \
163+
-H 'Content-Type: application/json' \
164+
-d '{
165+
\"input\": [\"Multimodal test text\"]
166+
}'" \
167+
'"text_embedding"'
168+
169+
echo -e "\n${YELLOW}=== Test Suite 3: Contextual Embeddings (voyage-context-3) ===${NC}\n"
170+
171+
run_test "Create contextual embeddings endpoint" \
172+
"curl -s -X PUT '${ES_URL}/_inference/text_embedding/voyage-contextual-test' \
173+
-u '${ES_USER}:${ES_PASSWORD}' \
174+
-H 'Content-Type: application/json' \
175+
-d '{
176+
\"service\": \"voyageai\",
177+
\"service_settings\": {
178+
\"api_key\": \"${VOYAGE_API_KEY}\",
179+
\"model_id\": \"voyage-context-3\"
180+
}
181+
}'" \
182+
'"model_id":"voyage-context-3"'
183+
184+
run_test "Get contextual embeddings configuration" \
185+
"curl -s -X GET '${ES_URL}/_inference/text_embedding/voyage-contextual-test' \
186+
-u '${ES_USER}:${ES_PASSWORD}'" \
187+
'"model_id":"voyage-context-3"'
188+
189+
run_test "Inference with contextual embeddings" \
190+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-contextual-test' \
191+
-u '${ES_USER}:${ES_PASSWORD}' \
192+
-H 'Content-Type: application/json' \
193+
-d '{
194+
\"input\": [\"Contextual test sentence\", \"Another test\"]
195+
}'" \
196+
'"text_embedding"'
197+
198+
echo -e "\n${YELLOW}=== Test Suite 4: v3.5 Models ===${NC}\n"
199+
200+
run_test "Create voyage-3.5 embeddings endpoint" \
201+
"curl -s -X PUT '${ES_URL}/_inference/text_embedding/voyage-35-test' \
202+
-u '${ES_USER}:${ES_PASSWORD}' \
203+
-H 'Content-Type: application/json' \
204+
-d '{
205+
\"service\": \"voyageai\",
206+
\"service_settings\": {
207+
\"api_key\": \"${VOYAGE_API_KEY}\",
208+
\"model_id\": \"voyage-3.5\"
209+
}
210+
}'" \
211+
'"model_id":"voyage-3.5"'
212+
213+
run_test "Inference with voyage-3.5" \
214+
"curl -s -X POST '${ES_URL}/_inference/text_embedding/voyage-35-test' \
215+
-u '${ES_USER}:${ES_PASSWORD}' \
216+
-H 'Content-Type: application/json' \
217+
-d '{
218+
\"input\": [\"Test with v3.5 model\"]
219+
}'" \
220+
'"text_embedding"'
221+
222+
echo -e "\n${YELLOW}=== Test Suite 5: Rerank (Regression Test) ===${NC}\n"
223+
224+
run_test "Create rerank endpoint" \
225+
"curl -s -X PUT '${ES_URL}/_inference/rerank/voyage-rerank-test' \
226+
-u '${ES_USER}:${ES_PASSWORD}' \
227+
-H 'Content-Type: application/json' \
228+
-d '{
229+
\"service\": \"voyageai\",
230+
\"service_settings\": {
231+
\"api_key\": \"${VOYAGE_API_KEY}\",
232+
\"model_id\": \"rerank-2\"
233+
}
234+
}'" \
235+
'"model_id":"rerank-2"'
236+
237+
run_test "Rerank inference" \
238+
"curl -s -X POST '${ES_URL}/_inference/rerank/voyage-rerank-test' \
239+
-u '${ES_USER}:${ES_PASSWORD}' \
240+
-H 'Content-Type: application/json' \
241+
-d '{
242+
\"query\": \"What is the capital of France?\",
243+
\"input\": [
244+
\"Paris is the capital of France\",
245+
\"London is the capital of England\",
246+
\"Berlin is the capital of Germany\"
247+
]
248+
}'" \
249+
'"rerank"'
250+
251+
# Print summary
252+
echo -e "\n${YELLOW}=== Test Summary ===${NC}"
253+
echo -e "Total tests run: $TESTS_RUN"
254+
echo -e "${GREEN}Tests passed: $TESTS_PASSED${NC}"
255+
if [ $TESTS_FAILED -gt 0 ]; then
256+
echo -e "${RED}Tests failed: $TESTS_FAILED${NC}"
257+
exit 1
258+
else
259+
echo -e "${GREEN}All tests passed!${NC}"
260+
exit 0
261+
fi

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/VoyageAIEmbeddingsRequestManager.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.http.sender;
9+
10+
import org.apache.logging.log4j.LogManager;
11+
import org.apache.logging.log4j.Logger;
12+
import org.elasticsearch.action.ActionListener;
13+
import org.elasticsearch.inference.InferenceServiceResults;
14+
import org.elasticsearch.threadpool.ThreadPool;
15+
import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
16+
import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
17+
import org.elasticsearch.xpack.inference.external.request.voyageai.VoyageAIEmbeddingsRequest;
18+
import org.elasticsearch.xpack.inference.external.response.voyageai.VoyageAIEmbeddingsResponseEntity;
19+
import org.elasticsearch.xpack.inference.external.voyageai.VoyageAIResponseHandler;
20+
import org.elasticsearch.xpack.inference.services.voyageai.embeddings.VoyageAIEmbeddingsModel;
21+
22+
import java.util.List;
23+
import java.util.Objects;
24+
import java.util.function.Supplier;
25+
26+
public class VoyageAIEmbeddingsRequestManager extends VoyageAIRequestManager {
27+
private static final Logger logger = LogManager.getLogger(VoyageAIEmbeddingsRequestManager.class);
28+
private static final ResponseHandler HANDLER = createEmbeddingsHandler();
29+
30+
private static ResponseHandler createEmbeddingsHandler() {
31+
return new VoyageAIResponseHandler("voyageai text embedding", VoyageAIEmbeddingsResponseEntity::fromResponse);
32+
}
33+
34+
public static VoyageAIEmbeddingsRequestManager of(VoyageAIEmbeddingsModel model, ThreadPool threadPool) {
35+
return new VoyageAIEmbeddingsRequestManager(Objects.requireNonNull(model), Objects.requireNonNull(threadPool));
36+
}
37+
38+
private final VoyageAIEmbeddingsModel model;
39+
40+
private VoyageAIEmbeddingsRequestManager(VoyageAIEmbeddingsModel model, ThreadPool threadPool) {
41+
super(threadPool, model);
42+
this.model = Objects.requireNonNull(model);
43+
}
44+
45+
@Override
46+
public void execute(
47+
InferenceInputs inferenceInputs,
48+
RequestSender requestSender,
49+
Supplier<Boolean> hasRequestCompletedFunction,
50+
ActionListener<InferenceServiceResults> listener
51+
) {
52+
List<String> docsInput = DocumentsOnlyInput.of(inferenceInputs).getInputs();
53+
VoyageAIEmbeddingsRequest request = new VoyageAIEmbeddingsRequest(docsInput, model);
54+
55+
execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
56+
}
57+
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/VoyageAIModel.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ public abstract class VoyageAIModel extends RateLimitGroupingModel {
3434
Map<String, String> tempMap = new HashMap<>();
3535
tempMap.put("voyage-3.5", "embed_medium");
3636
tempMap.put("voyage-3.5-lite", "embed_small");
37+
tempMap.put("voyage-context-3", "embed_context");
3738
tempMap.put("voyage-multimodal-3", "embed_multimodal");
3839
tempMap.put("voyage-3-large", "embed_large");
3940
tempMap.put("voyage-code-3", "embed_large");

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/VoyageAIRequestManager.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.external.http.sender;
9+
10+
import org.elasticsearch.threadpool.ThreadPool;
11+
import org.elasticsearch.xpack.inference.services.voyageai.VoyageAIModel;
12+
13+
import java.util.Map;
14+
import java.util.Objects;
15+
16+
abstract class VoyageAIRequestManager extends BaseRequestManager {
17+
private static final String DEFAULT_MODEL_FAMILY = "default_model_family";
18+
private static final Map<String, String> MODEL_TO_MODEL_FAMILY = Map.of(
19+
"voyage-multimodal-3",
20+
"embed_multimodal",
21+
"voyage-3-large",
22+
"embed_large",
23+
"voyage-code-3",
24+
"embed_large",
25+
"voyage-3",
26+
"embed_medium",
27+
"voyage-3-lite",
28+
"embed_small",
29+
"voyage-finance-2",
30+
"embed_large",
31+
"voyage-law-2",
32+
"embed_large",
33+
"voyage-code-2",
34+
"embed_large",
35+
"rerank-2",
36+
"rerank_large",
37+
"rerank-2-lite",
38+
"rerank_small"
39+
);
40+
41+
protected VoyageAIRequestManager(ThreadPool threadPool, VoyageAIModel model) {
42+
super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
43+
}
44+
45+
record RateLimitGrouping(int apiKeyHash) {
46+
public static RateLimitGrouping of(VoyageAIModel model) {
47+
Objects.requireNonNull(model);
48+
String modelId = model.getServiceSettings().modelId();
49+
String modelFamily = MODEL_TO_MODEL_FAMILY.getOrDefault(modelId, DEFAULT_MODEL_FAMILY);
50+
51+
return new RateLimitGrouping(modelFamily.hashCode());
52+
}
53+
}
54+
}

0 commit comments

Comments
 (0)