136 lines
3.3 KiB
Go
136 lines
3.3 KiB
Go
package openapi
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestSessionToolTurns(t *testing.T) {
|
|
callCount := 0
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
callCount++
|
|
var resp ChatCompletionResponse
|
|
|
|
if callCount == 1 {
|
|
// First call: Model decides to call a tool
|
|
resp = ChatCompletionResponse{
|
|
ID: "resp1",
|
|
Choices: []Choice{
|
|
{
|
|
Message: Message{
|
|
Role: RoleAssistant,
|
|
ToolCalls: []ToolCall{
|
|
{
|
|
ID: "call_1",
|
|
Type: "function",
|
|
Function: ToolCallFunction{
|
|
Name: "get_weather",
|
|
Arguments: `{"location":"Stockholm"}`,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
FinishReason: "tool_calls",
|
|
},
|
|
},
|
|
}
|
|
} else {
|
|
// Second call: Model gives final response after tool result
|
|
resp = ChatCompletionResponse{
|
|
ID: "resp2",
|
|
Choices: []Choice{
|
|
{
|
|
Message: Message{
|
|
Role: RoleAssistant,
|
|
Content: "The weather in Stockholm is sunny.",
|
|
},
|
|
FinishReason: "stop",
|
|
},
|
|
},
|
|
}
|
|
}
|
|
json.NewEncoder(w).Encode(resp)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := NewClient(server.URL, "test-key")
|
|
session := NewSession(client, "test-model")
|
|
|
|
type weatherArgs struct {
|
|
Location string `json:"location"`
|
|
}
|
|
|
|
toolExecuted := false
|
|
session.RegisterTool("get_weather", "Get weather",
|
|
map[string]any{"type": "object", "properties": map[string]any{"location": map[string]any{"type": "string"}}},
|
|
func(ctx context.Context, s *Session, args json.RawMessage) (any, error) {
|
|
var w weatherArgs
|
|
json.Unmarshal(args, &w)
|
|
if w.Location != "Stockholm" {
|
|
t.Errorf("expected Stockholm, got %s", w.Location)
|
|
}
|
|
toolExecuted = true
|
|
return map[string]any{"temp": 20, "condition": "sunny"}, nil
|
|
})
|
|
|
|
hooksCalled := make(map[string]bool)
|
|
session.SetHooks(SessionHooks{
|
|
BeforeMessage: func(ctx context.Context, s *Session, msg *Message) error {
|
|
hooksCalled["BeforeMessage"] = true
|
|
return nil
|
|
},
|
|
AfterMessage: func(ctx context.Context, s *Session, msg *Message) error {
|
|
hooksCalled["AfterMessage"] = true
|
|
return nil
|
|
},
|
|
BeforeToolCall: func(ctx context.Context, s *Session, tc *ToolCall) error {
|
|
hooksCalled["BeforeToolCall"] = true
|
|
return nil
|
|
},
|
|
AfterToolCall: func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error {
|
|
hooksCalled["AfterToolCall"] = true
|
|
return nil
|
|
},
|
|
OnDone: func(ctx context.Context, s *Session) error {
|
|
hooksCalled["OnDone"] = true
|
|
return nil
|
|
},
|
|
})
|
|
|
|
ctx := context.Background()
|
|
chunks, errs := session.Chat(ctx, "How is the weather in Stockholm?")
|
|
|
|
var finalContent string
|
|
for chunk := range chunks {
|
|
if len(chunk.Choices) > 0 {
|
|
finalContent += chunk.Choices[0].Delta.Content
|
|
}
|
|
}
|
|
|
|
if err := <-errs; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if !toolExecuted {
|
|
t.Error("expected tool to be executed")
|
|
}
|
|
|
|
hookNames := []string{"BeforeMessage", "AfterMessage", "BeforeToolCall", "AfterToolCall", "OnDone"}
|
|
for _, name := range hookNames {
|
|
if !hooksCalled[name] {
|
|
t.Errorf("expected hook %s to be called", name)
|
|
}
|
|
}
|
|
|
|
if finalContent != "The weather in Stockholm is sunny." {
|
|
t.Errorf("expected final response, got %s", finalContent)
|
|
}
|
|
|
|
if len(session.Messages()) != 4 {
|
|
t.Errorf("expected 4 messages in history, got %d", len(session.Messages()))
|
|
}
|
|
}
|