llm

package
v1.2.1 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jul 1, 2025 License: MIT Imports: 10 Imported by: 0

Documentation

Overview

Package llm provides integration with Large Language Models for advanced PI validation. It supports multiple LLM providers through a unified interface and uses AI to analyze code context and reduce false positives in PI detection.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func CreateEnhancedDetector

func CreateEnhancedDetector(baseDetector detection.Detector, config *detection.Config) (detection.Detector, error)

CreateEnhancedDetector wraps a detector with LLM validation if enabled

func NewValidatorFromConfig

func NewValidatorFromConfig(config *detection.Config) (detection.LLMValidator, error)

NewValidatorFromConfig creates an LLM validator from detection config with security validation

Types

type Capabilities

type Capabilities struct {
	MaxTokens         int  `json:"max_tokens"`
	SupportsStreaming bool `json:"supports_streaming"`
	SupportsJSON      bool `json:"supports_json"`
	SupportsFunctions bool `json:"supports_functions"`
}

Capabilities describes what the LLM client supports

type Client

type Client interface {
	// Complete generates a completion for the given prompt
	Complete(ctx context.Context, prompt string, options *CompletionOptions) (string, error)

	// Stream generates a streaming completion
	Stream(ctx context.Context, prompt string, options *CompletionOptions, callback StreamCallback) error

	// GetCapabilities returns the client's capabilities
	GetCapabilities() Capabilities
}

Client is a generic interface for LLM operations

func NewClientAdapter

func NewClientAdapter(validator interface{}) Client

NewClientAdapter creates a new adapter

type ClientAdapter

type ClientAdapter struct {
	// contains filtered or unexported fields
}

ClientAdapter adapts a detection.LLMValidator to the generic Client interface

func (*ClientAdapter) Complete

func (a *ClientAdapter) Complete(ctx context.Context, prompt string, options *CompletionOptions) (string, error)

Complete implements Client.Complete by wrapping the validator

func (*ClientAdapter) GetCapabilities

func (a *ClientAdapter) GetCapabilities() Capabilities

GetCapabilities implements Client.GetCapabilities

func (*ClientAdapter) Stream

func (a *ClientAdapter) Stream(ctx context.Context, prompt string, options *CompletionOptions, callback StreamCallback) error

Stream implements Client.Stream

type CompletionOptions

type CompletionOptions struct {
	MaxTokens   int      `json:"max_tokens,omitempty"`
	Temperature float32  `json:"temperature,omitempty"`
	TopP        float32  `json:"top_p,omitempty"`
	Stop        []string `json:"stop,omitempty"`
}

CompletionOptions contains options for completion requests

type Config

type Config struct {
	Enabled     bool
	Provider    string
	Endpoint    string
	Model       string
	APIKey      string
	MaxTokens   int
	Temperature float32
	Timeout     time.Duration
}

Config holds LLM client configuration

type EndpointValidator

type EndpointValidator struct {
	// contains filtered or unexported fields
}

EndpointValidator validates LLM endpoints for security compliance

func NewEndpointValidator

func NewEndpointValidator(config ValidationConfig) *EndpointValidator

NewEndpointValidator creates a new endpoint validator

func (*EndpointValidator) TestEndpoint

func (v *EndpointValidator) TestEndpoint(ctx context.Context, endpoint string) error

TestEndpoint performs a health check on the endpoint

func (*EndpointValidator) ValidateConfig

func (v *EndpointValidator) ValidateConfig(config Config) error

ValidateConfig validates the LLM configuration

func (*EndpointValidator) ValidateEndpoint

func (v *EndpointValidator) ValidateEndpoint(endpoint string) error

ValidateEndpoint validates an LLM endpoint URL

func (*EndpointValidator) ValidateProviderConfig

func (v *EndpointValidator) ValidateProviderConfig(provider string, config Config) error

ValidateProviderConfig validates provider-specific configurations

type LMStudioClient

type LMStudioClient struct {
	// contains filtered or unexported fields
}

LMStudioClient implements the LLMValidator interface for LM Studio

func NewLMStudioClient

func NewLMStudioClient(config Config) (*LMStudioClient, error)

NewLMStudioClient creates a new LM Studio client

func (*LMStudioClient) Complete

func (c *LMStudioClient) Complete(ctx context.Context, prompt string, options *CompletionOptions) (string, error)

Complete implements the Client interface for generic completions

func (*LMStudioClient) GetCapabilities

func (c *LMStudioClient) GetCapabilities() Capabilities

GetCapabilities implements the Client interface

func (*LMStudioClient) HealthCheck

func (c *LMStudioClient) HealthCheck(ctx context.Context) error

HealthCheck verifies the LLM service is accessible

func (*LMStudioClient) Stream

func (c *LMStudioClient) Stream(ctx context.Context, prompt string, options *CompletionOptions, callback StreamCallback) error

Stream implements the Client interface for streaming completions

func (*LMStudioClient) ValidateFinding

ValidateFinding validates a PI finding using LLM

type ProviderConstructor

type ProviderConstructor func(config Config) (detection.LLMValidator, error)

ProviderConstructor is a function that creates an LLM validator

type SecureLLMConfig

type SecureLLMConfig struct {
	// contains filtered or unexported fields
}

SecureLLMConfig provides a secure configuration builder

func NewSecureLLMConfig

func NewSecureLLMConfig() *SecureLLMConfig

NewSecureLLMConfig creates a new secure configuration with safe defaults

func (*SecureLLMConfig) Build

func (c *SecureLLMConfig) Build() Config

Build returns the configuration

func (*SecureLLMConfig) WithMaxTokens

func (c *SecureLLMConfig) WithMaxTokens(tokens int) *SecureLLMConfig

WithMaxTokens sets max tokens (clamped to safe range)

func (*SecureLLMConfig) WithModel

func (c *SecureLLMConfig) WithModel(model string) *SecureLLMConfig

WithModel sets the model name

func (*SecureLLMConfig) WithProvider

func (c *SecureLLMConfig) WithProvider(provider string) *SecureLLMConfig

WithProvider sets the provider (validated against allowed list)

func (*SecureLLMConfig) WithTemperature

func (c *SecureLLMConfig) WithTemperature(temp float32) *SecureLLMConfig

WithTemperature sets temperature (clamped to valid range)

type StreamCallback

type StreamCallback func(chunk string) error

StreamCallback is called for each chunk of streaming response

type ValidationConfig

type ValidationConfig struct {
	// AllowedHosts is a list of allowed hostnames (default: localhost only)
	AllowedHosts []string

	// AllowedNetworks is a list of allowed IP networks in CIDR notation
	AllowedNetworks []string

	// AllowPrivateIPs allows RFC1918 private IP addresses
	AllowPrivateIPs bool

	// RequireHTTPS enforces HTTPS for non-localhost endpoints
	RequireHTTPS bool

	// MaxResponseTime is the maximum allowed response time for health checks
	MaxResponseTime time.Duration

	// EnableStrictMode enforces localhost-only connections
	EnableStrictMode bool
}

ValidationConfig configures endpoint validation

func DefaultValidationConfig

func DefaultValidationConfig() ValidationConfig

DefaultValidationConfig returns a secure default configuration

type ValidatorFactory

type ValidatorFactory struct {
	// contains filtered or unexported fields
}

ValidatorFactory creates LLM validators with proper security validation

func GetDefaultFactory

func GetDefaultFactory() *ValidatorFactory

GetDefaultFactory returns the default validator factory with security enforcement

func NewValidatorFactory

func NewValidatorFactory() *ValidatorFactory

NewValidatorFactory creates a new validator factory with security enforcement

func (*ValidatorFactory) CreateValidator

func (f *ValidatorFactory) CreateValidator(config Config) (detection.LLMValidator, error)

CreateValidator creates an LLM validator with security validation

func (*ValidatorFactory) RegisterProvider

func (f *ValidatorFactory) RegisterProvider(name string, constructor ProviderConstructor)

RegisterProvider registers a new LLM provider

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL