Initial commit
This commit is contained in:
@@ -0,0 +1,253 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCreateChatCompletion(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Errorf("expected path /chat/completions, got %s", r.URL.Path)
|
||||
}
|
||||
if r.Header.Get("Authorization") != "Bearer test-key" {
|
||||
t.Errorf("expected bearer token, got %s", r.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
var req ChatCompletionRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !req.CachePrompt {
|
||||
t.Error("expected cache_prompt to be true")
|
||||
}
|
||||
|
||||
resp := ChatCompletionResponse{
|
||||
ID: "test-id",
|
||||
Model: "test-model",
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "Hello!",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
CachePrompt: true,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content != "Hello!" {
|
||||
t.Errorf("expected Hello!, got %s", resp.Choices[0].Message.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithReasoning(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "Hello!",
|
||||
ReasoningContent: "Thinking...",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.ReasoningContent != "Thinking..." {
|
||||
t.Errorf("expected Thinking..., got %s", resp.Choices[0].Message.ReasoningContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithToolCalls(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "call_123",
|
||||
Type: "function",
|
||||
Function: ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Stockholm"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "What is the weather in Stockholm?"},
|
||||
},
|
||||
Tools: []Tool{
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "get_weather",
|
||||
Description: "Get the current weather",
|
||||
Parameters: json.RawMessage(`{"type":"object","properties":{"location":{"type":"string"}}}`),
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if len(resp.Choices[0].Message.ToolCalls) != 1 {
|
||||
t.Fatalf("expected 1 tool call, got %d", len(resp.Choices[0].Message.ToolCalls))
|
||||
}
|
||||
|
||||
tc := resp.Choices[0].Message.ToolCalls[0]
|
||||
if tc.Function.Name != "get_weather" {
|
||||
t.Errorf("expected get_weather, got %s", tc.Function.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithToolResponse(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var req ChatCompletionRequest
|
||||
json.NewDecoder(r.Body).Decode(&req)
|
||||
|
||||
// Check if the tool response is present
|
||||
foundToolResponse := false
|
||||
for _, msg := range req.Messages {
|
||||
if msg.Role == RoleTool && msg.ToolCallID == "call_123" {
|
||||
foundToolResponse = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !foundToolResponse {
|
||||
t.Error("expected tool response message in request")
|
||||
}
|
||||
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "The weather in Stockholm is 15°C.",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "What is the weather in Stockholm?"},
|
||||
{
|
||||
Role: RoleAssistant,
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "call_123",
|
||||
Type: "function",
|
||||
Function: ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Stockholm"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Role: RoleTool,
|
||||
ToolCallID: "call_123",
|
||||
Content: `{"temp":15}`,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionStreamWithReasoning(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"reasoning_content":"Thinking"}}]} `)
|
||||
fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"Hello"}}]} `)
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
chunks, errs := client.CreateChatCompletionStream(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
})
|
||||
|
||||
var reasoning, content string
|
||||
for chunk := range chunks {
|
||||
if len(chunk.Choices) > 0 {
|
||||
reasoning += chunk.Choices[0].Delta.ReasoningContent
|
||||
content += chunk.Choices[0].Delta.Content
|
||||
}
|
||||
}
|
||||
|
||||
if err := <-errs; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if reasoning != "Thinking" {
|
||||
t.Errorf("expected Thinking, got %s", reasoning)
|
||||
}
|
||||
if content != "Hello" {
|
||||
t.Errorf("expected Hello, got %s", content)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user