diff options
Diffstat (limited to 'internal/llm')
| -rw-r--r-- | internal/llm/client.go | 37 | ||||
| -rw-r--r-- | internal/llm/client_test.go | 40 |
2 files changed, 77 insertions, 0 deletions
diff --git a/internal/llm/client.go b/internal/llm/client.go index 613ebe5..e3a6102 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -34,6 +34,38 @@ type Client struct { type Message struct { Role string `json:"role"` Content string `json:"content"` + // ToolCalls is set on assistant messages that request tool invocations. + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + // ToolCallID links a role:"tool" result message to the call it answers. + ToolCallID string `json:"tool_call_id,omitempty"` + // Name carries the tool name on role:"tool" result messages. + Name string `json:"name,omitempty"` +} + +// Tool declares a function the model may call (OpenAI function-calling format). +type Tool struct { + Type string `json:"type"` // always "function" + Function ToolFunction `json:"function"` +} + +// ToolFunction describes a callable function and its JSON-schema parameters. +type ToolFunction struct { + Name string `json:"name"` + Description string `json:"description"` + Parameters map[string]any `json:"parameters,omitempty"` +} + +// ToolCall is a function invocation the model emitted in its response. +type ToolCall struct { + ID string `json:"id"` + Type string `json:"type"` // "function" + Function ToolCallFunction `json:"function"` +} + +// ToolCallFunction holds the called function's name and raw JSON arguments. +type ToolCallFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` } // ChatRequest captures the parameters of a single Chat or ChatStream call. @@ -45,6 +77,7 @@ type ChatRequest struct { Temperature *float64 MaxTokens int ResponseJSON bool + Tools []Tool } // ChatResponse is the aggregated result of a chat completion. @@ -54,6 +87,7 @@ type ChatResponse struct { OutputTokens int Model string FinishReason string + ToolCalls []ToolCall } // Chat performs a non-streaming chat completion. Rate-limit errors (HTTP 429, @@ -87,6 +121,7 @@ func (c *Client) Chat(ctx context.Context, req ChatRequest) (*ChatResponse, erro OutputTokens: oai.Usage.CompletionTokens, Model: oai.Model, FinishReason: oai.Choices[0].FinishReason, + ToolCalls: oai.Choices[0].Message.ToolCalls, } return nil }) @@ -133,6 +168,7 @@ func (c *Client) buildRequestBody(req ChatRequest, stream bool) ([]byte, error) Model: model, Messages: req.Messages, Stream: stream, + Tools: req.Tools, } if req.Temperature != nil { payload.Temperature = req.Temperature @@ -301,6 +337,7 @@ type openAIRequest struct { Stream bool `json:"stream,omitempty"` StreamOptions *streamOptions `json:"stream_options,omitempty"` ResponseFormat *responseFormat `json:"response_format,omitempty"` + Tools []Tool `json:"tools,omitempty"` } type streamOptions struct { diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go index 8257836..7533a8a 100644 --- a/internal/llm/client_test.go +++ b/internal/llm/client_test.go @@ -157,3 +157,43 @@ func TestErrFromStatus_RateLimitMarker(t *testing.T) { t.Errorf("error should embed retry-after, got: %v", err) } } + +func TestChat_ToolsRoundTrip(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body openAIRequest + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode body: %v", err) + } + if len(body.Tools) != 1 || body.Tools[0].Function.Name != "do_thing" { + t.Errorf("tools not forwarded in request: %+v", body.Tools) + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintln(w, `{ + "model": "m", + "choices": [{"message": {"role": "assistant", "content": "", + "tool_calls": [{"id": "call_1", "type": "function", + "function": {"name": "do_thing", "arguments": "{\"x\":1}"}}]}, + "finish_reason": "tool_calls"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 4} + }`) + })) + defer srv.Close() + + c := &Client{Endpoint: srv.URL + "/v1", Model: "m"} + resp, err := c.Chat(context.Background(), ChatRequest{ + Messages: []Message{{Role: "user", Content: "go"}}, + Tools: []Tool{{Type: "function", Function: ToolFunction{ + Name: "do_thing", Description: "d", Parameters: map[string]any{"type": "object"}, + }}}, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } + tc := resp.ToolCalls[0] + if tc.ID != "call_1" || tc.Function.Name != "do_thing" || tc.Function.Arguments != `{"x":1}` { + t.Errorf("unexpected tool call parsed: %+v", tc) + } +} |