Skip to content

Commit 75ff054

Browse files
zetxqxzetxqx
authored
vertexai parser add streamRawPredictServiceMethod and rawPredict support (#1261)
* add streamRawPredictServiceMethod support for vertex Signed-off-by: bobzetian <bobzetian@google.com> * add document and did some testing Signed-off-by: bobzetian <bobzetian@google.com> * reinfe readme Signed-off-by: bobzetian <bobzetian@google.com> * add rawPredict support Signed-off-by: bobzetian <bobzetian@google.com> * fix(sessionaffinity): resolve imports and signature mismatches Signed-off-by: bobzetian <bobzetian@google.com> * refactor(vertexai): consolidate raw predict parsing logic Consolidate streamRawPredictServiceMethod and rawPredictServiceMethod parsing logic in Vertex AI parser to reduce code duplication. Introduced a local `rawRequest` interface to enforce compile-time type safety for protobuf messages containing HttpBody data. Done by: gemini-cli TEST=go test ./pkg/epp/framework/plugins/requesthandling/parsers/vertexai/... && make lint && make test-unit Signed-off-by: bobzetian <bobzetian@google.com> * refactor(vertexai): further consolidate ChatCompletions parsing logic Generalized the parser helper method to `parseVertexRequest` to handle ChatCompletionsRequest as well, eliminating all remaining boilerplate for gRPC payload extraction in the Vertex AI parser. Also removed the now unused `errors` import. Done by: gemini-cli TEST=go test ./pkg/epp/framework/plugins/requesthandling/parsers/vertexai/... && make lint && make test-unit Signed-off-by: bobzetian <bobzetian@google.com> * refactor(vertexai): rename rawRequest interface to httpBodyMessage Renamed the `rawRequest` interface to `httpBodyMessage` to better reflect its generic capability and usage across all Vertex AI gRPC requests that wrap an HttpBody. Done by: gemini-cli TEST=go test ./pkg/epp/framework/plugins/requesthandling/parsers/vertexai/... && make lint && make test-unit Signed-off-by: bobzetian <bobzetian@google.com> --------- Signed-off-by: bobzetian <bobzetian@google.com>
1 parent d25f9fe commit 75ff054

3 files changed

Lines changed: 179 additions & 32 deletions

File tree

‎pkg/epp/framework/plugins/requesthandling/parsers/vertexai/README.md‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# Vertex AI Parser Plugin
2+
3+
The Vertex AI Parser plugin (`vertexai-parser`) implements the `fwkrh.Parser` interface for the Vertex AI gRPC API. It enables the `llm-d-router` to parse and route incoming Vertex AI gRPC requests and extract metrics from their responses.
4+
5+
## How it Works
6+
7+
Vertex AI's flexible prediction services wrap standard OpenAI-compatible payloads inside gRPC protobuf envelopes. Instead of fully re-implementing the JSON parsing, the Vertex AI parser acts as a wrapper that:
8+
1. **Parses the outer gRPC frame** to extract the raw protobuf payload.
9+
2. **Unmarshals the protobuf request** into the corresponding Vertex AI message type.
10+
3. **Extracts the inner JSON payload** from the embedded `HttpBody`.
11+
4. **Delegates to the OpenAI parser** by cloning the HTTP headers, setting the path to an OpenAI-compatible route, and passing the extracted JSON.
12+
5. **Packages the results**, including the original protobuf request in the `Payload` metadata so downstream plugins can access the full gRPC context.
13+
14+
## Supported gRPC Methods
15+
16+
The parser automatically matches incoming requests based on the `:path` header suffix:
17+
18+
| gRPC Method Suffix | Protocol Message | Inner OpenAI Path | Description |
19+
| :--- | :--- | :--- | :--- |
20+
| `PredictionService/ChatCompletions` | `aiplatformpb.ChatCompletionsRequest` | `/chat/completions` | OpenAI-compatible Chat Completions service. |
21+
| `PredictionService/StreamRawPredict` | `aiplatformpb.StreamRawPredictRequest` | `/responses` | Streaming raw prediction service. |
22+
| `PredictionService/RawPredict` | `aiplatformpb.RawPredictRequest` | `/responses` | Non-streaming raw prediction service. |
23+
24+
## Response Parsing
25+
26+
For responses, the parser:
27+
1. Extracts the gRPC payload and unmarshals it into `httpbody.HttpBody`.
28+
2. Extracts the raw JSON data from the body.
29+
3. Delegates to the OpenAI parser to extract token usage metrics (`prompt_tokens`, `completion_tokens`, `total_tokens`) which are then recorded by the router.
30+
31+
## Configuration
32+
33+
To enable the Vertex AI parser, configure it in your `EndpointPickerConfig` under the `requestHandler` section:
34+
35+
```yaml
36+
apiVersion: llm-d.ai/v1alpha1
37+
kind: EndpointPickerConfig
38+
requestHandler:
39+
parser:
40+
pluginRef: vertexai-parser
41+
plugins:
42+
- name: vertexai-parser
43+
type: vertexai-parser
44+
```
45+
46+
No additional parameters are required.

‎pkg/epp/framework/plugins/requesthandling/parsers/vertexai/vertexai.go‎

Lines changed: 61 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ package vertexai
1919
import (
2020
"context"
2121
"encoding/json"
22-
"errors"
2322
"fmt"
2423
"maps"
2524
"strings"
@@ -38,8 +37,24 @@ import (
3837
const (
3938
VertexAIParserType = "vertexai-parser"
4039

41-
chatCompletionsMethod = "PredictionService/ChatCompletions"
40+
// chatCompletionsMethod is the gRPC method path suffix for Vertex AI's OpenAI-compatible
41+
// ChatCompletions service (maps to aiplatformpb.ChatCompletionsRequest).
42+
// See: https://github.com/googleapis/googleapis/blob/89c3153888201c9e80bc5ec78d6ffca0debe6b52/google/cloud/aiplatform/v1beta1/prediction_service.proto#L234 for definition.
43+
chatCompletionsMethod = "PredictionService/ChatCompletions"
44+
// streamRawPredictServiceMethod is the gRPC method path suffix for Vertex AI's flexible,
45+
// low-level raw prediction streaming service (maps to aiplatformpb.StreamRawPredictRequest).
46+
// See: https://github.com/googleapis/googleapis/blob/89c3153888201c9e80bc5ec78d6ffca0debe6b52/google/cloud/aiplatform/v1beta1/prediction_service.proto#L84 for definition.
47+
streamRawPredictServiceMethod = "PredictionService/StreamRawPredict"
48+
// rawPredictServiceMethod is the gRPC method path suffix for Vertex AI's flexible,
49+
// low-level raw prediction service (maps to aiplatformpb.RawPredictRequest).
50+
// See: https://github.com/googleapis/googleapis/blob/89c3153888201c9e80bc5ec78d6ffca0debe6b52/google/cloud/aiplatform/v1beta1/prediction_service.proto#L71 for definition.
51+
rawPredictServiceMethod = "PredictionService/RawPredict"
52+
// openAIChatCompletionsPath is the standard OpenAI endpoint path for Chat Completions,
53+
// used to route extracted JSON payloads to the OpenAI parser.
4254
openAIChatCompletionsPath = "/chat/completions"
55+
// openAIResponsesPath is the OpenAI-compatible path for raw responses, used to route
56+
// extracted StreamRawPredict JSON payloads to the OpenAI parser.
57+
openAIResponsesPath = "/responses"
4358
)
4459

4560
// compile-time type validation
@@ -85,40 +100,18 @@ func (p *VertexAIParser) WithName(name string) *VertexAIParser {
85100
// ChatCompletionsRequest protobuf message. This message embeds an HttpBody containing the
86101
// actual request payload as an OpenAI-compatible JSON string. The parser extracts this JSON
87102
// data and delegates the parsing to the OpenAI parser.
88-
// See: https://github.com/googleapis/googleapis/blob/89c3153888201c9e80bc5ec78d6ffca0debe6b52/google/cloud/aiplatform/v1beta1/prediction_service.proto#L235
89103
func (p *VertexAIParser) ParseRequest(ctx context.Context, body []byte, headers map[string]string) (*fwkrh.ParseResult, error) {
90104
path := headers[parsers.MethodPathKey]
91105

92106
switch {
93107
case strings.HasSuffix(path, chatCompletionsMethod):
94-
parsedPayload, err := grpcutil.ParseGrpcPayload(body)
95-
if err != nil {
96-
return nil, fmt.Errorf("invalid or unsupported gRPC payload: %w", err)
97-
}
98-
99-
req := &aiplatformpb.ChatCompletionsRequest{}
100-
if err := proto.Unmarshal(parsedPayload, req); err != nil {
101-
return nil, fmt.Errorf("unmarshaling ChatCompletionsRequest: %w", err)
102-
}
103-
104-
httpBody := req.GetHttpBody()
105-
if httpBody == nil {
106-
return nil, errors.New("ChatCompletionsRequest has no HttpBody")
107-
}
108-
jsonBytes := httpBody.GetData()
109-
110-
// Use OpenAI parser to parse the JSON payload
111-
// Clone headers and set path to /chat/completions to make OpenAI parser recognize it
112-
headersCopy := maps.Clone(headers)
113-
headersCopy[parsers.MethodPathKey] = openAIChatCompletionsPath
114-
parseResult, err := p.openAIParser.ParseRequest(ctx, jsonBytes, headersCopy)
115-
if err != nil {
116-
return nil, fmt.Errorf("parsing ChatCompletionsRequest: %w", err)
117-
}
118-
119-
inferenceRequestBody := parseResult.Body
120-
inferenceRequestBody.Payload = fwkrh.PayloadProto{Message: req}
121-
return &fwkrh.ParseResult{Body: inferenceRequestBody, Skip: parseResult.Skip}, nil
108+
return p.parseVertexRequest(ctx, body, headers, &aiplatformpb.ChatCompletionsRequest{}, "ChatCompletionsRequest", openAIChatCompletionsPath)
109+
110+
case strings.HasSuffix(path, streamRawPredictServiceMethod):
111+
return p.parseVertexRequest(ctx, body, headers, &aiplatformpb.StreamRawPredictRequest{}, "StreamRawPredictRequest", openAIResponsesPath)
112+
113+
case strings.HasSuffix(path, rawPredictServiceMethod):
114+
return p.parseVertexRequest(ctx, body, headers, &aiplatformpb.RawPredictRequest{}, "RawPredictRequest", openAIResponsesPath)
122115

123116
default:
124117
return &fwkrh.ParseResult{Skip: true}, nil
@@ -134,7 +127,7 @@ func (p *VertexAIParser) ParseResponse(ctx context.Context, body []byte, headers
134127

135128
parsedPayload, err := grpcutil.ParseGrpcPayload(body)
136129
if err != nil {
137-
return nil, fmt.Errorf("parsing gRPC payload: %w", err)
130+
return nil, fmt.Errorf("parsing gRPC response payload: %w", err)
138131
}
139132

140133
respMsg := &httpbody.HttpBody{}
@@ -145,3 +138,39 @@ func (p *VertexAIParser) ParseResponse(ctx context.Context, body []byte, headers
145138

146139
return p.openAIParser.ParseResponse(ctx, jsonBytes, headers, false)
147140
}
141+
142+
type httpBodyMessage interface {
143+
proto.Message
144+
GetHttpBody() *httpbody.HttpBody
145+
}
146+
147+
// parseVertexRequest is a generic helper to parse Vertex AI gRPC requests that wrap an HttpBody payload.
148+
func (p *VertexAIParser) parseVertexRequest(ctx context.Context, body []byte, headers map[string]string, req httpBodyMessage, typeName string, targetPath string) (*fwkrh.ParseResult, error) {
149+
parsedPayload, err := grpcutil.ParseGrpcPayload(body)
150+
if err != nil {
151+
return nil, fmt.Errorf("invalid or unsupported gRPC payload: %w", err)
152+
}
153+
154+
if err := proto.Unmarshal(parsedPayload, req); err != nil {
155+
return nil, fmt.Errorf("unmarshaling %s: %w", typeName, err)
156+
}
157+
158+
httpBody := req.GetHttpBody()
159+
if httpBody == nil {
160+
return nil, fmt.Errorf("%s has no HttpBody", typeName)
161+
}
162+
jsonBytes := httpBody.GetData()
163+
164+
// Use OpenAI parser to parse the JSON payload
165+
// Clone headers and set path to targetPath to make OpenAI parser recognize it
166+
headersCopy := maps.Clone(headers)
167+
headersCopy[parsers.MethodPathKey] = targetPath
168+
parseResult, err := p.openAIParser.ParseRequest(ctx, jsonBytes, headersCopy)
169+
if err != nil {
170+
return nil, fmt.Errorf("parsing %s: %w", typeName, err)
171+
}
172+
173+
inferenceRequestBody := parseResult.Body
174+
inferenceRequestBody.Payload = fwkrh.PayloadProto{Message: req}
175+
return &fwkrh.ParseResult{Body: inferenceRequestBody, Skip: parseResult.Skip}, nil
176+
}

‎pkg/epp/framework/plugins/requesthandling/parsers/vertexai/vertexai_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,32 @@ func TestParseRequest(t *testing.T) {
4848
t.Fatalf("Failed to create gRPC frame: %v", err)
4949
}
5050

51+
streamRawPredictResponsesPayload := []byte(`{"input":"Hello from stream raw predict"}`)
52+
53+
streamRawPredictReqMsg := &aiplatformpb.StreamRawPredictRequest{
54+
HttpBody: &httpbody.HttpBody{
55+
Data: streamRawPredictResponsesPayload,
56+
},
57+
}
58+
59+
validStreamRawPredictBody, err := createGrpcFrame(streamRawPredictReqMsg)
60+
if err != nil {
61+
t.Fatalf("Failed to create gRPC frame: %v", err)
62+
}
63+
64+
rawPredictResponsesPayload := []byte(`{"input":"Hello from raw predict"}`)
65+
66+
rawPredictReqMsg := &aiplatformpb.RawPredictRequest{
67+
HttpBody: &httpbody.HttpBody{
68+
Data: rawPredictResponsesPayload,
69+
},
70+
}
71+
72+
validRawPredictBody, err := createGrpcFrame(rawPredictReqMsg)
73+
if err != nil {
74+
t.Fatalf("Failed to create gRPC frame: %v", err)
75+
}
76+
5177
tests := []struct {
5278
name string
5379
body []byte
@@ -75,6 +101,40 @@ func TestParseRequest(t *testing.T) {
75101
Skip: false,
76102
},
77103
},
104+
{
105+
name: "Success - StreamRawPredict",
106+
body: validStreamRawPredictBody,
107+
headers: map[string]string{
108+
":path": "/google.cloud.aiplatform.v1beta1.PredictionService/StreamRawPredict",
109+
"content-type": "application/grpc",
110+
},
111+
wantResult: &fwkrh.ParseResult{
112+
Body: &fwkrh.InferenceRequestBody{
113+
Responses: &fwkrh.ResponsesRequest{
114+
Input: "Hello from stream raw predict",
115+
},
116+
Payload: fwkrh.PayloadProto{Message: streamRawPredictReqMsg},
117+
},
118+
Skip: false,
119+
},
120+
},
121+
{
122+
name: "Success - RawPredict",
123+
body: validRawPredictBody,
124+
headers: map[string]string{
125+
":path": "/google.cloud.aiplatform.v1beta1.PredictionService/RawPredict",
126+
"content-type": "application/grpc",
127+
},
128+
wantResult: &fwkrh.ParseResult{
129+
Body: &fwkrh.InferenceRequestBody{
130+
Responses: &fwkrh.ResponsesRequest{
131+
Input: "Hello from raw predict",
132+
},
133+
Payload: fwkrh.PayloadProto{Message: rawPredictReqMsg},
134+
},
135+
Skip: false,
136+
},
137+
},
78138
{
79139
name: "Unsupported Path",
80140
body: []byte{},
@@ -95,6 +155,18 @@ func TestParseRequest(t *testing.T) {
95155
headers: map[string]string{":path": "/google.cloud.aiplatform.v1beta1.PredictionService/ChatCompletions", "content-type": "application/grpc"},
96156
wantErr: "unmarshaling ChatCompletionsRequest",
97157
},
158+
{
159+
name: "Invalid proto message - StreamRawPredict",
160+
body: []byte{0, 0, 0, 0, 1, 0xFF}, // Valid header, invalid payload
161+
headers: map[string]string{":path": "/google.cloud.aiplatform.v1beta1.PredictionService/StreamRawPredict", "content-type": "application/grpc"},
162+
wantErr: "unmarshaling StreamRawPredictRequest",
163+
},
164+
{
165+
name: "Invalid proto message - RawPredict",
166+
body: []byte{0, 0, 0, 0, 1, 0xFF}, // Valid header, invalid payload
167+
headers: map[string]string{":path": "/google.cloud.aiplatform.v1beta1.PredictionService/RawPredict", "content-type": "application/grpc"},
168+
wantErr: "unmarshaling RawPredictRequest",
169+
},
98170
}
99171

100172
for _, tc := range tests {

0 commit comments

Comments
 (0)