yao/agent/llm/providers/openai/openai_test.go
Max 3dd63e5530 Refactor model capabilities to use OpenAI struct
- Updated the model capabilities throughout the agent to utilize the new gouOpenAI.Capabilities struct instead of the previous ModelCapabilities.
- Adjusted related methods and types to ensure compatibility with the new capabilities structure, enhancing clarity and maintainability.
- Improved context handling and message processing by directly integrating OpenAI capabilities, streamlining the overall architecture.
2025-12-02 11:33:09 +08:00

1538 lines
41 KiB
Go

package openai_test
import (
gocontext "context"
"encoding/json"
"strings"
"testing"
"time"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/agent/output/message"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/test"
)
// TestOpenAIStreamBasic tests basic streaming completion with short output
func TestOpenAIStreamBasic(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector from real configuration
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Create LLM instance with capabilities
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Prepare messages with concise prompt
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'Hello' in one word.",
},
}
// Set short max tokens to ensure quick response
maxTokens := 5
options.MaxTokens = &maxTokens
// Create context
ctx := newTestContext("test-stream-basic", "openai.gpt-4o")
// Track streaming chunks
var chunks []string
handler := func(chunkType message.StreamChunkType, data []byte) int {
chunks = append(chunks, string(data))
t.Logf("Stream chunk [%s]: %s", chunkType, string(data))
return 0 // Continue
}
// Call Stream
response, err := llmInstance.Stream(ctx, messages, options, handler)
if err != nil {
t.Fatalf("Stream failed: %v", err)
}
// Validate response
if response == nil {
t.Fatal("Response is nil")
}
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
if response.Content == "" {
t.Error("Response content is empty")
}
if response.FinishReason == "" {
t.Error("FinishReason is empty")
}
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
if response.Usage.TotalTokens == 0 {
t.Error("Response Usage.TotalTokens is 0")
}
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
if len(chunks) == 0 {
t.Error("No streaming chunks received")
}
t.Logf("Final response: %+v", response)
t.Logf("Total chunks received: %d", len(chunks))
}
// TestOpenAIPostBasic tests basic non-streaming completion
func TestOpenAIPostBasic(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Create LLM instance
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Prepare messages with concise prompt
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Reply with only the word 'OK'.",
},
}
// Set short max tokens
maxTokens := 5
options.MaxTokens = &maxTokens
// Create context
ctx := newTestContext("test-stream-basic", "openai.gpt-4o")
// Call Post
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post failed: %v", err)
}
// Validate response
if response == nil {
t.Fatal("Response is nil")
}
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
if response.Content == "" {
t.Error("Response content is empty")
}
if response.FinishReason == "" {
t.Error("FinishReason is empty")
}
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
if response.Usage.TotalTokens == 0 {
t.Error("Response Usage.TotalTokens is 0")
}
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
t.Logf("Response: %+v", response)
}
// TestOpenAIStreamWithToolCalls tests streaming with tool calls and JSON schema validation
func TestOpenAIStreamWithToolCalls(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Create LLM instance with tool call capabilities
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
// Define a simple weather tool with JSON schema
weatherTool := map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "get_weather",
"description": "Get the current weather for a location",
"parameters": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"location": map[string]interface{}{
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": map[string]interface{}{
"type": "string",
"enum": []string{"celsius", "fahrenheit"},
},
},
"required": []string{"location"},
},
},
}
options.Tools = []map[string]interface{}{weatherTool}
options.ToolChoice = "auto"
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Prepare messages that should trigger tool call
messages := []context.Message{
{
Role: context.RoleUser,
Content: "What's the weather in Tokyo? Use celsius.",
},
}
// Create context
ctx := newTestContext("test-stream-basic", "openai.gpt-4o")
// Track streaming chunks
var toolCallChunks int
handler := func(chunkType message.StreamChunkType, data []byte) int {
if chunkType == message.ChunkToolCall {
toolCallChunks++
}
t.Logf("Stream chunk [%s]: %s", chunkType, string(data))
return 0 // Continue
}
// Call Stream
response, err := llmInstance.Stream(ctx, messages, options, handler)
if err != nil {
t.Fatalf("Stream with tool calls failed: %v", err)
}
// Validate response
if response == nil {
t.Fatal("Response is nil")
}
// Should have tool calls
if len(response.ToolCalls) == 0 {
t.Error("Expected tool calls but got none")
} else {
t.Logf("Received %d tool call(s)", len(response.ToolCalls))
for i, tc := range response.ToolCalls {
t.Logf("Tool call %d: %s(%s)", i, tc.Function.Name, tc.Function.Arguments)
// Validate tool call has required fields
if tc.ID == "" {
t.Errorf("Tool call %d missing ID", i)
}
if tc.Function.Name == "" {
t.Errorf("Tool call %d missing function name", i)
}
if tc.Function.Arguments == "" {
t.Errorf("Tool call %d missing arguments", i)
}
}
}
if response.FinishReason != context.FinishReasonToolCalls {
t.Logf("Warning: Expected finish_reason='tool_calls', got '%s'", response.FinishReason)
}
if toolCallChunks == 0 {
t.Error("No tool call chunks received during streaming")
}
t.Logf("Final response: %+v", response)
}
// TestOpenAIPostWithToolCalls tests non-streaming with tool calls
func TestOpenAIPostWithToolCalls(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Create LLM instance with tool call capabilities
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
}
// Define a calculation tool
calcTool := map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "calculate",
"description": "Perform a mathematical calculation",
"parameters": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"expression": map[string]interface{}{
"type": "string",
"description": "The mathematical expression to evaluate",
},
},
"required": []string{"expression"},
},
},
}
options.Tools = []map[string]interface{}{calcTool}
options.ToolChoice = "auto"
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Prepare messages
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Calculate 15 * 8",
},
}
// Create context
ctx := newTestContext("test-stream-basic", "openai.gpt-4o")
// Call Post
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post with tool calls failed: %v", err)
}
// Validate response
if response == nil {
t.Fatal("Response is nil")
}
// Validate response metadata
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
if response.FinishReason != "tool_calls" {
t.Errorf("FinishReason is %s, expected tool_calls", response.FinishReason)
}
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
if response.Usage.TotalTokens == 0 {
t.Error("Response Usage.TotalTokens is 0")
}
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
// Should have tool calls
if len(response.ToolCalls) == 0 {
t.Error("Expected tool calls but got none")
} else {
tc := response.ToolCalls[0]
// Validate tool call structure
if tc.ID == "" {
t.Error("Tool call ID is empty")
}
if tc.Type != context.ToolTypeFunction {
t.Errorf("Tool call Type is %s, expected %s", tc.Type, context.ToolTypeFunction)
}
if tc.Function.Name != "calculate" {
t.Errorf("Tool call function name is %s, expected calculate", tc.Function.Name)
}
if tc.Function.Arguments == "" {
t.Error("Tool call arguments are empty")
}
t.Logf("Received %d tool call(s)", len(response.ToolCalls))
for i, tc := range response.ToolCalls {
t.Logf("Tool call %d: %s(%s)", i, tc.Function.Name, tc.Function.Arguments)
}
}
t.Logf("Response: %+v", response)
}
// TestOpenAIStreamWithInvalidToolCall tests that invalid tool calls trigger validation error
func TestOpenAIStreamWithInvalidToolCall(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Create LLM instance
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
// Define a strict tool that requires specific format
strictTool := map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "send_email",
"description": "Send an email",
"parameters": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"to": map[string]interface{}{
"type": "string",
"pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$",
},
"subject": map[string]interface{}{
"type": "string",
"minLength": 1,
},
"body": map[string]interface{}{
"type": "string",
"minLength": 1,
},
},
"required": []string{"to", "subject", "body"},
},
},
}
options.Tools = []map[string]interface{}{strictTool}
options.ToolChoice = map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "send_email",
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Prepare messages with incomplete information (should cause validation error)
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Send email to invalid-email without subject",
},
}
// Create context
ctx := newTestContext("test-stream-basic", "openai.gpt-4o")
handler := func(chunkType message.StreamChunkType, data []byte) int {
return 0 // Continue
}
// Call Stream - should succeed but may trigger validation if tool call is malformed
response, err := llmInstance.Stream(ctx, messages, options, handler)
// The API might return a valid tool call despite the bad prompt,
// so we just log the result
if err != nil {
t.Logf("Stream failed as expected with validation error: %v", err)
} else {
t.Logf("Stream succeeded, response: %+v", response)
if len(response.ToolCalls) > 0 {
t.Logf("Tool calls: %v", response.ToolCalls)
}
}
}
// TestOpenAIStreamRetry tests the retry mechanism with invalid API key
func TestOpenAIStreamRetry(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector with invalid API key to trigger 401 error (non-retryable)
connDSL := `{
"type": "openai",
"options": {
"model": "gpt-4o",
"key": "sk-invalid-key-should-fail-auth",
"host": "https://api.openai.com"
}
}`
conn, err := connector.New("openai", "test-retry", []byte(connDSL))
if err != nil {
t.Fatalf("Failed to create test connector: %v", err)
}
// Create LLM instance
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true, // Need this to select OpenAI provider
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Test",
},
}
ctx := newTestContext("test-retry", "test-retry")
// This should fail quickly without retry (401 is non-retryable)
_, err = llmInstance.Stream(ctx, messages, options, nil)
if err == nil {
t.Fatal("Expected error due to invalid API key, but got success")
}
// Verify it's an error related to invalid API key
// Could be: 401, unauthorized, authentication error, or no data (empty response)
errMsg := err.Error()
hasExpectedError := strings.Contains(strings.ToLower(errMsg), "401") ||
strings.Contains(strings.ToLower(errMsg), "unauthorized") ||
strings.Contains(strings.ToLower(errMsg), "authentication") ||
strings.Contains(strings.ToLower(errMsg), "incorrect api key") ||
strings.Contains(strings.ToLower(errMsg), "no data received")
if !hasExpectedError {
t.Errorf("Expected authentication or empty response error, got: %v", err)
}
// Should mention non-retryable (these errors should not trigger retry)
if !strings.Contains(strings.ToLower(errMsg), "non-retryable") {
t.Errorf("Error should indicate non-retryable: %v", err)
}
t.Logf("Failed as expected with error: %v", err)
}
// TestOpenAIStreamChunkTypes tests that stream handler receives correct chunk types
func TestOpenAIStreamChunkTypes(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'test' in one word.",
},
}
ctx := newTestContext("test-chunk-types", "openai.gpt-4o")
// Track chunk types
chunkTypes := make(map[message.StreamChunkType]int)
handler := func(chunkType message.StreamChunkType, data []byte) int {
chunkTypes[chunkType]++
t.Logf("Received chunk type: %s, data length: %d", chunkType, len(data))
return 1 // Continue
}
response, err := llmInstance.Stream(ctx, messages, options, handler)
if err != nil {
t.Fatalf("Stream failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate chunk types received
if chunkTypes[message.ChunkText] == 0 {
t.Error("Expected to receive ChunkText, but got 0")
}
t.Logf("Chunk types received: %+v", chunkTypes)
}
// TestOpenAIStreamErrorCallback tests that errors are sent to stream handler
func TestOpenAIStreamErrorCallback(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Create connector with invalid API key to trigger error
connDSL := `{
"type": "openai",
"options": {
"model": "gpt-4o",
"key": "sk-invalid-for-error-test",
"host": "https://api.openai.com"
}
}`
conn, err := connector.New("openai", "test-error-callback", []byte(connDSL))
if err != nil {
t.Fatalf("Failed to create test connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Test",
},
}
ctx := newTestContext("test-error-callback", "test-error-callback")
// Track if error chunk was received
receivedError := false
var errorMessage string
handler := func(chunkType message.StreamChunkType, data []byte) int {
if chunkType == message.ChunkError {
receivedError = true
errorMessage = string(data)
t.Logf("Received error chunk: %s", errorMessage)
}
return 1 // Continue
}
// This should fail and send error to handler
_, err = llmInstance.Stream(ctx, messages, options, handler)
if err == nil {
t.Fatal("Expected error due to invalid API key")
}
// Verify error was sent to handler
if !receivedError {
t.Error("Expected to receive ChunkError in handler, but didn't")
}
if errorMessage == "" {
t.Error("Error message in chunk is empty")
}
t.Logf("Error callback test passed. Error: %v", err)
}
// TestOpenAIToolCallValidationRetry tests automatic tool call validation retry with LLM feedback
func TestOpenAIToolCallValidationRetry(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
Tools: []map[string]interface{}{
{
"type": "function",
"function": map[string]interface{}{
"name": "test_strict_validation",
"description": "A function with very strict validation rules",
"parameters": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"status": map[string]interface{}{
"type": "string",
"description": "Must be exactly 'active' or 'inactive'",
"enum": []string{"active", "inactive"},
},
"priority": map[string]interface{}{
"type": "integer",
"description": "Must be between 1 and 5",
"minimum": 1,
"maximum": 5,
},
},
"required": []string{"status", "priority"},
},
},
},
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
ctx := newTestContext("test-tool-validation-retry", "openai.gpt-4o")
// Try to make LLM call with intentionally unclear requirements
// This may or may not trigger validation, depending on LLM behavior
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Call test_strict_validation function with status='pending' and priority=10",
},
}
// The Provider will automatically:
// 1. Call LLM
// 2. If validation fails, add error feedback to conversation
// 3. Retry up to 3 times with feedback
// 4. Return success or validation error after max retries
response, err := llmInstance.Stream(ctx, messages, options, nil)
if err != nil {
// Check if it's a validation error after retries
if strings.Contains(err.Error(), "tool call validation failed after") &&
strings.Contains(err.Error(), "retries") {
t.Logf("✓ Automatic validation retry exhausted: %v", err)
} else if strings.Contains(err.Error(), "validation") {
t.Logf("✓ Validation failed: %v", err)
} else {
t.Logf("Request failed (non-validation): %v", err)
}
} else if response != nil {
if len(response.ToolCalls) > 0 {
t.Logf("✓ Tool call succeeded (possibly after auto-retry): %+v", response.ToolCalls[0])
// Verify the tool call arguments are valid
tc := response.ToolCalls[0]
var args map[string]interface{}
if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err == nil {
if status, ok := args["status"].(string); ok {
if status != "active" && status != "inactive" {
t.Errorf("Status should be 'active' or 'inactive', got: %s", status)
}
}
if priority, ok := args["priority"].(float64); ok {
if priority < 1 || priority > 5 {
t.Errorf("Priority should be between 1-5, got: %v", priority)
}
}
}
} else {
t.Log("✓ Response returned but no tool calls")
}
}
t.Log("Automatic tool call validation retry test completed")
}
// TestOpenAIJSONMode tests JSON mode response formatting
func TestOpenAIJSONMode(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSON,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Generate a JSON object with fields: name (string), age (number), city (string). Use values: John, 30, New York",
},
}
ctx := newTestContext("test-json-mode", "openai.gpt-4o")
// Test streaming with JSON mode
response, err := llmInstance.Stream(ctx, messages, options, nil)
if err != nil {
t.Fatalf("Stream with JSON mode failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate response
contentStr, ok := response.Content.(string)
if !ok || contentStr == "" {
t.Error("Response content is empty or not a string")
}
// Try to parse as JSON
var jsonData map[string]interface{}
if err := json.Unmarshal([]byte(contentStr), &jsonData); err != nil {
t.Errorf("Response is not valid JSON: %v\nContent: %s", err, contentStr)
} else {
t.Logf("✓ Response is valid JSON: %+v", jsonData)
// Verify expected fields exist
if _, hasName := jsonData["name"]; !hasName {
t.Error("JSON response missing 'name' field")
}
if _, hasAge := jsonData["age"]; !hasAge {
t.Error("JSON response missing 'age' field")
}
if _, hasCity := jsonData["city"]; !hasCity {
t.Error("JSON response missing 'city' field")
}
}
// Validate metadata
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
t.Log("JSON mode test completed successfully")
}
// TestOpenAIJSONModePost tests JSON mode with non-streaming
func TestOpenAIJSONModePost(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSON,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Return a JSON with: status='success', count=42",
},
}
ctx := newTestContext("test-json-mode-post", "openai.gpt-4o")
// Test non-streaming with JSON mode
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post with JSON mode failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate response content is JSON
contentStr, ok := response.Content.(string)
if !ok || contentStr == "" {
t.Error("Response content is empty or not a string")
}
var jsonData map[string]interface{}
if err := json.Unmarshal([]byte(contentStr), &jsonData); err != nil {
t.Errorf("Response is not valid JSON: %v\nContent: %s", err, contentStr)
} else {
t.Logf("✓ Response is valid JSON: %+v", jsonData)
}
// Validate metadata
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
if response.Usage.TotalTokens == 0 {
t.Error("Response Usage.TotalTokens is 0")
}
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
t.Log("JSON mode Post test completed successfully")
}
// TestOpenAIJSONSchema tests JSON mode with strict schema validation
func TestOpenAIJSONSchema(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Define a strict JSON schema
// Note: For OpenAI strict mode, 'required' must include ALL properties
schema := map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"user": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"name": map[string]interface{}{
"type": "string",
"description": "User's full name",
},
"email": map[string]interface{}{
"type": "string",
"description": "User's email address",
},
"age": map[string]interface{}{
"type": "integer",
"description": "User's age",
},
"isActive": map[string]interface{}{
"type": "boolean",
"description": "Whether user is active",
},
},
"required": []string{"name", "email", "age", "isActive"},
"additionalProperties": false,
},
},
"required": []string{"user"},
"additionalProperties": false,
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSONSchema,
JSONSchema: &context.JSONSchema{
Name: "user_info",
Description: "User information schema",
Schema: schema,
Strict: func() *bool { v := true; return &v }(),
},
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Create user info for: Alice Smith, alice@example.com, age 28, active user",
},
}
ctx := newTestContext("test-json-schema", "openai.gpt-4o")
// Test streaming with JSON schema
response, err := llmInstance.Stream(ctx, messages, options, nil)
if err != nil {
t.Fatalf("Stream with JSON schema failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate response content
contentStr, ok := response.Content.(string)
if !ok || contentStr == "" {
t.Fatal("Response content is empty or not a string")
}
// Parse as JSON
var jsonData map[string]interface{}
if err := json.Unmarshal([]byte(contentStr), &jsonData); err != nil {
t.Fatalf("Response is not valid JSON: %v\nContent: %s", err, contentStr)
}
t.Logf("✓ Response is valid JSON: %+v", jsonData)
// Verify structure matches schema
user, hasUser := jsonData["user"].(map[string]interface{})
if !hasUser {
t.Fatal("JSON response missing 'user' object")
}
// Verify required fields
if _, hasName := user["name"]; !hasName {
t.Error("User object missing required 'name' field")
}
if _, hasEmail := user["email"]; !hasEmail {
t.Error("User object missing required 'email' field")
}
// Verify field types
if name, ok := user["name"].(string); ok {
t.Logf("✓ name: %s (string)", name)
} else {
t.Error("name is not a string")
}
if email, ok := user["email"].(string); ok {
t.Logf("✓ email: %s (string)", email)
} else {
t.Error("email is not a string")
}
if age, ok := user["age"].(float64); ok {
if age < 0 || age > 150 {
t.Errorf("age %v is out of range [0, 150]", age)
}
t.Logf("✓ age: %v (integer, in range)", age)
}
if isActive, ok := user["isActive"].(bool); ok {
t.Logf("✓ isActive: %v (boolean)", isActive)
}
// Validate metadata
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
t.Log("JSON schema test completed successfully")
}
// TestOpenAIJSONSchemaPost tests JSON schema with non-streaming
func TestOpenAIJSONSchemaPost(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Simple schema for testing
// Note: For OpenAI strict mode, 'required' must include ALL properties
schema := map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"status": map[string]interface{}{
"type": "string",
"enum": []string{"success", "error", "pending"},
},
"message": map[string]interface{}{
"type": "string",
},
"code": map[string]interface{}{
"type": "integer",
},
},
"required": []string{"status", "message", "code"},
"additionalProperties": false,
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSONSchema,
JSONSchema: &context.JSONSchema{
Name: "api_response",
Description: "API response format",
Schema: schema,
Strict: func() *bool { v := true; return &v }(),
},
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Generate an API response with status 'success', message 'Operation completed', and code 200",
},
}
ctx := newTestContext("test-json-schema-post", "openai.gpt-4o")
// Test non-streaming with JSON schema
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post with JSON schema failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate response content
contentStr, ok := response.Content.(string)
if !ok || contentStr == "" {
t.Fatal("Response content is empty or not a string")
}
// Parse and validate JSON
var jsonData map[string]interface{}
if err := json.Unmarshal([]byte(contentStr), &jsonData); err != nil {
t.Fatalf("Response is not valid JSON: %v\nContent: %s", err, contentStr)
}
t.Logf("✓ Response is valid JSON: %+v", jsonData)
// Verify required fields
status, hasStatus := jsonData["status"].(string)
if !hasStatus {
t.Fatal("Missing required 'status' field")
}
// Verify enum constraint
validStatuses := map[string]bool{"success": true, "error": true, "pending": true}
if !validStatuses[status] {
t.Errorf("status '%s' is not in enum [success, error, pending]", status)
}
if _, hasMessage := jsonData["message"].(string); !hasMessage {
t.Error("Missing required 'message' field")
}
// Validate metadata
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
t.Log("JSON schema Post test completed successfully")
}
// TestOpenAIProxySupport tests that HTTP proxy configuration is respected
func TestOpenAIProxySupport(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// This test verifies proxy support exists in the connector configuration
// Actual proxy testing requires a real proxy server setup
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
settings := conn.Setting()
t.Logf("Connector settings: %+v", settings)
// Verify host field exists in settings (host is the API endpoint)
if host, hasHost := settings["host"]; hasHost {
t.Logf("API host configured: %v", host)
} else {
t.Log("Host field not in settings (will use default)")
}
// The actual HTTP proxy functionality is implemented via environment variables
// (HTTP_PROXY, HTTPS_PROXY, NO_PROXY) and handled by http.GetTransport
t.Log("HTTP proxy support is implemented via http.GetTransport using environment variables")
}
// TestOpenAIStreamLifecycleEvents tests that LLM-level lifecycle events are sent correctly
// LLM layer sends group_start/end for individual messages (thinking, text, tool_call)
// Note: stream_start/end and Agent-level blocks are handled at Agent level
func TestOpenAIStreamLifecycleEvents(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'hello' in one word",
},
}
ctx := newTestContext("test-lifecycle", "openai.gpt-4o")
// Track lifecycle events (group_start/end at LLM layer represent message boundaries)
var events []string
var groupStartReceived, groupEndReceived bool
handler := func(chunkType message.StreamChunkType, data []byte) int {
events = append(events, string(chunkType))
switch chunkType {
case message.ChunkStreamStart:
t.Error("❌ LLM layer should NOT send stream_start (now sent at Agent level)")
case message.ChunkStreamEnd:
t.Error("❌ LLM layer should NOT send stream_end (now sent at Agent level)")
case message.ChunkMessageStart:
groupStartReceived = true
var startData message.EventMessageStartData
if err := json.Unmarshal(data, &startData); err == nil {
t.Logf("✓ group_start (message start): type=%s, id=%s", startData.Type, startData.MessageID)
if startData.MessageID == "" {
t.Error("group_start missing message_id")
}
} else {
t.Errorf("Failed to parse group_start data: %v", err)
}
case message.ChunkMessageEnd:
groupEndReceived = true
var endData message.EventMessageEndData
if err := json.Unmarshal(data, &endData); err == nil {
t.Logf("✓ group_end (message end): type=%s, chunks=%d, duration=%dms",
endData.Type, endData.ChunkCount, endData.DurationMs)
if endData.ChunkCount <= 0 {
t.Error("group_end should have chunk_count > 0")
}
} else {
t.Errorf("Failed to parse group_end data: %v", err)
}
case message.ChunkText:
t.Logf(" text chunk: %s", string(data))
}
return 0
}
response, err := llmInstance.Stream(ctx, messages, options, handler)
if err != nil {
t.Fatalf("Stream failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate that LLM-level message lifecycle events were received
if !groupStartReceived {
t.Error("group_start (message start) event was not received")
}
if !groupEndReceived {
t.Error("group_end (message end) event was not received")
}
// Validate event order: group_start should come before group_end
if len(events) < 2 {
t.Errorf("Expected at least 2 events (message start/end), got %d", len(events))
}
t.Logf("Total events received: %d", len(events))
t.Log("LLM message lifecycle events test completed successfully")
t.Log("Note: LLM layer group_start/end represent message boundaries (thinking, text, tool_call)")
t.Log(" Agent-level block boundaries and stream_start/end are handled at Agent level")
}
// TestOpenAIStreamContextCancellation tests that stream respects context cancellation
func TestOpenAIStreamContextCancellation(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Write a very long essay about the history of computing", // Long task
},
}
// Create a context with a very short timeout
ctx := newTestContext("test-cancel", "openai.gpt-4o")
goCtx, cancel := gocontext.WithTimeout(gocontext.Background(), 100*time.Millisecond)
defer cancel()
ctx.Context = goCtx
var receivedChunks int
handler := func(chunkType message.StreamChunkType, data []byte) int {
if chunkType == message.ChunkText || chunkType == message.ChunkToolCall {
receivedChunks++
}
// Note: stream_end is now sent at Agent level, not LLM level
if chunkType == message.ChunkStreamEnd {
t.Error("❌ LLM layer should NOT send stream_end (now sent at Agent level)")
}
return 0
}
response, err := llmInstance.Stream(ctx, messages, options, handler)
// Should get an error due to context cancellation
if err == nil {
t.Error("Expected error due to context cancellation, but got nil")
} else {
t.Logf("✓ Got expected cancellation error: %v", err)
// Check if error message indicates cancellation
errStr := err.Error()
if !strings.Contains(errStr, "context") && !strings.Contains(errStr, "cancel") {
t.Errorf("Error should mention context/cancellation: %v", err)
}
}
// Response should be nil due to cancellation
if response != nil {
t.Logf("Warning: Response is not nil despite cancellation (partial response)")
}
t.Logf("Received %d chunks before cancellation", receivedChunks)
t.Log("Context cancellation test completed successfully")
t.Log("Note: stream_end for cancellation is now sent at Agent level")
}
// TestOpenAIStreamWithTemperature tests different temperature settings
func TestOpenAIStreamWithTemperature(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
temperature := 0.7 // Moderate temperature
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true, // Need this to select OpenAI provider
},
Temperature: &temperature,
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'yes' in one word.",
},
}
ctx := newTestContext("test-temperature", "openai.gpt-4o")
// Use callback to collect chunks
chunkCount := 0
var callback message.StreamFunc = func(chunkType message.StreamChunkType, data []byte) int {
chunkCount++
return 1 // Continue
}
response, err := llmInstance.Stream(ctx, messages, options, callback)
if err != nil {
t.Fatalf("Stream failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Validate response data
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
if response.Content == "" {
t.Error("Response Content is empty")
}
if response.FinishReason == "" {
t.Error("Response FinishReason is empty")
}
if response.Usage == nil {
t.Error("Response Usage is nil")
} else {
if response.Usage.TotalTokens == 0 {
t.Error("Response Usage.TotalTokens is 0")
}
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
if chunkCount == 0 {
t.Error("No chunks received")
}
t.Logf("Response with temperature=0.7: %+v", response)
t.Logf("Total chunks received: %d", chunkCount)
}
// ============================================================================
// Helper Functions
// ============================================================================
// newTestContext creates a real Context for testing OpenAI provider
func newTestContext(chatID, connectorID string) *context.Context {
return &context.Context{
Context: gocontext.Background(),
Space: plan.NewMemorySharedSpace(),
ChatID: chatID,
AssistantID: "test-assistant",
Connector: connectorID,
Locale: "en-us",
Theme: "light",
Client: context.Client{
Type: "web",
UserAgent: "OpenAIProviderTest/1.0",
IP: "127.0.0.1",
},
Referer: context.RefererAPI,
Accept: context.AcceptStandard,
Route: "/api/test",
Metadata: make(map[string]interface{}),
Authorized: &types.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client",
UserID: "test-user-123",
TeamID: "test-team-456",
TenantID: "test-tenant-789",
SessionID: "test-session-id",
Constraints: types.DataConstraints{
TeamOnly: true,
Extra: map[string]interface{}{
"test": "openai-provider",
},
},
},
}
}