Skip to content

Commit cfefd00

Browse files
yuluo-yxyuluo-yx
authored
optimize: use openai go sdk ChatCompletion replace map struct (#246)
Signed-off-by: yuluo-yx <[email protected]>
1 parent 164c391 commit cfefd00

File tree

4 files changed

+79
-74
lines changed

4 files changed

+79
-74
lines changed

‎src/semantic-router/pkg/extproc/caching_test.go‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
. "github.com/onsi/ginkgo/v2"
88
. "github.com/onsi/gomega"
9+
"github.com/openai/openai-go"
910

1011
ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1112

@@ -83,18 +84,18 @@ var _ = Describe("Caching Functionality", func() {
8384
}
8485

8586
// Simulate response processing
86-
openAIResponse := map[string]interface{}{
87-
"choices": []map[string]interface{}{
87+
openAIResponse := openai.ChatCompletion{
88+
Choices: []openai.ChatCompletionChoice{
8889
{
89-
"message": map[string]interface{}{
90-
"content": "Cached response",
90+
Message: openai.ChatCompletionMessage{
91+
Content: "Cached response.",
9192
},
9293
},
9394
},
94-
"usage": map[string]interface{}{
95-
"prompt_tokens": 10,
96-
"completion_tokens": 5,
97-
"total_tokens": 15,
95+
Usage: openai.CompletionUsage{
96+
PromptTokens: 10,
97+
CompletionTokens: 5,
98+
TotalTokens: 15,
9899
},
99100
}
100101

@@ -142,18 +143,18 @@ var _ = Describe("Caching Functionality", func() {
142143
Expect(err).To(Or(BeNil(), HaveOccurred()))
143144

144145
// Process response
145-
openAIResponse := map[string]interface{}{
146-
"choices": []map[string]interface{}{
146+
openAIResponse := openai.ChatCompletion{
147+
Choices: []openai.ChatCompletionChoice{
147148
{
148-
"message": map[string]interface{}{
149-
"content": "Machine learning is a subset of artificial intelligence...",
149+
Message: openai.ChatCompletionMessage{
150+
Content: "Machine learning is a subset of artificial intelligence...",
150151
},
151152
},
152153
},
153-
"usage": map[string]interface{}{
154-
"prompt_tokens": 20,
155-
"completion_tokens": 30,
156-
"total_tokens": 50,
154+
Usage: openai.CompletionUsage{
155+
PromptTokens: 20,
156+
CompletionTokens: 30,
157+
TotalTokens: 50,
157158
},
158159
}
159160

‎src/semantic-router/pkg/extproc/metrics_integration_test.go‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
. "github.com/onsi/ginkgo/v2"
88
. "github.com/onsi/gomega"
9+
"github.com/openai/openai-go"
910

1011
core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
1112
ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -86,23 +87,24 @@ var _ = Describe("Metrics recording", func() {
8687
beforePrompt := getHistogramSampleCount("llm_prompt_tokens_per_request", ctx.RequestModel)
8788
beforeCompletion := getHistogramSampleCount("llm_completion_tokens_per_request", ctx.RequestModel)
8889

89-
openAIResponse := map[string]interface{}{
90-
"id": "chatcmpl-xyz",
91-
"object": "chat.completion",
92-
"created": time.Now().Unix(),
93-
"model": ctx.RequestModel,
94-
"usage": map[string]interface{}{
95-
"prompt_tokens": 10,
96-
"completion_tokens": 5,
97-
"total_tokens": 15,
90+
openAIResponse := openai.ChatCompletion{
91+
ID: "chatcmpl-xyz",
92+
Object: "chat.completion",
93+
Created: time.Now().Unix(),
94+
Model: ctx.RequestModel,
95+
Usage: openai.CompletionUsage{
96+
PromptTokens: 10,
97+
CompletionTokens: 5,
98+
TotalTokens: 15,
9899
},
99-
"choices": []map[string]interface{}{
100+
Choices: []openai.ChatCompletionChoice{
100101
{
101-
"message": map[string]interface{}{"role": "assistant", "content": "Hello"},
102-
"finish_reason": "stop",
102+
Message: openai.ChatCompletionMessage{Role: "assistant", Content: "Hello"},
103+
FinishReason: "stop",
103104
},
104105
},
105106
}
107+
106108
respBodyJSON, err := json.Marshal(openAIResponse)
107109
Expect(err).NotTo(HaveOccurred())
108110

‎src/semantic-router/pkg/extproc/request_processing_test.go‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66

77
. "github.com/onsi/ginkgo/v2"
88
. "github.com/onsi/gomega"
9+
"github.com/openai/openai-go"
910

1011
core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
1112
ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -360,23 +361,24 @@ var _ = Describe("Request Processing", func() {
360361

361362
Describe("handleResponseBody", func() {
362363
It("should process response body with token parsing", func() {
363-
openAIResponse := map[string]interface{}{
364-
"id": "chatcmpl-123",
365-
"object": "chat.completion",
366-
"created": time.Now().Unix(),
367-
"model": "model-a",
368-
"usage": map[string]interface{}{
369-
"prompt_tokens": 150,
370-
"completion_tokens": 50,
371-
"total_tokens": 200,
364+
365+
openAIResponse := openai.ChatCompletion{
366+
ID: "chatcmpl-123",
367+
Object: "chat.completion",
368+
Created: time.Now().Unix(),
369+
Model: "model-a",
370+
Usage: openai.CompletionUsage{
371+
PromptTokens: 150,
372+
CompletionTokens: 50,
373+
TotalTokens: 200,
372374
},
373-
"choices": []map[string]interface{}{
375+
Choices: []openai.ChatCompletionChoice{
374376
{
375-
"message": map[string]interface{}{
376-
"role": "assistant",
377-
"content": "This is a test response",
377+
Message: openai.ChatCompletionMessage{
378+
Role: "assistant",
379+
Content: "This is a test response",
378380
},
379-
"finish_reason": "stop",
381+
FinishReason: "stop",
380382
},
381383
},
382384
}

‎src/semantic-router/pkg/utils/http/response.go‎

Lines changed: 32 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
99
ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1010
typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
11+
"github.com/openai/openai-go"
1112
"github.com/vllm-project/semantic-router/src/semantic-router/pkg/metrics"
1213
"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability"
1314
)
@@ -18,26 +19,26 @@ func CreatePIIViolationResponse(model string, deniedPII []string) *ext_proc.Proc
1819
metrics.RecordPIIViolations(model, deniedPII)
1920

2021
// Create OpenAI-compatible response format for PII violations
21-
openAIResponse := map[string]interface{}{
22-
"id": fmt.Sprintf("chatcmpl-pii-violation-%d", time.Now().Unix()),
23-
"object": "chat.completion",
24-
"created": time.Now().Unix(),
25-
"model": model,
26-
"system_fingerprint": "router_pii_policy",
27-
"choices": []map[string]interface{}{
22+
unixTimeStep := time.Now().Unix()
23+
openAIResponse := openai.ChatCompletion{
24+
ID: fmt.Sprintf("chatcmpl-pii-violation-%d", unixTimeStep),
25+
Object: "chat.completion",
26+
Created: unixTimeStep,
27+
Model: model,
28+
Choices: []openai.ChatCompletionChoice{
2829
{
29-
"index": 0,
30-
"message": map[string]interface{}{
31-
"role": "assistant",
32-
"content": fmt.Sprintf("I cannot process this request as it contains personally identifiable information (%v) that is not allowed for the '%s' model according to the configured privacy policy. Please remove any sensitive information and try again.", deniedPII, model),
30+
Index: 0,
31+
Message: openai.ChatCompletionMessage{
32+
Role: "assistant",
33+
Content: fmt.Sprintf("I cannot process this request as it contains personally identifiable information (%v) that is not allowed for the '%s' model according to the configured privacy policy. Please remove any sensitive information and try again.", deniedPII, model),
3334
},
34-
"finish_reason": "content_filter",
35+
FinishReason: "content_filter",
3536
},
3637
},
37-
"usage": map[string]interface{}{
38-
"prompt_tokens": 0,
39-
"completion_tokens": 0,
40-
"total_tokens": 0,
38+
Usage: openai.CompletionUsage{
39+
PromptTokens: 0,
40+
CompletionTokens: 0,
41+
TotalTokens: 0,
4142
},
4243
}
4344

@@ -81,26 +82,25 @@ func CreatePIIViolationResponse(model string, deniedPII []string) *ext_proc.Proc
8182
// CreateJailbreakViolationResponse creates an HTTP response for jailbreak detection violations
8283
func CreateJailbreakViolationResponse(jailbreakType string, confidence float32) *ext_proc.ProcessingResponse {
8384
// Create OpenAI-compatible response format for jailbreak violations
84-
openAIResponse := map[string]interface{}{
85-
"id": fmt.Sprintf("chatcmpl-jailbreak-blocked-%d", time.Now().Unix()),
86-
"object": "chat.completion",
87-
"created": time.Now().Unix(),
88-
"model": "security-filter",
89-
"system_fingerprint": "router_prompt_guard",
90-
"choices": []map[string]interface{}{
85+
openAIResponse := openai.ChatCompletion{
86+
ID: fmt.Sprintf("chatcmpl-jailbreak-blocked-%d", time.Now().Unix()),
87+
Object: "chat.completion",
88+
Created: time.Now().Unix(),
89+
Model: "security-filter",
90+
Choices: []openai.ChatCompletionChoice{
9191
{
92-
"index": 0,
93-
"message": map[string]interface{}{
94-
"role": "assistant",
95-
"content": fmt.Sprintf("I cannot process this request as it appears to contain a potential jailbreak attempt (type: %s, confidence: %.3f). Please rephrase your request in a way that complies with our usage policies.", jailbreakType, confidence),
92+
Index: 0,
93+
Message: openai.ChatCompletionMessage{
94+
Role: "assistant",
95+
Content: fmt.Sprintf("I cannot process this request as it appears to contain a potential jailbreak attempt (type: %s, confidence: %.3f). Please rephrase your request in a way that complies with our usage policies.", jailbreakType, confidence),
9696
},
97-
"finish_reason": "content_filter",
97+
FinishReason: "content_filter",
9898
},
9999
},
100-
"usage": map[string]interface{}{
101-
"prompt_tokens": 0,
102-
"completion_tokens": 0,
103-
"total_tokens": 0,
100+
Usage: openai.CompletionUsage{
101+
PromptTokens: 0,
102+
CompletionTokens: 0,
103+
TotalTokens: 0,
104104
},
105105
}
106106

0 commit comments

Comments
 (0)