first commit
This commit is contained in:
Beyhan Oğur
323
plugins/jsonparser/plugin_test.go
Normal file
323
plugins/jsonparser/plugin_test.go
Normal file
@@ -0,0 +1,323 @@
|
||||
package jsonparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
bifrost "github.com/maximhq/bifrost/core"
|
||||
"github.com/maximhq/bifrost/core/schemas"
|
||||
)
|
||||
|
||||
// BaseAccount implements the schemas.Account interface for testing purposes.
|
||||
// It provides mock implementations of the required methods to test the JSON parser plugin
|
||||
// with a basic OpenAI configuration.
|
||||
type BaseAccount struct{}
|
||||
|
||||
// GetConfiguredProviders returns a list of supported providers for testing.
|
||||
// Currently only supports OpenAI for simplicity in testing.
|
||||
func (baseAccount *BaseAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) {
|
||||
return []schemas.ModelProvider{schemas.OpenAI}, nil
|
||||
}
|
||||
|
||||
// GetKeysForProvider returns a mock API key configuration for testing.
|
||||
// Uses the OPENAI_API_KEY environment variable for authentication.
|
||||
func (baseAccount *BaseAccount) GetKeysForProvider(ctx context.Context, providerKey schemas.ModelProvider) ([]schemas.Key, error) {
|
||||
return []schemas.Key{
|
||||
{
|
||||
Value: *schemas.NewEnvVar("env.OPENAI_API_KEY"),
|
||||
Models: []string{"gpt-4o-mini", "gpt-4-turbo"},
|
||||
Weight: 1.0,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetConfigForProvider returns default provider configuration for testing.
|
||||
// Uses standard network and concurrency settings.
|
||||
func (baseAccount *BaseAccount) GetConfigForProvider(providerKey schemas.ModelProvider) (*schemas.ProviderConfig, error) {
|
||||
return &schemas.ProviderConfig{
|
||||
NetworkConfig: schemas.DefaultNetworkConfig,
|
||||
ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// TestJsonParserPluginEndToEnd tests the integration of the JSON parser plugin with Bifrost.
|
||||
// It performs the following steps:
|
||||
// 1. Initializes the JSON parser plugin with AllRequests usage
|
||||
// 2. Sets up a test Bifrost instance with the plugin
|
||||
// 3. Makes a test chat completion request with streaming enabled
|
||||
// 4. Verifies that the plugin processes the streaming response correctly
|
||||
//
|
||||
// Required environment variables:
|
||||
// - OPENAI_API_KEY: Your OpenAI API key for the test request
|
||||
func TestJsonParserPluginEndToEnd(t *testing.T) {
|
||||
ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
|
||||
// Check if OpenAI API key is set
|
||||
if os.Getenv("OPENAI_API_KEY") == "" {
|
||||
t.Skip("OPENAI_API_KEY is not set, skipping end-to-end test")
|
||||
}
|
||||
|
||||
// Initialize the JSON parser plugin for all requests
|
||||
plugin, err := Init(PluginConfig{
|
||||
Usage: AllRequests,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
MaxAge: 30 * time.Minute,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Error initializing JSON parser plugin: %v", err)
|
||||
}
|
||||
|
||||
account := BaseAccount{}
|
||||
|
||||
// Initialize Bifrost with the plugin
|
||||
client, err := bifrost.Init(ctx, schemas.BifrostConfig{
|
||||
Account: &account,
|
||||
LLMPlugins: []schemas.LLMPlugin{plugin},
|
||||
Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Error initializing Bifrost: %v", err)
|
||||
}
|
||||
defer client.Shutdown()
|
||||
|
||||
// Make a test responses request with streaming enabled
|
||||
// Request JSON output to test the parser
|
||||
var responseFormat interface{} = map[string]interface{}{
|
||||
"type": "json_object",
|
||||
}
|
||||
|
||||
request := &schemas.BifrostChatRequest{
|
||||
Provider: schemas.OpenAI,
|
||||
Model: "gpt-4o-mini",
|
||||
Input: []schemas.ChatMessage{
|
||||
{
|
||||
Role: schemas.ChatMessageRoleUser,
|
||||
Content: &schemas.ChatMessageContent{
|
||||
ContentStr: bifrost.Ptr("Return a JSON object with name, age, and city fields. Example: {\"name\": \"John\", \"age\": 30, \"city\": \"New York\"}"),
|
||||
},
|
||||
},
|
||||
},
|
||||
Params: &schemas.ChatParameters{
|
||||
ResponseFormat: &responseFormat,
|
||||
},
|
||||
}
|
||||
// Make the streaming request
|
||||
responseChan, bifrostErr := client.ChatCompletionStreamRequest(ctx, request)
|
||||
|
||||
if bifrostErr != nil {
|
||||
t.Fatalf("Error in Bifrost request: %v", bifrostErr)
|
||||
}
|
||||
|
||||
// Process streaming responses
|
||||
if responseChan != nil {
|
||||
t.Logf("Streaming response channel received")
|
||||
|
||||
// Read from the channel to see the streaming responses
|
||||
responseCount := 0
|
||||
|
||||
for streamResponse := range responseChan {
|
||||
responseCount++
|
||||
|
||||
if streamResponse.BifrostError != nil {
|
||||
t.Logf("Streaming response error: %v", streamResponse.BifrostError)
|
||||
}
|
||||
|
||||
if streamResponse.BifrostChatResponse != nil {
|
||||
if streamResponse.BifrostChatResponse.Choices != nil {
|
||||
for _, outputMsg := range streamResponse.BifrostChatResponse.Choices {
|
||||
if outputMsg.ChatStreamResponseChoice != nil && outputMsg.ChatStreamResponseChoice.Delta.Content != nil {
|
||||
content := *outputMsg.ChatStreamResponseChoice.Delta.Content
|
||||
if content != "" {
|
||||
t.Logf("Chunk %d: %s", responseCount, content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("Stream completed after %d responses", responseCount)
|
||||
} else {
|
||||
t.Logf("No streaming response channel received")
|
||||
}
|
||||
|
||||
t.Log("End-to-end test completed - check logs for JSON parsing behavior")
|
||||
}
|
||||
|
||||
// TestJsonParserPluginPerRequest tests the per-request configuration of the JSON parser plugin.
|
||||
// It tests how the plugin behaves when enabled via context for specific requests.
|
||||
//
|
||||
// Required environment variables:
|
||||
// - OPENAI_API_KEY: Your OpenAI API key for the test request
|
||||
func TestJsonParserPluginPerRequest(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
// Check if OpenAI API key is set
|
||||
if os.Getenv("OPENAI_API_KEY") == "" {
|
||||
t.Skip("OPENAI_API_KEY is not set, skipping per-request test")
|
||||
}
|
||||
|
||||
// Initialize the JSON parser plugin for per-request usage
|
||||
plugin, err := Init(PluginConfig{
|
||||
Usage: PerRequest,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
MaxAge: 30 * time.Minute,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Error initializing JSON parser plugin: %v", err)
|
||||
}
|
||||
|
||||
account := BaseAccount{}
|
||||
|
||||
// Initialize Bifrost with the plugin
|
||||
client, err := bifrost.Init(ctx, schemas.BifrostConfig{
|
||||
Account: &account,
|
||||
LLMPlugins: []schemas.LLMPlugin{plugin},
|
||||
Logger: bifrost.NewDefaultLogger(schemas.LogLevelDebug),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Error initializing Bifrost: %v", err)
|
||||
}
|
||||
defer client.Shutdown()
|
||||
|
||||
// Test request with plugin enabled via context
|
||||
var responseFormat interface{} = map[string]interface{}{
|
||||
"type": "json_object",
|
||||
}
|
||||
|
||||
request := &schemas.BifrostChatRequest{
|
||||
Provider: schemas.OpenAI,
|
||||
Model: "gpt-4o-mini",
|
||||
Input: []schemas.ChatMessage{
|
||||
{
|
||||
Role: schemas.ChatMessageRoleUser,
|
||||
Content: &schemas.ChatMessageContent{
|
||||
ContentStr: bifrost.Ptr("Return a JSON object with name and age fields."),
|
||||
},
|
||||
},
|
||||
},
|
||||
Params: &schemas.ChatParameters{
|
||||
ResponseFormat: &responseFormat,
|
||||
},
|
||||
}
|
||||
|
||||
// Create context with plugin enabled
|
||||
newContext := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline).WithValue(EnableStreamingJSONParser, true)
|
||||
|
||||
// Make the streaming request
|
||||
responseChan, bifrostErr := client.ChatCompletionStreamRequest(newContext, request)
|
||||
|
||||
if bifrostErr != nil {
|
||||
t.Logf("Error in Bifrost request: %v", bifrostErr)
|
||||
}
|
||||
|
||||
// Process streaming responses
|
||||
if responseChan != nil {
|
||||
t.Logf("Streaming response channel received for per-request test")
|
||||
|
||||
// Read from the channel to see the streaming responses
|
||||
responseCount := 0
|
||||
|
||||
for streamResponse := range responseChan {
|
||||
responseCount++
|
||||
|
||||
if streamResponse.BifrostError != nil {
|
||||
t.Logf("Streaming response error: %v", streamResponse.BifrostError)
|
||||
}
|
||||
|
||||
if streamResponse.BifrostChatResponse != nil {
|
||||
for _, choice := range streamResponse.BifrostChatResponse.Choices {
|
||||
if choice.ChatStreamResponseChoice != nil && choice.ChatStreamResponseChoice.Delta.Content != nil {
|
||||
content := *choice.ChatStreamResponseChoice.Delta.Content
|
||||
if content != "" {
|
||||
t.Logf("Per-request chunk %d: %s", responseCount, content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("Per-request stream completed after %d responses", responseCount)
|
||||
} else {
|
||||
t.Logf("No streaming response channel received for per-request test")
|
||||
}
|
||||
|
||||
t.Log("Per-request test completed - check logs for JSON parsing behavior")
|
||||
}
|
||||
|
||||
func TestParsePartialJSON(t *testing.T) {
|
||||
plugin, err := Init(PluginConfig{
|
||||
Usage: AllRequests,
|
||||
CleanupInterval: 5 * time.Minute,
|
||||
MaxAge: 30 * time.Minute,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Error initializing JSON parser plugin: %v", err)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "Already valid JSON object",
|
||||
input: `{"name": "John", "age": 30}`,
|
||||
expected: `{"name": "John", "age": 30}`,
|
||||
},
|
||||
{
|
||||
name: "Partial JSON object missing closing brace",
|
||||
input: `{"name": "John", "age": 30, "city": "New York"`,
|
||||
expected: `{"name": "John", "age": 30, "city": "New York"}`,
|
||||
},
|
||||
{
|
||||
name: "Partial JSON array missing closing bracket",
|
||||
input: `["apple", "banana", "cherry"`,
|
||||
expected: `["apple", "banana", "cherry"]`,
|
||||
},
|
||||
{
|
||||
name: "Nested partial JSON",
|
||||
input: `{"user": {"name": "John", "details": {"age": 30, "city": "NY"`,
|
||||
expected: `{"user": {"name": "John", "details": {"age": 30, "city": "NY"}}}`,
|
||||
},
|
||||
{
|
||||
name: "Partial JSON with string containing newline",
|
||||
input: `{"message": "Hello\nWorld"`,
|
||||
expected: `{"message": "Hello\nWorld"}`,
|
||||
},
|
||||
{
|
||||
name: "Empty string",
|
||||
input: "",
|
||||
expected: "{}",
|
||||
},
|
||||
{
|
||||
name: "Whitespace only",
|
||||
input: " \n\t ",
|
||||
expected: "{}",
|
||||
},
|
||||
{
|
||||
name: "Non-JSON string",
|
||||
input: "This is not JSON",
|
||||
expected: "This is not JSON",
|
||||
},
|
||||
{
|
||||
name: "Partial JSON with escaped quotes",
|
||||
input: `{"message": "He said \"Hello\""`,
|
||||
expected: `{"message": "He said \"Hello\""}`,
|
||||
},
|
||||
{
|
||||
name: "Complex nested structure",
|
||||
input: `{"data": {"users": [{"id": 1, "name": "John"}, {"id": 2, "name": "Jane"`,
|
||||
expected: `{"data": {"users": [{"id": 1, "name": "John"}, {"id": 2, "name": "Jane"}]}}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := plugin.parsePartialJSON(tt.input)
|
||||
if result != tt.expected {
|
||||
t.Errorf("parsePartialJSON(%q) = %q, want %q", tt.input, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user