Add LLM interface

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-11 22:33:12 +01:00
parent 307ca48394
commit 158c36d9ab
22 changed files with 2822 additions and 91 deletions
--- a/pkg/llm/bedrock/provider.go
+++ b/pkg/llm/bedrock/provider.go
@@ -0,0 +1,446 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package bedrock
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"go.probo.inc/probo/pkg/llm"
+)
+
+type (
+	Provider struct {
+		client *bedrockruntime.Client
+	}
+
+	Option func(*bedrockruntime.Options)
+)
+
+// WithBaseEndpoint overrides the Bedrock service endpoint URL.
+func WithBaseEndpoint(url string) Option {
+	return func(o *bedrockruntime.Options) { o.BaseEndpoint = &url }
+}
+
+func NewProvider(cfg aws.Config, opts ...Option) *Provider {
+	fns := make([]func(*bedrockruntime.Options), len(opts))
+	for i, o := range opts {
+		fns[i] = func(bo *bedrockruntime.Options) { o(bo) }
+	}
+
+	client := bedrockruntime.NewFromConfig(cfg, fns...)
+	return &Provider{client: client}
+}
+
+func (p *Provider) ChatCompletion(ctx context.Context, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
+	input := buildInput(req)
+
+	output, err := p.client.Converse(ctx, input)
+	if err != nil {
+		return nil, mapError(err)
+	}
+
+	return mapResponse(output, req.Model), nil
+}
+
+func (p *Provider) ChatCompletionStream(ctx context.Context, req *llm.ChatCompletionRequest) (llm.ChatCompletionStream, error) {
+	input := &bedrockruntime.ConverseStreamInput{
+		ModelId:         aws.String(req.Model),
+		Messages:        buildMessages(req.Messages),
+		InferenceConfig: buildInferenceConfig(req),
+	}
+
+	system := buildSystem(req.Messages)
+	if len(system) > 0 {
+		input.System = system
+	}
+
+	if len(req.Tools) > 0 {
+		toolConfig := buildToolConfig(req)
+		input.ToolConfig = toolConfig
+	}
+
+	output, err := p.client.ConverseStream(ctx, input)
+	if err != nil {
+		return nil, mapError(err)
+	}
+
+	return newBedrockStream(output.GetStream()), nil
+}
+
+func buildInput(req *llm.ChatCompletionRequest) *bedrockruntime.ConverseInput {
+	input := &bedrockruntime.ConverseInput{
+		ModelId:         aws.String(req.Model),
+		Messages:        buildMessages(req.Messages),
+		InferenceConfig: buildInferenceConfig(req),
+	}
+
+	system := buildSystem(req.Messages)
+	if len(system) > 0 {
+		input.System = system
+	}
+
+	if len(req.Tools) > 0 {
+		toolConfig := buildToolConfig(req)
+		input.ToolConfig = toolConfig
+	}
+
+	return input
+}
+
+func buildInferenceConfig(req *llm.ChatCompletionRequest) *types.InferenceConfiguration {
+	cfg := &types.InferenceConfiguration{}
+
+	if req.MaxTokens != nil {
+		v := int32(*req.MaxTokens)
+		cfg.MaxTokens = &v
+	}
+	if req.Temperature != nil {
+		v := float32(*req.Temperature)
+		cfg.Temperature = &v
+	}
+	if req.TopP != nil {
+		v := float32(*req.TopP)
+		cfg.TopP = &v
+	}
+	if len(req.StopSequences) > 0 {
+		cfg.StopSequences = req.StopSequences
+	}
+
+	return cfg
+}
+
+func buildSystem(messages []llm.Message) []types.SystemContentBlock {
+	var system []types.SystemContentBlock
+	for _, msg := range messages {
+		if msg.Role == llm.RoleSystem {
+			system = append(
+				system,
+				&types.SystemContentBlockMemberText{
+					Value: msg.Text(),
+				},
+			)
+		}
+	}
+	return system
+}
+
+func buildMessages(messages []llm.Message) []types.Message {
+	var out []types.Message
+
+	for _, msg := range messages {
+		switch msg.Role {
+		case llm.RoleSystem:
+			continue
+		case llm.RoleUser:
+			var content []types.ContentBlock
+			for _, p := range msg.Parts {
+				if tp, ok := p.(llm.TextPart); ok {
+					content = append(content, &types.ContentBlockMemberText{Value: tp.Text})
+				}
+			}
+			out = append(
+				out, types.Message{
+					Role:    types.ConversationRoleUser,
+					Content: content,
+				},
+			)
+
+		case llm.RoleAssistant:
+			var content []types.ContentBlock
+			if text := msg.Text(); text != "" {
+				content = append(content, &types.ContentBlockMemberText{Value: text})
+			}
+			for _, tc := range msg.ToolCalls {
+				var input any
+				_ = json.Unmarshal([]byte(tc.Function.Arguments), &input)
+				content = append(
+					content,
+					&types.ContentBlockMemberToolUse{
+						Value: types.ToolUseBlock{
+							ToolUseId: aws.String(tc.ID),
+							Name:      aws.String(tc.Function.Name),
+							Input:     document.NewLazyDocument(input),
+						},
+					},
+				)
+			}
+			out = append(
+				out, types.Message{
+					Role:    types.ConversationRoleAssistant,
+					Content: content,
+				},
+			)
+
+		case llm.RoleTool:
+			out = append(out, types.Message{
+				Role: types.ConversationRoleUser,
+				Content: []types.ContentBlock{
+					&types.ContentBlockMemberToolResult{
+						Value: types.ToolResultBlock{
+							ToolUseId: aws.String(msg.ToolCallID),
+							Content: []types.ToolResultContentBlock{
+								&types.ToolResultContentBlockMemberText{Value: msg.Text()},
+							},
+						},
+					},
+				},
+			})
+		}
+	}
+
+	return out
+}
+
+func buildToolConfig(req *llm.ChatCompletionRequest) *types.ToolConfiguration {
+	config := &types.ToolConfiguration{}
+
+	tools := make([]types.Tool, len(req.Tools))
+	for i, t := range req.Tools {
+		spec := types.ToolSpecification{
+			Name:        aws.String(t.Name),
+			Description: aws.String(t.Description),
+		}
+		if t.Parameters != nil {
+			var schema any
+			_ = json.Unmarshal(t.Parameters, &schema)
+			spec.InputSchema = &types.ToolInputSchemaMemberJson{
+				Value: document.NewLazyDocument(schema),
+			}
+		}
+		tools[i] = &types.ToolMemberToolSpec{Value: spec}
+	}
+	config.Tools = tools
+
+	if req.ToolChoice != nil {
+		config.ToolChoice = buildToolChoice(req.ToolChoice)
+	}
+
+	return config
+}
+
+func buildToolChoice(tc *llm.ToolChoice) types.ToolChoice {
+	switch tc.Type {
+	case llm.ToolChoiceAuto:
+		return &types.ToolChoiceMemberAuto{Value: types.AutoToolChoice{}}
+	case llm.ToolChoiceRequired:
+		return &types.ToolChoiceMemberAny{Value: types.AnyToolChoice{}}
+	case llm.ToolChoiceFunction:
+		return &types.ToolChoiceMemberTool{
+			Value: types.SpecificToolChoice{
+				Name: aws.String(tc.Function),
+			},
+		}
+	case llm.ToolChoiceNone:
+		// Bedrock doesn't have a "none" tool choice; omit tools instead.
+		return nil
+	default:
+		return nil
+	}
+}
+
+func mapResponse(output *bedrockruntime.ConverseOutput, model string) *llm.ChatCompletionResponse {
+	resp := &llm.ChatCompletionResponse{
+		Model:        model,
+		FinishReason: mapStopReason(output.StopReason),
+		Message: llm.Message{
+			Role: llm.RoleAssistant,
+		},
+	}
+
+	if output.Usage != nil {
+		resp.Usage = llm.Usage{
+			InputTokens:  int(aws.ToInt32(output.Usage.InputTokens)),
+			OutputTokens: int(aws.ToInt32(output.Usage.OutputTokens)),
+		}
+	}
+
+	// Extract message content from the response output union.
+	if msgOutput, ok := output.Output.(*types.ConverseOutputMemberMessage); ok {
+		for _, block := range msgOutput.Value.Content {
+			switch b := block.(type) {
+			case *types.ContentBlockMemberText:
+				resp.Message.Parts = append(resp.Message.Parts, llm.TextPart{Text: b.Value})
+			case *types.ContentBlockMemberToolUse:
+				var args any
+				if b.Value.Input != nil {
+					_ = b.Value.Input.UnmarshalSmithyDocument(&args)
+				}
+				argsJSON, _ := json.Marshal(args)
+				resp.Message.ToolCalls = append(resp.Message.ToolCalls, llm.ToolCall{
+					ID: aws.ToString(b.Value.ToolUseId),
+					Function: llm.FunctionCall{
+						Name:      aws.ToString(b.Value.Name),
+						Arguments: string(argsJSON),
+					},
+				})
+			}
+		}
+	}
+
+	return resp
+}
+
+func mapStopReason(reason types.StopReason) llm.FinishReason {
+	switch reason {
+	case types.StopReasonEndTurn, types.StopReasonStopSequence:
+		return llm.FinishReasonStop
+	case types.StopReasonMaxTokens:
+		return llm.FinishReasonLength
+	case types.StopReasonToolUse:
+		return llm.FinishReasonToolCalls
+	case types.StopReasonContentFiltered, types.StopReasonGuardrailIntervened:
+		return llm.FinishReasonContentFilter
+	default:
+		return llm.FinishReasonStop
+	}
+}
+
+func mapError(err error) error {
+	var respErr *smithyhttp.ResponseError
+	if !errors.As(err, &respErr) {
+		// Check for common error types by message content.
+		msg := err.Error()
+		if strings.Contains(msg, "throttling") || strings.Contains(msg, "ThrottlingException") {
+			return &llm.ErrRateLimit{Err: err}
+		}
+		return err
+	}
+
+	switch respErr.HTTPStatusCode() {
+	case 429:
+		return &llm.ErrRateLimit{Err: err}
+	case 401, 403:
+		return &llm.ErrAuthentication{Err: err}
+	case 400:
+		msg := err.Error()
+		if strings.Contains(msg, "context") || strings.Contains(msg, "token") {
+			return &llm.ErrContextLength{Err: err}
+		}
+		return err
+	default:
+		return err
+	}
+}
+
+// bedrockStream adapts a Bedrock ConverseStream to our ChatCompletionStream interface.
+type bedrockStream struct {
+	eventStream *bedrockruntime.ConverseStreamEventStream
+	events      <-chan types.ConverseStreamOutput
+	current     llm.ChatCompletionStreamEvent
+	err         error
+	toolIndex   int
+}
+
+func newBedrockStream(eventStream *bedrockruntime.ConverseStreamEventStream) *bedrockStream {
+	return &bedrockStream{
+		eventStream: eventStream,
+		events:      eventStream.Events(),
+	}
+}
+
+func (s *bedrockStream) Next() bool {
+	for event := range s.events {
+		mapped, ok := s.mapEvent(event)
+		if ok {
+			s.current = mapped
+			return true
+		}
+	}
+
+	if err := s.eventStream.Err(); err != nil {
+		s.err = mapError(err)
+	}
+	return false
+}
+
+func (s *bedrockStream) Event() llm.ChatCompletionStreamEvent {
+	return s.current
+}
+
+func (s *bedrockStream) Err() error {
+	return s.err
+}
+
+func (s *bedrockStream) Close() error {
+	return s.eventStream.Close()
+}
+
+func (s *bedrockStream) mapEvent(event types.ConverseStreamOutput) (llm.ChatCompletionStreamEvent, bool) {
+	switch e := event.(type) {
+	case *types.ConverseStreamOutputMemberContentBlockStart:
+		if start, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok {
+			return llm.ChatCompletionStreamEvent{
+				Delta: llm.MessageDelta{
+					ToolCalls: []llm.ToolCallDelta{{
+						Index: s.toolIndex,
+						ID:    aws.ToString(start.Value.ToolUseId),
+						Name:  aws.ToString(start.Value.Name),
+					}},
+				},
+			}, true
+		}
+		return llm.ChatCompletionStreamEvent{}, false
+
+	case *types.ConverseStreamOutputMemberContentBlockDelta:
+		switch d := e.Value.Delta.(type) {
+		case *types.ContentBlockDeltaMemberText:
+			return llm.ChatCompletionStreamEvent{
+				Delta: llm.MessageDelta{Content: d.Value},
+			}, true
+		case *types.ContentBlockDeltaMemberToolUse:
+			return llm.ChatCompletionStreamEvent{
+				Delta: llm.MessageDelta{
+					ToolCalls: []llm.ToolCallDelta{{
+						Index:     s.toolIndex,
+						Arguments: aws.ToString(d.Value.Input),
+					}},
+				},
+			}, true
+		}
+		return llm.ChatCompletionStreamEvent{}, false
+
+	case *types.ConverseStreamOutputMemberContentBlockStop:
+		s.toolIndex++
+		return llm.ChatCompletionStreamEvent{}, false
+
+	case *types.ConverseStreamOutputMemberMessageStop:
+		fr := mapStopReason(e.Value.StopReason)
+		return llm.ChatCompletionStreamEvent{
+			FinishReason: &fr,
+		}, true
+
+	case *types.ConverseStreamOutputMemberMetadata:
+		if e.Value.Usage != nil {
+			return llm.ChatCompletionStreamEvent{
+				Usage: &llm.Usage{
+					InputTokens:  int(aws.ToInt32(e.Value.Usage.InputTokens)),
+					OutputTokens: int(aws.ToInt32(e.Value.Usage.OutputTokens)),
+				},
+			}, true
+		}
+		return llm.ChatCompletionStreamEvent{}, false
+
+	default:
+		return llm.ChatCompletionStreamEvent{}, false
+	}
+}