diff --git a/chains/llm_test.go b/chains/llm_test.go index da0922b8b..fd09940e2 100644 --- a/chains/llm_test.go +++ b/chains/llm_test.go @@ -12,6 +12,7 @@ import ( "github.com/tmc/langchaingo/callbacks" "github.com/tmc/langchaingo/httputil" "github.com/tmc/langchaingo/internal/httprr" + "github.com/tmc/langchaingo/llms" "github.com/tmc/langchaingo/llms/googleai" "github.com/tmc/langchaingo/llms/openai" "github.com/tmc/langchaingo/prompts" @@ -65,6 +66,31 @@ func TestLLMChain(t *testing.T) { require.True(t, strings.Contains(result, "Paris")) } +type errorLanguageModel struct { + err error +} + +func (m *errorLanguageModel) Call(_ context.Context, _ string, _ ...llms.CallOption) (string, error) { + return "", m.err +} + +func (m *errorLanguageModel) GenerateContent(_ context.Context, _ []llms.MessageContent, _ ...llms.CallOption) (*llms.ContentResponse, error) { + return nil, m.err +} + +func TestLLMChainPropagatesContentFilterError(t *testing.T) { + t.Parallel() + + chain := NewLLMChain( + &errorLanguageModel{err: llms.NewError(llms.ErrCodeContentFilter, "bedrock", "blocked")}, + prompts.NewPromptTemplate("{{.text}}", []string{"text"}), + ) + + _, err := chain.Call(context.Background(), map[string]any{"text": "unsafe prompt"}) + require.Error(t, err) + require.True(t, llms.IsContentFilterError(err)) +} + func TestLLMChainWithChatPromptTemplate(t *testing.T) { ctx := context.Background() t.Parallel() diff --git a/chains/options.go b/chains/options.go index c6df23b09..5865f4dfa 100644 --- a/chains/options.go +++ b/chains/options.go @@ -63,6 +63,10 @@ type chainCallOption struct { RepetitionPenalty float64 repetitionPenaltySet bool + // Safety configures provider-defined safety controls in an LLM call. + SafetyConfig map[string]any + safetyConfigSet bool + // CallbackHandler is the callback handler for Chain CallbackHandler callbacks.Handler } @@ -146,6 +150,16 @@ func WithRepetitionPenalty(repetitionPenalty float64) ChainCallOption { } } +// WithSafetyConfig configures provider-defined safety controls for the LLM call. +func WithSafetyConfig(config map[string]any) ChainCallOption { + return func(o *chainCallOption) { + if config != nil { + o.SafetyConfig = config + o.safetyConfigSet = true + } + } +} + // WithStopWords is an option for setting the stop words for LLM.Call. func WithStopWords(stopWords []string) ChainCallOption { return func(o *chainCallOption) { @@ -208,6 +222,9 @@ func GetLLMCallOptions(options ...ChainCallOption) []llms.CallOption { //nolint: if opts.repetitionPenaltySet { chainCallOption = append(chainCallOption, llms.WithRepetitionPenalty(opts.RepetitionPenalty)) } + if opts.safetyConfigSet { + chainCallOption = append(chainCallOption, llms.WithSafetyConfig(opts.SafetyConfig)) + } chainCallOption = append(chainCallOption, llms.WithStreamingFunc(opts.StreamingFunc)) return chainCallOption diff --git a/llms/bedrock/bedrockllm.go b/llms/bedrock/bedrockllm.go index 12363a090..4ee109a9b 100644 --- a/llms/bedrock/bedrockllm.go +++ b/llms/bedrock/bedrockllm.go @@ -103,37 +103,45 @@ func processMessages(messages []llms.MessageContent) ([]bedrockclient.Message, e for _, m := range messages { for _, part := range m.Parts { - switch part := part.(type) { + var cacheControl *llms.CacheControl + if cached, ok := part.(llms.CachedContent); ok { + cacheControl = cached.CacheControl + part = cached.ContentPart + } + + switch p := part.(type) { case llms.TextContent: bedrockMsgs = append(bedrockMsgs, bedrockclient.Message{ - Role: m.Role, - Content: part.Text, - Type: "text", + Role: m.Role, + Content: p.Text, + Type: "text", + CacheControl: cacheControl, }) case llms.BinaryContent: bedrockMsgs = append(bedrockMsgs, bedrockclient.Message{ - Role: m.Role, - Content: string(part.Data), - MimeType: part.MIMEType, - Type: "image", + Role: m.Role, + Content: string(p.Data), + MimeType: p.MIMEType, + Type: "image", + CacheControl: cacheControl, }) case llms.ToolCall: - // Handle tool calls from AI messages bedrockMsgs = append(bedrockMsgs, bedrockclient.Message{ - Role: m.Role, - Content: "", // Content will be empty for tool calls - Type: "tool_call", - ToolCallID: part.ID, - ToolName: part.FunctionCall.Name, - ToolArgs: part.FunctionCall.Arguments, + Role: m.Role, + Content: "", // Content will be empty for tool calls + Type: "tool_call", + ToolCallID: p.ID, + ToolName: p.FunctionCall.Name, + ToolArgs: p.FunctionCall.Arguments, + CacheControl: cacheControl, }) case llms.ToolCallResponse: - // Handle tool result messages bedrockMsgs = append(bedrockMsgs, bedrockclient.Message{ - Role: m.Role, - Content: part.Content, - Type: "tool_result", - ToolUseID: part.ToolCallID, + Role: m.Role, + Content: p.Content, + Type: "tool_result", + ToolUseID: p.ToolCallID, + CacheControl: cacheControl, }) default: return nil, errors.New("unsupported message type") diff --git a/llms/bedrock/caching.go b/llms/bedrock/caching.go new file mode 100644 index 000000000..e2558a801 --- /dev/null +++ b/llms/bedrock/caching.go @@ -0,0 +1,19 @@ +package bedrock + +import ( + "time" + + "github.com/tmc/langchaingo/llms" +) + +// EphemeralCache returns an ephemeral cache control for Bedrock prompt +// caching with no ttl field on the wire. Bedrock applies its default +// 5-minute behavior. +func EphemeralCache() *llms.CacheControl { + return &llms.CacheControl{Type: "ephemeral"} +} + +// EphemeralCacheOneHour returns an ephemeral cache control with ttl=1h. +func EphemeralCacheOneHour() *llms.CacheControl { + return &llms.CacheControl{Type: "ephemeral", Duration: time.Hour} +} diff --git a/llms/bedrock/internal/bedrockclient/bedrockclient.go b/llms/bedrock/internal/bedrockclient/bedrockclient.go index fcb91a1dc..d8309be61 100644 --- a/llms/bedrock/internal/bedrockclient/bedrockclient.go +++ b/llms/bedrock/internal/bedrockclient/bedrockclient.go @@ -32,6 +32,9 @@ type Message struct { ToolArgs string `json:"tool_args,omitempty"` // Tool result fields ToolUseID string `json:"tool_use_id,omitempty"` + // CacheControl marks this message as a prompt-cache breakpoint when set. + // Providers that don't support prompt caching silently ignore the field. + CacheControl *llms.CacheControl `json:"-"` } func getProvider(modelID string) string { diff --git a/llms/bedrock/internal/bedrockclient/bedrockclient_integration_test.go b/llms/bedrock/internal/bedrockclient/bedrockclient_integration_test.go index 084fdacb5..be60a1c6d 100644 --- a/llms/bedrock/internal/bedrockclient/bedrockclient_integration_test.go +++ b/llms/bedrock/internal/bedrockclient/bedrockclient_integration_test.go @@ -174,8 +174,10 @@ func TestClient_CreateCompletion(t *testing.T) { }, StopReason: AnthropicCompletionReasonEndTurn, Usage: struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ InputTokens: 10, OutputTokens: 5, @@ -429,16 +431,20 @@ func TestClient_CreateCompletion_Streaming(t *testing.T) { StopReason any `json:"stop_reason"` StopSequence any `json:"stop_sequence"` Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` } `json:"usage"` }{ ID: "msg-123", Type: "message", Role: "assistant", Usage: struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ InputTokens: 10, }, @@ -485,7 +491,9 @@ func TestClient_CreateCompletion_Streaming(t *testing.T) { StopReason: AnthropicCompletionReasonEndTurn, }, Usage: struct { - OutputTokens int `json:"output_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ OutputTokens: 15, }, diff --git a/llms/bedrock/internal/bedrockclient/bedrockclient_test.go b/llms/bedrock/internal/bedrockclient/bedrockclient_test.go index 8f036db25..eeed7a600 100644 --- a/llms/bedrock/internal/bedrockclient/bedrockclient_test.go +++ b/llms/bedrock/internal/bedrockclient/bedrockclient_test.go @@ -258,7 +258,7 @@ func TestProcessInputMessagesAnthropic(t *testing.T) { name string messages []Message expectedMsgs int - expectedSystem string + expectedSystem interface{} expectError bool errorContains string }{ @@ -268,7 +268,7 @@ func TestProcessInputMessagesAnthropic(t *testing.T) { {Role: llms.ChatMessageTypeHuman, Type: "text", Content: "Hello"}, }, expectedMsgs: 1, - expectedSystem: "", + expectedSystem: nil, }, { name: "system message extracted", @@ -304,7 +304,7 @@ func TestProcessInputMessagesAnthropic(t *testing.T) { {Role: llms.ChatMessageTypeHuman, Type: "text", Content: "How are you?"}, }, expectedMsgs: 3, - expectedSystem: "", + expectedSystem: nil, }, { name: "multiple messages same role chunked together", @@ -314,7 +314,7 @@ func TestProcessInputMessagesAnthropic(t *testing.T) { {Role: llms.ChatMessageTypeAI, Type: "text", Content: "Response"}, }, expectedMsgs: 2, - expectedSystem: "", + expectedSystem: nil, }, { name: "function role converted to user", @@ -322,7 +322,7 @@ func TestProcessInputMessagesAnthropic(t *testing.T) { {Role: llms.ChatMessageTypeFunction, Type: "text", Content: "Function call"}, }, expectedMsgs: 1, - expectedSystem: "", + expectedSystem: nil, }, } @@ -673,8 +673,10 @@ func TestAnthropicResponseParsing(t *testing.T) { StopReason: AnthropicCompletionReasonEndTurn, StopSequence: "", Usage: struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ InputTokens: 10, OutputTokens: 15, @@ -746,8 +748,10 @@ func TestAnthropicStreamingResponseChunk(t *testing.T) { StopReason any `json:"stop_reason"` StopSequence any `json:"stop_sequence"` Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` } `json:"usage"` }{ ID: "msg-123", @@ -755,8 +759,10 @@ func TestAnthropicStreamingResponseChunk(t *testing.T) { Role: "assistant", Model: "claude-3", Usage: struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ InputTokens: 25, }, @@ -792,7 +798,9 @@ func TestAnthropicStreamingResponseChunk(t *testing.T) { StopReason: AnthropicCompletionReasonEndTurn, }, Usage: struct { - OutputTokens int `json:"output_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ OutputTokens: 12, }, diff --git a/llms/bedrock/internal/bedrockclient/guardrail.go b/llms/bedrock/internal/bedrockclient/guardrail.go new file mode 100644 index 000000000..45c3e3f57 --- /dev/null +++ b/llms/bedrock/internal/bedrockclient/guardrail.go @@ -0,0 +1,147 @@ +package bedrockclient + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/tmc/langchaingo/llms" +) + +const bedrockGuardrailIntervened = "INTERVENED" + +type bedrockGuardrailParams struct { + identifier *string + version *string + tagSuffix *string + streamProcessingMode *string +} + +func getRequiredGuardrailParam(safety map[string]any, key string) (*string, error) { + value, ok := safety[key] + if !ok { + return nil, fmt.Errorf("bedrock guardrail %s is required when safety config is provided", key) + } + + str, ok := value.(string) + if !ok || strings.TrimSpace(str) == "" { + return nil, fmt.Errorf("bedrock guardrail %s must be a non-empty string", key) + } + + return aws.String(str), nil +} + +func getOptionalGuardrailParam(safety map[string]any, key string) (*string, error) { + value, ok := safety[key] + if !ok { + return nil, nil + } + + str, ok := value.(string) + if !ok || strings.TrimSpace(str) == "" { + return nil, fmt.Errorf("bedrock guardrail %s must be a non-empty string if provided", key) + } + + return aws.String(str), nil +} + +func getGuardrailParams(options llms.CallOptions) (bedrockGuardrailParams, error) { + if options.SafetyConfig == nil { + return bedrockGuardrailParams{}, nil + } + + identifier, err := getRequiredGuardrailParam(options.SafetyConfig, "identifier") + if err != nil { + return bedrockGuardrailParams{}, err + } + version, err := getRequiredGuardrailParam(options.SafetyConfig, "version") + if err != nil { + return bedrockGuardrailParams{}, err + } + + tagSuffix, err := getOptionalGuardrailParam(options.SafetyConfig, "tagSuffix") + if err != nil { + return bedrockGuardrailParams{}, err + } + streamProcessingMode, err := getOptionalGuardrailParam(options.SafetyConfig, "streamProcessingMode") + if err != nil { + return bedrockGuardrailParams{}, err + } + + return bedrockGuardrailParams{ + identifier: identifier, + version: version, + tagSuffix: tagSuffix, + streamProcessingMode: streamProcessingMode, + }, nil +} + +func handleGuardrailParams(input guardrailConfigurableAdapter, options llms.CallOptions) error { + guardrailParams, err := getGuardrailParams(options) + if err != nil { + return fmt.Errorf("failed to get guardrail parameters: %w", err) + } + + if guardrailParams.identifier != nil && guardrailParams.version != nil { + input.setGuardrailIdentifier(guardrailParams.identifier) + input.setGuardrailVersion(guardrailParams.version) + } + + if guardrailParams.tagSuffix != nil || guardrailParams.streamProcessingMode != nil { + var m map[string]any + if err := json.Unmarshal(input.body(), &m); err != nil { + return fmt.Errorf("failed to unmarshal input body when adding guardrail parameters: %w", err) + } + + var config = make(map[string]string) + if guardrailParams.tagSuffix != nil { + config["tagSuffix"] = *guardrailParams.tagSuffix + } + if guardrailParams.streamProcessingMode != nil { + config["streamProcessingMode"] = *guardrailParams.streamProcessingMode + } + m["amazon-bedrock-guardrailConfig"] = config + + newBody, err := json.Marshal(m) + if err != nil { + return fmt.Errorf("failed to marshal input body when adding guardrail parameters: %w", err) + } + input.setBody(newBody) + } + return nil +} + +func checkGuardrailError(response any) error { + var action string + var trace json.RawMessage + + switch r := response.(type) { + case []byte: + var resp guardrailResponse + if err := json.Unmarshal(r, &resp); err != nil { + // If we can't parse the response, we assume it's not a guardrail error and return nil + return nil + } + action = resp.AmazonBedrockGuardrailAction + trace = resp.AmazonBedrockTrace + case streamingCompletionResponseChunk: + action = r.AmazonBedrockGuardrailAction + trace = r.AmazonBedrockTrace + case *streamingCompletionResponseChunk: + action = r.AmazonBedrockGuardrailAction + trace = r.AmazonBedrockTrace + default: + return nil + } + + if action == bedrockGuardrailIntervened { + err := llms.NewError(llms.ErrCodeContentFilter, "bedrock", "content blocked by Bedrock guardrail") + err.WithDetail("guardrail_action", action) + if len(trace) > 0 { + err.WithDetail("trace", string(trace)) + } + return err + } + return nil +} diff --git a/llms/bedrock/internal/bedrockclient/invoke.go b/llms/bedrock/internal/bedrockclient/invoke.go new file mode 100644 index 000000000..79341462a --- /dev/null +++ b/llms/bedrock/internal/bedrockclient/invoke.go @@ -0,0 +1,158 @@ +package bedrockclient + +import ( + "context" + "encoding/json" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + "github.com/tmc/langchaingo/llms" +) + +const ( + bedrockAcceptAll = "*/*" + bedrockJSONContentType = "application/json" +) + +type guardrailResponse struct { + AmazonBedrockGuardrailAction string `json:"amazon-bedrock-guardrailAction"` + AmazonBedrockTrace json.RawMessage `json:"amazon-bedrock-trace"` +} + +type streamingCompletionResponseChunk struct { + Type string `json:"type"` + Index int `json:"index"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + StopReason string `json:"stop_reason"` + StopSequence any `json:"stop_sequence"` + } `json:"delta"` + AmazonBedrockInvocationMetrics struct { + InputTokenCount int `json:"inputTokenCount"` + OutputTokenCount int `json:"outputTokenCount"` + InvocationLatency int `json:"invocationLatency"` + FirstByteLatency int `json:"firstByteLatency"` + } `json:"amazon-bedrock-invocationMetrics"` + Usage struct { + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + } `json:"usage"` + Message struct { + ID string `json:"id"` + Type string `json:"type"` + Role string `json:"role"` + Content []any `json:"content"` + Model string `json:"model"` + StopReason any `json:"stop_reason"` + StopSequence any `json:"stop_sequence"` + Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` + } `json:"usage"` + } `json:"message"` + guardrailResponse +} + +type invokeModelAPI interface { + InvokeModel(ctx context.Context, params *bedrockruntime.InvokeModelInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) +} + +type invokeModelWithResponseStreamAPI interface { + InvokeModelWithResponseStream(ctx context.Context, params *bedrockruntime.InvokeModelWithResponseStreamInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelWithResponseStreamOutput, error) +} + +type guardrailConfigurableAdapter interface { + body() []byte + setGuardrailIdentifier(*string) + setGuardrailVersion(*string) + setBody([]byte) +} + +type guardrailConfigurableInput struct { + *bedrockruntime.InvokeModelInput +} + +var _ guardrailConfigurableAdapter = (*guardrailConfigurableInput)(nil) + +func (i *guardrailConfigurableInput) setGuardrailIdentifier(identifier *string) { + i.GuardrailIdentifier = identifier +} +func (i *guardrailConfigurableInput) setGuardrailVersion(version *string) { + i.GuardrailVersion = version +} +func (i *guardrailConfigurableInput) body() []byte { return i.Body } +func (i *guardrailConfigurableInput) setBody(body []byte) { i.Body = body } + +type guardrailConfigurableStreamInput struct { + *bedrockruntime.InvokeModelWithResponseStreamInput +} + +var _ guardrailConfigurableAdapter = (*guardrailConfigurableStreamInput)(nil) + +func (i *guardrailConfigurableStreamInput) setGuardrailIdentifier(identifier *string) { + i.GuardrailIdentifier = identifier +} +func (i *guardrailConfigurableStreamInput) setGuardrailVersion(version *string) { + i.GuardrailVersion = version +} +func (i *guardrailConfigurableStreamInput) body() []byte { return i.Body } +func (i *guardrailConfigurableStreamInput) setBody(body []byte) { i.Body = body } + +func newInvokeModelInput(modelID string, body []byte, options llms.CallOptions) (*bedrockruntime.InvokeModelInput, error) { + input := &guardrailConfigurableInput{ + &bedrockruntime.InvokeModelInput{ + ModelId: aws.String(modelID), + Accept: aws.String(bedrockAcceptAll), + ContentType: aws.String(bedrockJSONContentType), + Body: body, + }, + } + if err := handleGuardrailParams(input, options); err != nil { + return nil, err + } + return input.InvokeModelInput, nil +} + +func newInvokeModelWithResponseStreamInput(modelID string, body []byte, options llms.CallOptions) (*bedrockruntime.InvokeModelWithResponseStreamInput, error) { + input := &guardrailConfigurableStreamInput{ + &bedrockruntime.InvokeModelWithResponseStreamInput{ + ModelId: aws.String(modelID), + Accept: aws.String(bedrockAcceptAll), + ContentType: aws.String(bedrockJSONContentType), + Body: body, + }, + } + if err := handleGuardrailParams(input, options); err != nil { + return nil, err + } + return input.InvokeModelWithResponseStreamInput, nil +} + +func invokeModel(ctx context.Context, client invokeModelAPI, modelID string, body []byte, options llms.CallOptions) (*bedrockruntime.InvokeModelOutput, error) { + input, err := newInvokeModelInput(modelID, body, options) + if err != nil { + return nil, err + } + + output, err := client.InvokeModel(ctx, input) + if err != nil { + return nil, err + } + if err := checkGuardrailError(output.Body); err != nil { + return nil, err + } + + return output, nil +} + +func invokeModelWithResponseStream(ctx context.Context, client invokeModelWithResponseStreamAPI, modelID string, body []byte, options llms.CallOptions) (*bedrockruntime.InvokeModelWithResponseStreamOutput, error) { + input, err := newInvokeModelWithResponseStreamInput(modelID, body, options) + if err != nil { + return nil, err + } + return client.InvokeModelWithResponseStream(ctx, input) +} diff --git a/llms/bedrock/internal/bedrockclient/invoke_test.go b/llms/bedrock/internal/bedrockclient/invoke_test.go new file mode 100644 index 000000000..3f6471dc4 --- /dev/null +++ b/llms/bedrock/internal/bedrockclient/invoke_test.go @@ -0,0 +1,123 @@ +package bedrockclient + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tmc/langchaingo/llms" +) + +func TestInvokeModelInputs(t *testing.T) { + t.Run("without safety config", func(t *testing.T) { + input, err := newInvokeModelInput("anthropic.claude-v2", []byte(`{"prompt":"hello"}`), llms.CallOptions{}) + require.NoError(t, err) + require.NotNil(t, input) + assert.Nil(t, input.GuardrailIdentifier) + assert.Nil(t, input.GuardrailVersion) + + streamInput, err := newInvokeModelWithResponseStreamInput("anthropic.claude-v2", []byte(`{"prompt":"hello"}`), llms.CallOptions{}) + require.NoError(t, err) + require.NotNil(t, streamInput) + assert.Nil(t, streamInput.GuardrailIdentifier) + assert.Nil(t, streamInput.GuardrailVersion) + }) + + t.Run("with safety config", func(t *testing.T) { + var options llms.CallOptions + llms.WithSafetyConfig(map[string]any{ + "identifier": "gr-123", + "version": "1", + })(&options) + + input, err := newInvokeModelInput("anthropic.claude-v2", []byte(`{"prompt":"hello"}`), options) + require.NoError(t, err) + require.NotNil(t, input.GuardrailIdentifier) + require.NotNil(t, input.GuardrailVersion) + assert.Equal(t, "gr-123", *input.GuardrailIdentifier) + assert.Equal(t, "1", *input.GuardrailVersion) + + streamInput, err := newInvokeModelWithResponseStreamInput("anthropic.claude-v2", []byte(`{"prompt":"hello"}`), options) + require.NoError(t, err) + require.NotNil(t, streamInput.GuardrailIdentifier) + require.NotNil(t, streamInput.GuardrailVersion) + assert.Equal(t, "gr-123", *streamInput.GuardrailIdentifier) + assert.Equal(t, "1", *streamInput.GuardrailVersion) + }) + + t.Run("requires safety version", func(t *testing.T) { + _, err := newInvokeModelInput("anthropic.claude-v2", nil, llms.CallOptions{ + SafetyConfig: map[string]any{"identifier": "gr-123"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "bedrock safety version") + }) + + t.Run("requires safety identifier", func(t *testing.T) { + _, err := newInvokeModelInput("anthropic.claude-v2", nil, llms.CallOptions{ + SafetyConfig: map[string]any{"version": "1"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "bedrock safety identifier") + }) + + t.Run("rejects invalid safety types", func(t *testing.T) { + _, err := newInvokeModelInput("anthropic.claude-v2", nil, llms.CallOptions{ + SafetyConfig: map[string]any{ + "identifier": 123, + "version": "1", + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "bedrock safety identifier") + }) + + t.Run("returns content filter error when guardrail intervenes", func(t *testing.T) { + resp, err := invokeModel(context.Background(), &mockBedrockClient{ + invokeFunc: func(ctx context.Context, params *bedrockruntime.InvokeModelInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + return &bedrockruntime.InvokeModelOutput{ + Body: []byte(`{"amazon-bedrock-guardrailAction":"INTERVENED","amazon-bedrock-trace":{"rule":"blocked"}}`), + }, nil + }, + }, "anthropic.claude-v2", []byte(`{"prompt":"hello"}`), llms.CallOptions{}) + + require.Nil(t, resp) + require.Error(t, err) + assert.True(t, llms.IsContentFilterError(err)) + + var llmErr *llms.Error + require.ErrorAs(t, err, &llmErr) + assert.Equal(t, "bedrock", llmErr.Provider) + assert.Equal(t, "INTERVENED", llmErr.Details["guardrail_action"]) + assert.Equal(t, map[string]any{"rule": "blocked"}, llmErr.Details["trace"]) + }) + + t.Run("returns content filter error for provider filtered response", func(t *testing.T) { + resp, err := invokeModel(context.Background(), &mockBedrockClient{ + invokeFunc: func(ctx context.Context, params *bedrockruntime.InvokeModelInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + return &bedrockruntime.InvokeModelOutput{ + Body: []byte(`{"results":[{"completionReason":"CONTENT_FILTERED"}]}`), + }, nil + }, + }, "amazon.titan-text-express-v1", []byte(`{"prompt":"hello"}`), llms.CallOptions{}) + + require.Nil(t, resp) + require.Error(t, err) + assert.True(t, llms.IsContentFilterError(err)) + }) + + t.Run("returns response when guardrail does not intervene", func(t *testing.T) { + resp, err := invokeModel(context.Background(), &mockBedrockClient{ + invokeFunc: func(ctx context.Context, params *bedrockruntime.InvokeModelInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + return &bedrockruntime.InvokeModelOutput{ + Body: []byte(`{"amazon-bedrock-guardrailAction":"NONE","outputText":"ok"}`), + }, nil + }, + }, "anthropic.claude-v2", []byte(`{"prompt":"hello"}`), llms.CallOptions{}) + + require.NoError(t, err) + require.NotNil(t, resp) + }) +} diff --git a/llms/bedrock/internal/bedrockclient/provider_ai21.go b/llms/bedrock/internal/bedrockclient/provider_ai21.go index 5ff482b15..a9aef212d 100644 --- a/llms/bedrock/internal/bedrockclient/provider_ai21.go +++ b/llms/bedrock/internal/bedrockclient/provider_ai21.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/tmc/langchaingo/llms" ) @@ -100,14 +99,7 @@ func createAi21Completion(ctx context.Context, client *bedrockruntime.Client, mo return nil, err } - modelInput := bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Body: body, - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - } - - resp, err := client.InvokeModel(ctx, &modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } diff --git a/llms/bedrock/internal/bedrockclient/provider_amazon.go b/llms/bedrock/internal/bedrockclient/provider_amazon.go index bae6eb825..36200c16f 100644 --- a/llms/bedrock/internal/bedrockclient/provider_amazon.go +++ b/llms/bedrock/internal/bedrockclient/provider_amazon.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/tmc/langchaingo/llms" ) @@ -79,13 +78,7 @@ func createAmazonCompletion(ctx context.Context, return nil, err } - modelInput := &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - resp, err := client.InvokeModel(ctx, modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } diff --git a/llms/bedrock/internal/bedrockclient/provider_anthropic.go b/llms/bedrock/internal/bedrockclient/provider_anthropic.go index c6e389808..7369ad1b1 100644 --- a/llms/bedrock/internal/bedrockclient/provider_anthropic.go +++ b/llms/bedrock/internal/bedrockclient/provider_anthropic.go @@ -8,7 +8,6 @@ import ( "errors" "fmt" - "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/types" "github.com/tmc/langchaingo/llms" @@ -46,6 +45,25 @@ type anthropicTextGenerationInputContent struct { ID string `json:"id,omitempty"` Name string `json:"name,omitempty"` Input interface{} `json:"input,omitempty"` + // CacheControl marks this content block as a prompt-cache breakpoint. + // The cached prefix extends from the start of the prompt up to and + // including this block. + CacheControl *anthropicCacheControl `json:"cache_control,omitempty"` +} + +// anthropicCacheControl maps llms.CacheControl onto Bedrock's wire format for +// Anthropic models. +type anthropicCacheControl struct { + Type string `json:"type"` + TTL string `json:"ttl,omitempty"` +} + +// anthropicSystemBlock is used when the system field is emitted as an array of +// blocks (only when at least one system part carries cache control). +type anthropicSystemBlock struct { + Type string `json:"type"` + Text string `json:"text"` + CacheControl *anthropicCacheControl `json:"cache_control,omitempty"` } type anthropicTextGenerationInputMessage struct { @@ -63,8 +81,10 @@ type anthropicTextGenerationInput struct { AnthropicVersion string `json:"anthropic_version"` // The maximum number of tokens to generate per result. Required MaxTokens int `json:"max_tokens"` - // The system prompt to use. Optional - System string `json:"system,omitempty"` + // The system prompt to use. Optional. + // Either a string (legacy form, no caching) or []anthropicSystemBlock + // when at least one system part carries cache control. + System interface{} `json:"system,omitempty"` // The messages to use. Required Messages []*anthropicTextGenerationInputMessage `json:"messages"` // The amount of randomness injected into the response. Optional, default = 1 @@ -81,6 +101,11 @@ type anthropicTextGenerationInput struct { Tools []BedrockTool `json:"tools,omitempty"` // Tool choice configuration. Optional ToolChoice *BedrockToolChoice `json:"tool_choice,omitempty"` + // Guardrail configuration. Optional + GuardrailConfig *struct { + TagSuffix string `json:"tagSuffix,omitempty"` + StreamProcessingMode string `json:"streamProcessingMode,omitempty"` + } `json:"amazon-bedrock-guardrailConfig,omitempty"` } // anthropicTextGenerationOutput is the generated output. @@ -100,8 +125,10 @@ type anthropicTextGenerationOutput struct { // Which custom stop sequence was matched, if any. StopSequence string `json:"stop_sequence"` Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` } `json:"usage"` } @@ -182,28 +209,38 @@ func createAnthropicCompletion(ctx context.Context, } } + // Add guardrail tag suffix and stream processing mode if provided in safety config + if options.SafetyConfig != nil { + input.GuardrailConfig = &struct { + TagSuffix string `json:"tagSuffix,omitempty"` + StreamProcessingMode string `json:"streamProcessingMode,omitempty"` + }{} + if tagSuffix, ok := options.SafetyConfig["tag_suffix"]; ok { + if str, ok := tagSuffix.(string); ok { + input.GuardrailConfig.TagSuffix = str + } else { + return nil, errors.New("tag_suffix in safety config must be a string") + } + } + if streamProcessingMode, ok := options.SafetyConfig["stream_processing_mode"]; ok { + if str, ok := streamProcessingMode.(string); ok { + input.GuardrailConfig.StreamProcessingMode = str + } else { + return nil, errors.New("stream_processing_mode in safety config must be a string") + } + } + } + body, err := json.Marshal(input) if err != nil { return nil, err } if options.StreamingFunc != nil { - modelInput := &bedrockruntime.InvokeModelWithResponseStreamInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - return parseStreamingCompletionResponse(ctx, client, modelInput, options) + return parseStreamingCompletionResponse(ctx, client, modelID, body, options) } - modelInput := &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - resp, err := client.InvokeModel(ctx, modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } @@ -224,8 +261,10 @@ func createAnthropicCompletion(ctx context.Context, choice := &llms.ContentChoice{ StopReason: output.StopReason, GenerationInfo: map[string]interface{}{ - "input_tokens": output.Usage.InputTokens, - "output_tokens": output.Usage.OutputTokens, + "input_tokens": output.Usage.InputTokens, + "output_tokens": output.Usage.OutputTokens, + "CacheCreationInputTokens": output.Usage.CacheCreationInputTokens, + "CacheReadInputTokens": output.Usage.CacheReadInputTokens, }, } @@ -263,41 +302,8 @@ func createAnthropicCompletion(ctx context.Context, }, nil } -type streamingCompletionResponseChunk struct { - Type string `json:"type"` - Index int `json:"index"` - Delta struct { - Type string `json:"type"` - Text string `json:"text"` - StopReason string `json:"stop_reason"` - StopSequence any `json:"stop_sequence"` - } `json:"delta"` - AmazonBedrockInvocationMetrics struct { - InputTokenCount int `json:"inputTokenCount"` - OutputTokenCount int `json:"outputTokenCount"` - InvocationLatency int `json:"invocationLatency"` - FirstByteLatency int `json:"firstByteLatency"` - } `json:"amazon-bedrock-invocationMetrics"` - Usage struct { - OutputTokens int `json:"output_tokens"` - } `json:"usage"` - Message struct { - ID string `json:"id"` - Type string `json:"type"` - Role string `json:"role"` - Content []any `json:"content"` - Model string `json:"model"` - StopReason any `json:"stop_reason"` - StopSequence any `json:"stop_sequence"` - Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - } `json:"usage"` - } `json:"message"` -} - -func parseStreamingCompletionResponse(ctx context.Context, client *bedrockruntime.Client, modelInput *bedrockruntime.InvokeModelWithResponseStreamInput, options llms.CallOptions) (*llms.ContentResponse, error) { - output, err := client.InvokeModelWithResponseStream(ctx, modelInput) +func parseStreamingCompletionResponse(ctx context.Context, client invokeModelWithResponseStreamAPI, modelID string, body []byte, options llms.CallOptions) (*llms.ContentResponse, error) { + output, err := invokeModelWithResponseStream(ctx, client, modelID, body, options) if err != nil { return nil, err } @@ -315,14 +321,18 @@ func parseStreamingCompletionResponse(ctx context.Context, client *bedrockruntim if v, ok := e.(*types.ResponseStreamMemberChunk); ok { var resp streamingCompletionResponseChunk - err := json.NewDecoder(bytes.NewReader(v.Value.Bytes)).Decode(&resp) - if err != nil { + if err := json.NewDecoder(bytes.NewReader(v.Value.Bytes)).Decode(&resp); err != nil { + return nil, err + } + if err := checkGuardrailError(&resp); err != nil { return nil, err } switch resp.Type { case "message_start": contentchoices[0].GenerationInfo["input_tokens"] = resp.Message.Usage.InputTokens + contentchoices[0].GenerationInfo["CacheCreationInputTokens"] = resp.Message.Usage.CacheCreationInputTokens + contentchoices[0].GenerationInfo["CacheReadInputTokens"] = resp.Message.Usage.CacheReadInputTokens case "content_block_delta": if err = options.StreamingFunc(ctx, []byte(resp.Delta.Text)); err != nil { return nil, err @@ -331,6 +341,12 @@ func parseStreamingCompletionResponse(ctx context.Context, client *bedrockruntim case "message_delta": contentchoices[0].StopReason = resp.Delta.StopReason contentchoices[0].GenerationInfo["output_tokens"] = resp.Usage.OutputTokens + if resp.Usage.CacheCreationInputTokens > 0 { + contentchoices[0].GenerationInfo["CacheCreationInputTokens"] = resp.Usage.CacheCreationInputTokens + } + if resp.Usage.CacheReadInputTokens > 0 { + contentchoices[0].GenerationInfo["CacheReadInputTokens"] = resp.Usage.CacheReadInputTokens + } } } } @@ -343,9 +359,14 @@ func parseStreamingCompletionResponse(ctx context.Context, client *bedrockruntim }, nil } -// process the input messages to anthropic supported input -// returns the input content and system prompt. -func processInputMessagesAnthropic(messages []Message) ([]*anthropicTextGenerationInputMessage, string, error) { +// converts the flattened Bedrock messages into +// Anthropic's request shape. The returned system value is one of: +// - nil: no system messages were provided. +// - string: concatenated system text (legacy form, backward-compatible when +// no system part carries cache control). +// - []anthropicSystemBlock: emitted only when at least one system part has +// a cache control, so that per-block cache_control markers are preserved. +func processInputMessagesAnthropic(messages []Message) ([]*anthropicTextGenerationInputMessage, interface{}, error) { chunkedMessages := make([][]Message, 0, len(messages)) currentChunk := make([]Message, 0, len(messages)) var lastRole llms.ChatMessageType @@ -364,23 +385,17 @@ func processInputMessagesAnthropic(messages []Message) ([]*anthropicTextGenerati } inputContents := make([]*anthropicTextGenerationInputMessage, 0, len(messages)) - var systemPrompt string + var systemParts []Message for _, chunk := range chunkedMessages { role, err := getAnthropicRole(chunk[0].Role) if err != nil { - return nil, "", err + return nil, nil, err } if role == AnthropicSystem { - if systemPrompt != "" { - return nil, "", errors.New("multiple system prompts") - } - for _, message := range chunk { - c := getAnthropicInputContent(message) - if c.Type != AnthropicMessageTypeText { - return nil, "", errors.New("system prompt must be text") - } - systemPrompt += c.Text + if len(systemParts) > 0 { + return nil, nil, errors.New("multiple system prompts") } + systemParts = append(systemParts, chunk...) continue } content := make([]anthropicTextGenerationInputContent, 0, len(chunk)) @@ -392,7 +407,66 @@ func processInputMessagesAnthropic(messages []Message) ([]*anthropicTextGenerati Content: content, }) } - return inputContents, systemPrompt, nil + + system, err := buildAnthropicSystem(systemParts) + if err != nil { + return nil, nil, err + } + return inputContents, system, nil +} + +// buildAnthropicSystem emits the system field in array form when any part +// carries cache control, otherwise concatenates into a single string for +// backward compatibility. +func buildAnthropicSystem(parts []Message) (interface{}, error) { + if len(parts) == 0 { + return nil, nil + } + for _, m := range parts { + if m.Type != AnthropicMessageTypeText { + return nil, errors.New("system prompt must be text") + } + } + hasCache := false + for _, m := range parts { + if m.CacheControl != nil { + hasCache = true + break + } + } + if !hasCache { + var s string + for _, m := range parts { + s += m.Content + } + return s, nil + } + blocks := make([]anthropicSystemBlock, 0, len(parts)) + for _, m := range parts { + blocks = append(blocks, anthropicSystemBlock{ + Type: AnthropicMessageTypeText, + Text: m.Content, + CacheControl: cacheControlToAnthropic(m.CacheControl), + }) + } + return blocks, nil +} + +// cacheControlToAnthropic converts the provider-agnostic llms.CacheControl +// into Anthropic's Bedrock wire format. +func cacheControlToAnthropic(cc *llms.CacheControl) *anthropicCacheControl { + if cc == nil { + return nil + } + t := cc.Type + if t == "" { + t = "ephemeral" + } + out := &anthropicCacheControl{Type: t} + if cc.Duration > 0 { + out.TTL = "1h" + } + return out } // process the role of the message to anthropic supported role. @@ -462,5 +536,6 @@ func getAnthropicInputContent(message Message) anthropicTextGenerationInputConte Input: input, } } + c.CacheControl = cacheControlToAnthropic(message.CacheControl) return c } diff --git a/llms/bedrock/internal/bedrockclient/provider_cohere.go b/llms/bedrock/internal/bedrockclient/provider_cohere.go index 1ededb2f1..511b3a996 100644 --- a/llms/bedrock/internal/bedrockclient/provider_cohere.go +++ b/llms/bedrock/internal/bedrockclient/provider_cohere.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/tmc/langchaingo/llms" ) @@ -84,13 +83,7 @@ func createCohereCompletion(ctx context.Context, return nil, err } - modelInput := &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - resp, err := client.InvokeModel(ctx, modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } diff --git a/llms/bedrock/internal/bedrockclient/provider_meta.go b/llms/bedrock/internal/bedrockclient/provider_meta.go index 737918712..fe10ab8da 100644 --- a/llms/bedrock/internal/bedrockclient/provider_meta.go +++ b/llms/bedrock/internal/bedrockclient/provider_meta.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/tmc/langchaingo/llms" ) @@ -64,14 +63,7 @@ func createMetaCompletion(ctx context.Context, return nil, err } - modelInput := &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - - resp, err := client.InvokeModel(ctx, modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } diff --git a/llms/bedrock/internal/bedrockclient/provider_nova.go b/llms/bedrock/internal/bedrockclient/provider_nova.go index d452a139b..ba2ceed4c 100644 --- a/llms/bedrock/internal/bedrockclient/provider_nova.go +++ b/llms/bedrock/internal/bedrockclient/provider_nova.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/tmc/langchaingo/llms" ) @@ -161,13 +160,7 @@ func createNovaCompletion(ctx context.Context, return nil, errors.New("streaming not implemented for nova") } - modelInput := &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(modelID), - Accept: aws.String("*/*"), - ContentType: aws.String("application/json"), - Body: body, - } - resp, err := client.InvokeModel(ctx, modelInput) + resp, err := invokeModel(ctx, client, modelID, body, options) if err != nil { return nil, err } diff --git a/llms/bedrock/internal/bedrockclient/tool_call_test.go b/llms/bedrock/internal/bedrockclient/tool_call_test.go index 5f0d08231..22754106e 100644 --- a/llms/bedrock/internal/bedrockclient/tool_call_test.go +++ b/llms/bedrock/internal/bedrockclient/tool_call_test.go @@ -56,8 +56,10 @@ func TestAnthropicToolCallSupport(t *testing.T) { }, StopReason: AnthropicCompletionReasonEndTurn, Usage: struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` }{ InputTokens: 100, OutputTokens: 50, diff --git a/llms/options.go b/llms/options.go index 87b35359c..42cc6dc0f 100644 --- a/llms/options.go +++ b/llms/options.go @@ -73,6 +73,9 @@ type CallOptions struct { // WebSearchOptions configures web search behavior for models that support it. // Currently supported by OpenAI models like gpt-4o-search-preview. WebSearchOptions *WebSearchOptions `json:"web_search_options,omitempty"` + + // Safety configures provider-defined safety controls for the invocation. + SafetyConfig map[string]any `json:"safety,omitempty"` } // Tool is a tool that can be used by the model. diff --git a/llms/options_test.go b/llms/options_test.go index 86cc88fcb..9bc039873 100644 --- a/llms/options_test.go +++ b/llms/options_test.go @@ -160,6 +160,24 @@ func TestCallOptions(t *testing.T) { //nolint:funlen // comprehensive test } }, }, + { + name: "WithSafetyConfig", + option: llms.WithSafetyConfig(map[string]any{ + "identifier": "safety-profile", + "version": "v1", + }), + verify: func(t *testing.T, opts llms.CallOptions) { + if opts.SafetyConfig == nil { + t.Fatal("SafetyConfig = nil, want non-nil") + } + if opts.SafetyConfig["identifier"] != "safety-profile" { + t.Errorf("SafetyConfig[identifier] = %v, want %v", opts.SafetyConfig["identifier"], "safety-profile") + } + if opts.SafetyConfig["version"] != "v1" { + t.Errorf("SafetyConfig[version] = %v, want %v", opts.SafetyConfig["version"], "v1") + } + }, + }, } for _, tt := range tests { diff --git a/llms/safety.go b/llms/safety.go new file mode 100644 index 000000000..53be94400 --- /dev/null +++ b/llms/safety.go @@ -0,0 +1,10 @@ +package llms + +// WithSafetyConfig adds safety configuration to call options. +func WithSafetyConfig(config map[string]any) CallOption { + return func(opts *CallOptions) { + if config != nil { + opts.SafetyConfig = config + } + } +} diff --git a/llms/safety_test.go b/llms/safety_test.go new file mode 100644 index 000000000..9dd8e45b0 --- /dev/null +++ b/llms/safety_test.go @@ -0,0 +1,60 @@ +package llms_test + +import ( + "testing" + + "github.com/tmc/langchaingo/llms" +) + +func TestWithSafetyConfig(t *testing.T) { + var opts llms.CallOptions + + llms.WithSafetyConfig(map[string]any{ + "identifier": "safe-profile", + "version": "v1", + })(&opts) + + if opts.SafetyConfig == nil { + t.Fatal("expected safety config to be present") + } + if opts.SafetyConfig["identifier"] != "safe-profile" { + t.Fatalf("identifier = %q, want %q", opts.SafetyConfig["identifier"], "safe-profile") + } + if opts.SafetyConfig["version"] != "v1" { + t.Fatalf("version = %q, want %q", opts.SafetyConfig["version"], "v1") + } +} + +func TestWithSafetyConfig_Nil(t *testing.T) { + var opts llms.CallOptions + + llms.WithSafetyConfig(nil)(&opts) + + if opts.SafetyConfig != nil { + t.Fatal("expected safety config to be nil") + } +} + +func TestWithSafetyConfig_Overwrite(t *testing.T) { + var opts llms.CallOptions + + llms.WithSafetyConfig(map[string]any{ + "identifier": "safe-profile", + "version": "v1", + })(&opts) + + llms.WithSafetyConfig(map[string]any{ + "identifier": "new-profile", + "version": "v2", + })(&opts) + + if opts.SafetyConfig == nil { + t.Fatal("expected safety config to be present") + } + if opts.SafetyConfig["identifier"] != "new-profile" { + t.Fatalf("identifier = %q, want %q", opts.SafetyConfig["identifier"], "new-profile") + } + if opts.SafetyConfig["version"] != "v2" { + t.Fatalf("version = %q, want %q", opts.SafetyConfig["version"], "v2") + } +}