Initial commit

pkg/openapi/session_test.go

@@ -0,0 +1,135 @@
+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()))
+	}
+}