feat(llm): add LLM client interface, types, providers (openai, llamacpp), and mock
This commit is contained in:
parent
ce64b68f12
commit
641481022f
6 changed files with 1263 additions and 0 deletions
121
pkg/llm/mock/mock.go
Normal file
121
pkg/llm/mock/mock.go
Normal file
|
|
@ -0,0 +1,121 @@
|
||||||
|
package mock
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"iter"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// MockLLMClient is a deterministic implementation of llm.LLMClient for testing.
|
||||||
|
type MockLLMClient struct {
|
||||||
|
GenerateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error)
|
||||||
|
StreamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error]
|
||||||
|
NameFunc func() string
|
||||||
|
CapabilitiesFunc func() llm.ProviderCapabilities
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockLLMClient) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
if m.GenerateFunc != nil {
|
||||||
|
return m.GenerateFunc(ctx, req)
|
||||||
|
}
|
||||||
|
return llm.CompletionResponse{
|
||||||
|
Content: "default response",
|
||||||
|
StopReason: llm.StopReasonEndTurn,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockLLMClient) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
|
if m.StreamFunc != nil {
|
||||||
|
return m.StreamFunc(ctx, req)
|
||||||
|
}
|
||||||
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
|
yield(llm.StreamChunk{
|
||||||
|
Delta: "default response",
|
||||||
|
}, nil)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockLLMClient) Name() string {
|
||||||
|
if m.NameFunc != nil {
|
||||||
|
return m.NameFunc()
|
||||||
|
}
|
||||||
|
return "mock"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *MockLLMClient) Capabilities() llm.ProviderCapabilities {
|
||||||
|
if m.CapabilitiesFunc != nil {
|
||||||
|
return m.CapabilitiesFunc()
|
||||||
|
}
|
||||||
|
return llm.ProviderCapabilities{
|
||||||
|
SupportsTools: true,
|
||||||
|
SupportsVision: false,
|
||||||
|
SupportsJSON: true,
|
||||||
|
MaxContextWindow: 128000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// New returns a MockLLMClient with default behaviors.
|
||||||
|
func New() *MockLLMClient {
|
||||||
|
return &MockLLMClient{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewWithGenerate returns a MockLLMClient that returns the given response.
|
||||||
|
func NewWithGenerate(resp llm.CompletionResponse) *MockLLMClient {
|
||||||
|
return &MockLLMClient{
|
||||||
|
GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
return resp, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewWithStream returns a MockLLMClient that streams the given chunks.
|
||||||
|
func NewWithStream(chunks []llm.StreamChunk) *MockLLMClient {
|
||||||
|
return &MockLLMClient{
|
||||||
|
StreamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
|
for _, c := range chunks {
|
||||||
|
if !yield(c, nil) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewWithMatch returns a MockLLMClient that matches input messages against patterns.
|
||||||
|
func NewWithMatch(responses []MatchResponse) *MockLLMClient {
|
||||||
|
return &MockLLMClient{
|
||||||
|
GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
for _, r := range responses {
|
||||||
|
if r.Match == "*" || matchesAny(r.Match, req.Messages) {
|
||||||
|
return llm.CompletionResponse{
|
||||||
|
Content: r.Response,
|
||||||
|
StopReason: llm.StopReasonEndTurn,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return llm.CompletionResponse{
|
||||||
|
Content: "no match",
|
||||||
|
StopReason: llm.StopReasonEndTurn,
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MatchResponse pairs a pattern with a response string.
|
||||||
|
type MatchResponse struct {
|
||||||
|
Match string // exact string or "*" for wildcard
|
||||||
|
Response string
|
||||||
|
}
|
||||||
|
|
||||||
|
func matchesAny(pattern string, msgs []llm.Message) bool {
|
||||||
|
for _, m := range msgs {
|
||||||
|
if m.Role == llm.RoleUser && strings.Contains(m.Content, pattern) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
322
pkg/llm/providers/llamacpp/client.go
Normal file
322
pkg/llm/providers/llamacpp/client.go
Normal file
|
|
@ -0,0 +1,322 @@
|
||||||
|
package llamacpp
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"iter"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Config holds the settings needed to create a llama.cpp client.
|
||||||
|
type Config struct {
|
||||||
|
BaseURL string // defaults to http://localhost:8080/v1
|
||||||
|
Model string
|
||||||
|
Timeout int // request timeout in seconds (0 = default)
|
||||||
|
TopK int // top-k sampling (0 = default)
|
||||||
|
TopP float32
|
||||||
|
Temperature float32
|
||||||
|
}
|
||||||
|
|
||||||
|
// Client implements llm.LLMClient for llama.cpp.
|
||||||
|
type Client struct {
|
||||||
|
baseURL string
|
||||||
|
http *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// New returns a new llama.cpp client.
|
||||||
|
func New(cfg Config) (*Client, error) {
|
||||||
|
baseURL := cfg.BaseURL
|
||||||
|
if baseURL == "" {
|
||||||
|
baseURL = "http://localhost:8080/v1"
|
||||||
|
}
|
||||||
|
|
||||||
|
return &Client{
|
||||||
|
baseURL: baseURL,
|
||||||
|
http: http.DefaultClient,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
|
payload, err := c.buildRequest(req, false)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("creating request: %w", err)
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := c.http.Do(httpReq)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("request failed: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
|
||||||
|
}
|
||||||
|
|
||||||
|
var apiResp llamaChatResponse
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return c.toResponse(apiResp), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
|
payload, err := c.buildRequest(req, true)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("creating request: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := c.http.Do(httpReq)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("request failed: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
scanner := bufio.NewScanner(resp.Body)
|
||||||
|
for scanner.Scan() {
|
||||||
|
line := scanner.Text()
|
||||||
|
if !strings.HasPrefix(line, "data: ") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
data := strings.TrimPrefix(line, "data: ")
|
||||||
|
if data == "[DONE]" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var event llamaStreamEvent
|
||||||
|
if err := json.Unmarshal([]byte(data), &event); err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, choice := range event.Choices {
|
||||||
|
chunk := llm.StreamChunk{
|
||||||
|
Delta: choice.Delta.Content,
|
||||||
|
}
|
||||||
|
if choice.FinishReason != "" {
|
||||||
|
chunk.FinishReason = choice.FinishReason
|
||||||
|
}
|
||||||
|
if !yield(chunk, nil) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := scanner.Err(); err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Name() string {
|
||||||
|
return "llama.cpp"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Capabilities() llm.ProviderCapabilities {
|
||||||
|
return llm.ProviderCapabilities{
|
||||||
|
SupportsTools: true,
|
||||||
|
SupportsVision: false,
|
||||||
|
SupportsJSON: true,
|
||||||
|
MaxContextWindow: 32768,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildRequest converts an llm.CompletionRequest to the llama.cpp API format.
|
||||||
|
func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
|
||||||
|
messages := make([]llamaMessage, len(req.Messages))
|
||||||
|
for i, m := range req.Messages {
|
||||||
|
messages[i] = llamaMessage{
|
||||||
|
Role: string(m.Role),
|
||||||
|
Content: m.Content,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var tools []llamaTool
|
||||||
|
for _, t := range req.Tools {
|
||||||
|
var tool llamaTool
|
||||||
|
if err := json.Unmarshal(t, &tool); err != nil {
|
||||||
|
return nil, fmt.Errorf("parsing tool: %w", err)
|
||||||
|
}
|
||||||
|
tools = append(tools, tool)
|
||||||
|
}
|
||||||
|
|
||||||
|
openReq := llamaChatRequest{
|
||||||
|
Model: req.Model,
|
||||||
|
Messages: messages,
|
||||||
|
Stream: stream,
|
||||||
|
}
|
||||||
|
if len(tools) > 0 {
|
||||||
|
openReq.Tools = tools
|
||||||
|
}
|
||||||
|
if req.ToolChoice != nil {
|
||||||
|
openReq.ToolChoice = req.ToolChoice
|
||||||
|
}
|
||||||
|
if req.Temperature != nil {
|
||||||
|
tmp := *req.Temperature
|
||||||
|
openReq.Temperature = tmp
|
||||||
|
}
|
||||||
|
if req.MaxTokens != nil {
|
||||||
|
tmp := *req.MaxTokens
|
||||||
|
openReq.MaxTokens = tmp
|
||||||
|
}
|
||||||
|
if len(req.Stop) > 0 {
|
||||||
|
openReq.Stop = req.Stop
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(openReq)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("marshaling request: %w", err)
|
||||||
|
}
|
||||||
|
return strings.NewReader(string(data)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// toResponse converts a llama.cpp API response to our CompletionResponse.
|
||||||
|
func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
|
||||||
|
choice := resp.Choices[0]
|
||||||
|
result := llm.CompletionResponse{
|
||||||
|
ID: resp.ID,
|
||||||
|
Model: resp.Model,
|
||||||
|
Content: choice.Message.Content,
|
||||||
|
StopReason: choice.FinishReason,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range choice.Message.ToolCalls {
|
||||||
|
result.ToolCalls = append(result.ToolCalls, llm.ToolCall{
|
||||||
|
ID: tc.ID,
|
||||||
|
Name: tc.Function.Name,
|
||||||
|
Arguments: json.RawMessage(tc.Function.Arguments),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
result.Usage = llm.TokenUsage{
|
||||||
|
InputTokens: resp.Usage.PromptTokens,
|
||||||
|
OutputTokens: resp.Usage.CompletionTokens,
|
||||||
|
TotalTokens: resp.Usage.TotalTokens,
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// llama.cpp API types
|
||||||
|
|
||||||
|
type llamaChatRequest struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Messages []llamaMessage `json:"messages"`
|
||||||
|
Tools []llamaTool `json:"tools,omitempty"`
|
||||||
|
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||||
|
Temperature float32 `json:"temperature,omitempty"`
|
||||||
|
MaxTokens int `json:"max_tokens,omitempty"`
|
||||||
|
TopK int `json:"top_k,omitempty"`
|
||||||
|
TopP float32 `json:"top_p,omitempty"`
|
||||||
|
Stop []string `json:"stop,omitempty"`
|
||||||
|
Stream bool `json:"stream"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaMessage struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaTool struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function json.RawMessage `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaChatResponse struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Choices []llamaChoice `json:"choices"`
|
||||||
|
Usage llamaUsage `json:"usage"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaChoice struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
Message llamaMessageResult `json:"message"`
|
||||||
|
FinishReason string `json:"finish_reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaMessageResult struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ToolCalls []llamaToolCall `json:"tool_calls"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaToolCall struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function llamaFunction `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaFunction struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaUsage struct {
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
TotalTokens int `json:"total_tokens"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stream event types
|
||||||
|
|
||||||
|
type llamaStreamEvent struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Choices []llamaStreamChoice `json:"choices"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaStreamChoice struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
Delta llamaStreamDelta `json:"delta"`
|
||||||
|
FinishReason string `json:"finish_reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaStreamDelta struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
ToolCalls []llamaStreamToolCall `json:"tool_calls"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaStreamToolCall struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function llamaStreamFunction `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type llamaStreamFunction struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
}
|
||||||
232
pkg/llm/providers/llamacpp/client_test.go
Normal file
232
pkg/llm/providers/llamacpp/client_test.go
Normal file
|
|
@ -0,0 +1,232 @@
|
||||||
|
package llamacpp
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestClient_Name(t *testing.T) {
|
||||||
|
client := &Client{}
|
||||||
|
if client.Name() != "llama.cpp" {
|
||||||
|
t.Errorf("expected 'llama.cpp', got %q", client.Name())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Capabilities(t *testing.T) {
|
||||||
|
client := &Client{}
|
||||||
|
caps := client.Capabilities()
|
||||||
|
if !caps.SupportsTools {
|
||||||
|
t.Error("expected SupportsTools to be true")
|
||||||
|
}
|
||||||
|
if caps.SupportsVision {
|
||||||
|
t.Error("expected SupportsVision to be false")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Generate(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(llamaChatResponse{
|
||||||
|
ID: "llama-123",
|
||||||
|
Model: "llama3",
|
||||||
|
Choices: []llamaChoice{
|
||||||
|
{
|
||||||
|
Index: 0,
|
||||||
|
FinishReason: "stop",
|
||||||
|
Message: llamaMessageResult{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: "Hello from llama.cpp!",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Usage: llamaUsage{
|
||||||
|
PromptTokens: 10,
|
||||||
|
CompletionTokens: 5,
|
||||||
|
TotalTokens: 15,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.Generate(context.Background(), llm.CompletionRequest{
|
||||||
|
Model: "llama3",
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "Hi"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if resp.Content != "Hello from llama.cpp!" {
|
||||||
|
t.Errorf("expected 'Hello from llama.cpp!', got %q", resp.Content)
|
||||||
|
}
|
||||||
|
if resp.ID != "llama-123" {
|
||||||
|
t.Errorf("expected 'llama-123', got %q", resp.ID)
|
||||||
|
}
|
||||||
|
if resp.Usage.InputTokens != 10 {
|
||||||
|
t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Generate_Error(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
w.Write([]byte("server error"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = client.Generate(context.Background(), llm.CompletionRequest{})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Generate_ToolCall(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(llamaChatResponse{
|
||||||
|
ID: "llama-tool-1",
|
||||||
|
Model: "llama3",
|
||||||
|
Choices: []llamaChoice{
|
||||||
|
{
|
||||||
|
Index: 0,
|
||||||
|
FinishReason: "tool_calls",
|
||||||
|
Message: llamaMessageResult{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: "",
|
||||||
|
ToolCalls: []llamaToolCall{
|
||||||
|
{
|
||||||
|
ID: "call-1",
|
||||||
|
Type: "function",
|
||||||
|
Function: llamaFunction{
|
||||||
|
Name: "calculate",
|
||||||
|
Arguments: `{"a":1,"b":2}`,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Usage: llamaUsage{
|
||||||
|
PromptTokens: 20,
|
||||||
|
CompletionTokens: 10,
|
||||||
|
TotalTokens: 30,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.Generate(context.Background(), llm.CompletionRequest{
|
||||||
|
Model: "llama3",
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "What is 1 + 2?"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if len(resp.ToolCalls) != 1 {
|
||||||
|
t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls))
|
||||||
|
}
|
||||||
|
if resp.ToolCalls[0].Name != "calculate" {
|
||||||
|
t.Errorf("expected 'calculate', got %q", resp.ToolCalls[0].Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Stream(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n"))
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n"))
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"!\"},\"finish_reason\":\"stop\"}]}\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
stream := client.Stream(context.Background(), llm.CompletionRequest{})
|
||||||
|
for chunk, err := range stream {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 3 {
|
||||||
|
t.Errorf("expected 3 chunks, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
if chunks[0].Delta != "Hello" {
|
||||||
|
t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
|
||||||
|
}
|
||||||
|
if chunks[1].Delta != " world" {
|
||||||
|
t.Errorf("expected ' world', got %q", chunks[1].Delta)
|
||||||
|
}
|
||||||
|
if chunks[2].Delta != "!" {
|
||||||
|
t.Errorf("expected '!', got %q", chunks[2].Delta)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Stream_FinishReason(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}]}\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
stream := client.Stream(context.Background(), llm.CompletionRequest{})
|
||||||
|
for chunk, err := range stream {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 1 {
|
||||||
|
t.Fatalf("expected 1 chunk, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
if chunks[0].FinishReason != "stop" {
|
||||||
|
t.Errorf("expected 'stop' finish reason, got %q", chunks[0].FinishReason)
|
||||||
|
}
|
||||||
|
}
|
||||||
327
pkg/llm/providers/openai/client.go
Normal file
327
pkg/llm/providers/openai/client.go
Normal file
|
|
@ -0,0 +1,327 @@
|
||||||
|
package openai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"iter"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Config holds the settings needed to create an OpenAI client.
|
||||||
|
type Config struct {
|
||||||
|
APIKey string
|
||||||
|
Model string
|
||||||
|
BaseURL string // defaults to https://api.openai.com/v1
|
||||||
|
}
|
||||||
|
|
||||||
|
// Client implements llm.LLMClient for OpenAI.
|
||||||
|
type Client struct {
|
||||||
|
apiKey string
|
||||||
|
baseURL string
|
||||||
|
http *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// New returns a new OpenAI client.
|
||||||
|
func New(cfg Config) (*Client, error) {
|
||||||
|
if cfg.APIKey == "" {
|
||||||
|
return nil, fmt.Errorf("openai: API key is required")
|
||||||
|
}
|
||||||
|
|
||||||
|
baseURL := cfg.BaseURL
|
||||||
|
if baseURL == "" {
|
||||||
|
baseURL = "https://api.openai.com/v1"
|
||||||
|
}
|
||||||
|
|
||||||
|
return &Client{
|
||||||
|
apiKey: cfg.APIKey,
|
||||||
|
baseURL: baseURL,
|
||||||
|
http: http.DefaultClient,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||||
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
|
payload, err := c.buildRequest(req)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("creating request: %w", err)
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := c.http.Do(httpReq)
|
||||||
|
if err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("request failed: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
|
||||||
|
}
|
||||||
|
|
||||||
|
var apiResp openaiChatResponse
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
|
||||||
|
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return c.toResponse(apiResp), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||||
|
return func(yield func(llm.StreamChunk, error) bool) {
|
||||||
|
endpoint := c.baseURL + "/chat/completions"
|
||||||
|
|
||||||
|
payload, err := c.buildRequest(req)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("creating request: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
httpReq.Header.Set("Accept", "text/event-stream")
|
||||||
|
|
||||||
|
resp, err := c.http.Do(httpReq)
|
||||||
|
if err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("request failed: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
scanner := bufio.NewScanner(resp.Body)
|
||||||
|
for scanner.Scan() {
|
||||||
|
line := scanner.Text()
|
||||||
|
if !strings.HasPrefix(line, "data: ") {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
data := strings.TrimPrefix(line, "data: ")
|
||||||
|
if data == "[DONE]" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var event openaiStreamEvent
|
||||||
|
if err := json.Unmarshal([]byte(data), &event); err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, choice := range event.Choices {
|
||||||
|
chunk := llm.StreamChunk{
|
||||||
|
Delta: choice.Delta.Content,
|
||||||
|
}
|
||||||
|
if choice.FinishReason != "" {
|
||||||
|
chunk.FinishReason = choice.FinishReason
|
||||||
|
}
|
||||||
|
if !yield(chunk, nil) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := scanner.Err(); err != nil {
|
||||||
|
yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Name() string {
|
||||||
|
return "openai"
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Client) Capabilities() llm.ProviderCapabilities {
|
||||||
|
return llm.ProviderCapabilities{
|
||||||
|
SupportsTools: true,
|
||||||
|
SupportsVision: true,
|
||||||
|
SupportsJSON: true,
|
||||||
|
MaxContextWindow: 128000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildRequest converts an llm.CompletionRequest to the OpenAI API format.
|
||||||
|
func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
||||||
|
// Convert messages to OpenAI format
|
||||||
|
messages := make([]openaiMessage, len(req.Messages))
|
||||||
|
for i, m := range req.Messages {
|
||||||
|
messages[i] = openaiMessage{
|
||||||
|
Role: string(m.Role),
|
||||||
|
Content: m.Content,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tools := make([]openaiTool, len(req.Tools))
|
||||||
|
for i, t := range req.Tools {
|
||||||
|
var tool openaiTool
|
||||||
|
if err := json.Unmarshal(t, &tool); err != nil {
|
||||||
|
return nil, fmt.Errorf("parsing tool %d: %w", i, err)
|
||||||
|
}
|
||||||
|
tools[i] = tool
|
||||||
|
}
|
||||||
|
|
||||||
|
openaiReq := openaiChatRequest{
|
||||||
|
Model: req.Model,
|
||||||
|
Messages: messages,
|
||||||
|
Stream: false,
|
||||||
|
}
|
||||||
|
if len(tools) > 0 {
|
||||||
|
openaiReq.Tools = tools
|
||||||
|
}
|
||||||
|
if req.ToolChoice != nil {
|
||||||
|
openaiReq.ToolChoice = req.ToolChoice
|
||||||
|
}
|
||||||
|
if req.Temperature != nil {
|
||||||
|
tmp := *req.Temperature
|
||||||
|
openaiReq.Temperature = &tmp
|
||||||
|
}
|
||||||
|
if req.MaxTokens != nil {
|
||||||
|
max := *req.MaxTokens
|
||||||
|
openaiReq.MaxTokens = &max
|
||||||
|
}
|
||||||
|
if len(req.Stop) > 0 {
|
||||||
|
openaiReq.Stop = req.Stop
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(openaiReq)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("marshaling request: %w", err)
|
||||||
|
}
|
||||||
|
return strings.NewReader(string(data)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// toResponse converts an OpenAI API response to our CompletionResponse.
|
||||||
|
func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
|
||||||
|
choice := resp.Choices[0]
|
||||||
|
result := llm.CompletionResponse{
|
||||||
|
ID: resp.ID,
|
||||||
|
Model: resp.Model,
|
||||||
|
Content: choice.Message.Content,
|
||||||
|
StopReason: choice.FinishReason,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range choice.Message.ToolCalls {
|
||||||
|
result.ToolCalls = append(result.ToolCalls, llm.ToolCall{
|
||||||
|
ID: tc.ID,
|
||||||
|
Name: tc.Function.Name,
|
||||||
|
Arguments: json.RawMessage(tc.Function.Arguments),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
result.Usage = llm.TokenUsage{
|
||||||
|
InputTokens: resp.Usage.PromptTokens,
|
||||||
|
OutputTokens: resp.Usage.CompletionTokens,
|
||||||
|
TotalTokens: resp.Usage.TotalTokens,
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// OpenAI API types
|
||||||
|
|
||||||
|
type openaiChatRequest struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Messages []openaiMessage `json:"messages"`
|
||||||
|
Tools []openaiTool `json:"tools,omitempty"`
|
||||||
|
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||||
|
Temperature *float32 `json:"temperature,omitempty"`
|
||||||
|
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||||
|
Stop []string `json:"stop,omitempty"`
|
||||||
|
Stream bool `json:"stream"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiMessage struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiTool struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function json.RawMessage `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiChatResponse struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Choices []openaiChoice `json:"choices"`
|
||||||
|
Usage openaiUsage `json:"usage"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiChoice struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
Message openaiMessageResult `json:"message"`
|
||||||
|
FinishReason string `json:"finish_reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiMessageResult struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ToolCalls []openaiToolCall `json:"tool_calls"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiToolCall struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function openaiFunction `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiFunction struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiUsage struct {
|
||||||
|
PromptTokens int `json:"prompt_tokens"`
|
||||||
|
CompletionTokens int `json:"completion_tokens"`
|
||||||
|
TotalTokens int `json:"total_tokens"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stream event types
|
||||||
|
|
||||||
|
type openaiStreamEvent struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Choices []openaiStreamChoice `json:"choices"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiStreamChoice struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
Delta openaiStreamDelta `json:"delta"`
|
||||||
|
FinishReason string `json:"finish_reason"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiStreamDelta struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
ToolCalls []openaiStreamToolCall `json:"tool_calls"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiStreamToolCall struct {
|
||||||
|
Index int `json:"index"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Function openaiStreamFunction `json:"function"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type openaiStreamFunction struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments string `json:"arguments"`
|
||||||
|
}
|
||||||
141
pkg/llm/providers/openai/client_test.go
Normal file
141
pkg/llm/providers/openai/client_test.go
Normal file
|
|
@ -0,0 +1,141 @@
|
||||||
|
package openai
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestClient_Name(t *testing.T) {
|
||||||
|
client := &Client{}
|
||||||
|
if client.Name() != "openai" {
|
||||||
|
t.Errorf("expected 'openai', got %q", client.Name())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Capabilities(t *testing.T) {
|
||||||
|
client := &Client{}
|
||||||
|
caps := client.Capabilities()
|
||||||
|
if !caps.SupportsTools {
|
||||||
|
t.Error("expected SupportsTools to be true")
|
||||||
|
}
|
||||||
|
if !caps.SupportsVision {
|
||||||
|
t.Error("expected SupportsVision to be true")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Generate(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
json.NewEncoder(w).Encode(openaiChatResponse{
|
||||||
|
ID: "test-123",
|
||||||
|
Model: "gpt-4",
|
||||||
|
Choices: []openaiChoice{
|
||||||
|
{
|
||||||
|
Index: 0,
|
||||||
|
FinishReason: "stop",
|
||||||
|
Message: openaiMessageResult{
|
||||||
|
Role: "assistant",
|
||||||
|
Content: "Hello!",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Usage: openaiUsage{
|
||||||
|
PromptTokens: 10,
|
||||||
|
CompletionTokens: 5,
|
||||||
|
TotalTokens: 15,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
APIKey: "test-key",
|
||||||
|
BaseURL: server.URL,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := client.Generate(context.Background(), llm.CompletionRequest{
|
||||||
|
Model: "gpt-4",
|
||||||
|
Messages: []llm.Message{
|
||||||
|
{Role: llm.RoleUser, Content: "Hi"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if resp.Content != "Hello!" {
|
||||||
|
t.Errorf("expected 'Hello!', got %q", resp.Content)
|
||||||
|
}
|
||||||
|
if resp.ID != "test-123" {
|
||||||
|
t.Errorf("expected 'test-123', got %q", resp.ID)
|
||||||
|
}
|
||||||
|
if resp.Usage.InputTokens != 10 {
|
||||||
|
t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Generate_Error(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusUnauthorized)
|
||||||
|
w.Write([]byte("invalid api key"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
APIKey: "bad-key",
|
||||||
|
BaseURL: server.URL,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, err = client.Generate(context.Background(), llm.CompletionRequest{})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClient_Stream(t *testing.T) {
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n"))
|
||||||
|
w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n"))
|
||||||
|
w.Write([]byte("data: [DONE]\n"))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client, err := New(Config{
|
||||||
|
APIKey: "test-key",
|
||||||
|
BaseURL: server.URL + "/v1",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var chunks []llm.StreamChunk
|
||||||
|
stream := client.Stream(context.Background(), llm.CompletionRequest{})
|
||||||
|
for chunk, err := range stream {
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
chunks = append(chunks, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(chunks) != 2 {
|
||||||
|
t.Errorf("expected 2 chunks, got %d", len(chunks))
|
||||||
|
}
|
||||||
|
if chunks[0].Delta != "Hello" {
|
||||||
|
t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
|
||||||
|
}
|
||||||
|
if chunks[1].Delta != " world" {
|
||||||
|
t.Errorf("expected ' world', got %q", chunks[1].Delta)
|
||||||
|
}
|
||||||
|
}
|
||||||
120
pkg/llm/types.go
Normal file
120
pkg/llm/types.go
Normal file
|
|
@ -0,0 +1,120 @@
|
||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"iter"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Role represents the role of a message participant.
|
||||||
|
type Role string
|
||||||
|
|
||||||
|
const (
|
||||||
|
RoleSystem Role = "system"
|
||||||
|
RoleUser Role = "user"
|
||||||
|
RoleAssistant Role = "assistant"
|
||||||
|
RoleTool Role = "tool"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Message is a single message in a conversation.
|
||||||
|
type Message struct {
|
||||||
|
Role Role `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||||
|
Name string `json:"name,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ProviderCapabilities describes what a model supports.
|
||||||
|
type ProviderCapabilities struct {
|
||||||
|
SupportsTools bool `json:"supports_tools"`
|
||||||
|
SupportsVision bool `json:"supports_vision"`
|
||||||
|
SupportsJSON bool `json:"supports_json"`
|
||||||
|
MaxContextWindow int `json:"max_context_window"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TokenUsage tracks token consumption for a response.
|
||||||
|
type TokenUsage struct {
|
||||||
|
InputTokens int `json:"input_tokens"`
|
||||||
|
OutputTokens int `json:"output_tokens"`
|
||||||
|
TotalTokens int `json:"total_tokens"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToolChoice controls how the model chooses which tools to call.
|
||||||
|
type ToolChoice string
|
||||||
|
|
||||||
|
const (
|
||||||
|
ToolChoiceAuto ToolChoice = "auto"
|
||||||
|
ToolChoiceNone ToolChoice = "none"
|
||||||
|
ToolChoiceRequired ToolChoice = "required"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ToolRef is used when ToolChoice must invoke a specific tool by name.
|
||||||
|
type ToolRef struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *ToolRef) MarshalJSON() ([]byte, error) {
|
||||||
|
type Alias ToolRef
|
||||||
|
return json.Marshal(&struct {
|
||||||
|
*Alias
|
||||||
|
Type string `json:"type"`
|
||||||
|
}{
|
||||||
|
Alias: (*Alias)(t),
|
||||||
|
Type: "function",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// CompletionRequest is sent to an LLM provider.
|
||||||
|
type CompletionRequest struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Messages []Message `json:"messages"`
|
||||||
|
Tools []json.RawMessage `json:"tools,omitempty"`
|
||||||
|
ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null
|
||||||
|
Temperature *float32 `json:"temperature,omitempty"`
|
||||||
|
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||||
|
Stop []string `json:"stop,omitempty"`
|
||||||
|
Metadata map[string]string `json:"metadata,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CompletionResponse is returned from an LLM provider.
|
||||||
|
type CompletionResponse struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
|
StopReason string `json:"stop_reason"`
|
||||||
|
Usage TokenUsage `json:"usage,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// StopReason values.
|
||||||
|
const (
|
||||||
|
StopReasonEndTurn = "end_turn"
|
||||||
|
StopReasonToolUse = "tool_use"
|
||||||
|
StopReasonMaxTokens = "max_tokens"
|
||||||
|
StopReasonStopSeq = "stop_sequence"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ToolCall represents a function invocation requested by the model.
|
||||||
|
type ToolCall struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Arguments json.RawMessage `json:"arguments"`
|
||||||
|
Thought string `json:"thought,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// StreamChunk is emitted by the iterator returned from Stream().
|
||||||
|
type StreamChunk struct {
|
||||||
|
Delta string `json:"delta"`
|
||||||
|
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||||
|
FinishReason string `json:"finish_reason,omitempty"`
|
||||||
|
Usage TokenUsage `json:"usage,omitempty"` // only present on final chunk
|
||||||
|
}
|
||||||
|
|
||||||
|
// LLMClient is the interface that every provider implements.
|
||||||
|
// Products must never call a provider SDK directly; they go through this interface.
|
||||||
|
type LLMClient interface {
|
||||||
|
Generate(ctx context.Context, req CompletionRequest) (CompletionResponse, error)
|
||||||
|
Stream(ctx context.Context, req CompletionRequest) iter.Seq2[StreamChunk, error]
|
||||||
|
Name() string
|
||||||
|
Capabilities() ProviderCapabilities
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue