From 38a277e33222ecb82e6ea7fcac519b269b84d9b7 Mon Sep 17 00:00:00 2001 From: presbrey Date: Tue, 4 Feb 2025 14:26:22 -0500 Subject: [PATCH] feat: Add OpenRouter API request and response structures with tests --- go.sum | 10 ++++ openrouter/openrouter-tool-calls.json | 35 +++++++++++ openrouter/openrouter_test.go | 41 +++++++++++++ openrouter/types.go | 84 +++++++++++++++++++++++++++ 4 files changed, 170 insertions(+) create mode 100644 go.sum create mode 100644 openrouter/openrouter-tool-calls.json create mode 100644 openrouter/openrouter_test.go create mode 100644 openrouter/types.go diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..713a0b4 --- /dev/null +++ b/go.sum @@ -0,0 +1,10 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/openrouter/openrouter-tool-calls.json b/openrouter/openrouter-tool-calls.json new file mode 100644 index 0000000..76ae79b --- /dev/null +++ b/openrouter/openrouter-tool-calls.json @@ -0,0 +1,35 @@ +{ + "id": "gen-1738521823-zoQZYL4vJW8s167DSle", + "provider": "OpenAI", + "model": "openai/gpt-4o-2024-11-20", + "object": "chat.completion", + "created": 1738521823, + "system_fingerprint": "fp_e53e529995", + "choices": [ + { + "logprobs": null, + "finish_reason": "tool_calls", + "native_finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_vdkI1mqVvMFDYDtx22Q6MAQl", + "type": "function", + "function": { + "name": "get_train_fare", + "arguments": "{\"date\":\"2025-05-25\",\"destination\":\"New York City\",\"class\":\"standard\",\"origin\":\"Boston\",\"passengers\":1}" + } + } + ] + } + } + ], + "usage": { + "prompt_tokens": 2296, + "completion_tokens": 43, + "total_tokens": 2339 + } + } \ No newline at end of file diff --git a/openrouter/openrouter_test.go b/openrouter/openrouter_test.go new file mode 100644 index 0000000..67334b4 --- /dev/null +++ b/openrouter/openrouter_test.go @@ -0,0 +1,41 @@ +package openrouter + +import ( + "encoding/json" + "os" + "testing" + + "github.com/presbrey/aichat" + "github.com/stretchr/testify/assert" +) + +func TestOpenRouterToolCalls(t *testing.T) { + b, err := os.ReadFile("openrouter-tool-calls.json") + assert.NoError(t, err) + + resp := &Response{} + assert.NoError(t, json.Unmarshal(b, resp)) + + chat := &aichat.Chat{} + chat.AddMessage(resp.Choices[0].Message) + + toolCallCount := 0 + chat.RangePendingToolCalls(func(tc *aichat.ToolCallContext) error { + toolCallCount++ + args, err := tc.Arguments() + assert.NoError(t, err) + assert.Equal(t, map[string]interface{}{ + "class": "standard", + "date": "2025-05-25", + "destination": "New York City", + "origin": "Boston", + "passengers": float64(1), + }, args) + return tc.Return(map[string]any{ + "foo": "bar", + }) + }) + assert.Equal(t, 1, toolCallCount) + + assert.Equal(t, `{"foo":"bar"}`, chat.LastMessageByRole("tool").Content) +} diff --git a/openrouter/types.go b/openrouter/types.go new file mode 100644 index 0000000..503ff54 --- /dev/null +++ b/openrouter/types.go @@ -0,0 +1,84 @@ +package openrouter + +import "github.com/presbrey/aichat" + +// Request represents a chat completion request to the OpenRouter API +type Request struct { + Messages []*aichat.Message `json:"messages,omitempty"` + Prompt string `json:"prompt,omitempty"` + Model string `json:"model,omitempty"` + ResponseFormat *struct { + Type string `json:"type"` + } `json:"response_format,omitempty"` + + Stop interface{} `json:"stop,omitempty"` // string or []string + Stream bool `json:"stream,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + + Tools any `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` // string or ToolChoice + + Seed int `json:"seed,omitempty"` + TopP float64 `json:"top_p,omitempty"` + TopK int `json:"top_k,omitempty"` + FreqPenalty float64 `json:"frequency_penalty,omitempty"` + PresPenalty float64 `json:"presence_penalty,omitempty"` + RepPenalty float64 `json:"repetition_penalty,omitempty"` + + LogitBias map[int]float64 `json:"logit_bias,omitempty"` + TopLogprobs int `json:"top_logprobs,omitempty"` + + MinP float64 `json:"min_p,omitempty"` + TopA float64 `json:"top_a,omitempty"` + + Prediction *struct { + Type string `json:"type"` + Content string `json:"content"` + } `json:"prediction,omitempty"` + Transforms []string `json:"transforms,omitempty"` + Models []string `json:"models,omitempty"` + Route string `json:"route,omitempty"` + + IncludeReasoning bool `json:"include_reasoning,omitempty"` +} + +// Response represents the API response structure +type Response struct { + Error *struct { + Message string `json:"message"` + Code int `json:"code"` + } `json:"error,omitempty"` + + UserID string `json:"user_id,omitempty"` + ID string `json:"id,omitempty"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + Object string `json:"object,omitempty"` + Created int64 `json:"created,omitempty"` + + SystemFingerprint string `json:"system_fingerprint,omitempty"` + + Choices []struct { + LogProbs interface{} `json:"logprobs"` + FinishReason string `json:"finish_reason"` + NativeFinishReason string `json:"native_finish_reason"` + Index int `json:"index"` + Message *aichat.Message `json:"message"` + } `json:"choices,omitempty"` + + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage,omitempty"` +} + +// ToolChoice represents the model's choice of tool usage +type ToolChoice struct { + Type string `json:"type,omitempty"` + + Function struct { + Name string `json:"name"` + } `json:"function,omitempty"` +}