ml

package
v1.0.10 Latest Latest
Warning

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

Go to latest
Published: Aug 10, 2025 License: MIT Imports: 11 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func EnhanceFindingWithML

func EnhanceFindingWithML(finding map[string]interface{}, models *MLModels) map[string]interface{}

EnhanceFindingWithML enhances a finding with ML predictions

func ExtractFeatures

func ExtractFeatures(finding map[string]interface{}, encoders map[string]*LabelEncoder) []float64

ExtractFeatures extracts features from a finding for ML prediction

func SaveEvaluationResults

func SaveEvaluationResults(evaluations []AIEvaluation, summary *AIEvaluationSummary, outputDir string) error

SaveEvaluationResults saves the evaluation results to JSON files

func SaveModelToJSON

func SaveModelToJSON(model MLModel) ([]byte, error)

SaveModelToJSON saves a model to JSON

Types

type AIEvaluation

type AIEvaluation struct {
	FindingID           string  `json:"finding_id"`
	IsValidFinding      bool    `json:"is_valid_finding"`
	Confidence          float64 `json:"confidence"`
	Severity            string  `json:"severity"`
	Priority            string  `json:"priority"`
	Reasoning           string  `json:"reasoning"`
	MigrationSuggestion string  `json:"migration_suggestion"`
	FalsePositiveReason string  `json:"false_positive_reason,omitempty"`
	TokensUsed          int     `json:"tokens_used"`
	Cost                float64 `json:"cost"`
}

AIEvaluation represents the AI's evaluation of a finding

type AIEvaluationRequest

type AIEvaluationRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Temperature float64   `json:"temperature"`
	MaxTokens   int       `json:"max_tokens"`
}

AIEvaluationRequest represents a request for AI evaluation

type AIEvaluationResponse

type AIEvaluationResponse struct {
	ID      string    `json:"id"`
	Object  string    `json:"object"`
	Created int64     `json:"created"`
	Model   string    `json:"model"`
	Choices []Choice  `json:"choices"`
	Usage   Usage     `json:"usage"`
	Error   *APIError `json:"error,omitempty"`
}

AIEvaluationResponse represents OpenAI API response

type AIEvaluationSummary

type AIEvaluationSummary struct {
	TotalFindings      int      `json:"total_findings"`
	ValidFindings      int      `json:"valid_findings"`
	FalsePositives     int      `json:"false_positives"`
	TotalTokensUsed    int      `json:"total_tokens_used"`
	TotalCost          float64  `json:"total_cost"`
	AverageConfidence  float64  `json:"average_confidence"`
	ProcessingTime     string   `json:"processing_time"`
	RecommendedActions []string `json:"recommended_actions"`
}

AIEvaluationSummary represents summary statistics

type APIError

type APIError struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    string `json:"code"`
}

APIError represents an API error

type Choice

type Choice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

Choice represents a response choice

type DecisionTree

type DecisionTree struct {
	Root         *DecisionTreeNode      `json:"root"`
	FeatureNames []string               `json:"feature_names"`
	Version      string                 `json:"version"`
	Metadata     map[string]interface{} `json:"metadata"`
}

DecisionTree implements a lightweight decision tree model

func NewDecisionTree

func NewDecisionTree(featureNames []string) *DecisionTree

NewDecisionTree creates a new decision tree model

func (*DecisionTree) GetFeatureNames

func (dt *DecisionTree) GetFeatureNames() []string

GetFeatureNames returns the feature names

func (*DecisionTree) GetModelType

func (dt *DecisionTree) GetModelType() string

GetModelType returns the model type

func (*DecisionTree) GetVersion

func (dt *DecisionTree) GetVersion() string

GetVersion returns the model version

func (*DecisionTree) Predict

func (dt *DecisionTree) Predict(features []float64) float64

Predict returns the prediction for given features

type DecisionTreeNode

type DecisionTreeNode struct {
	FeatureIndex int               `json:"feature_index"`
	Threshold    float64           `json:"threshold"`
	Value        float64           `json:"value"` // For leaf nodes
	Left         *DecisionTreeNode `json:"left"`
	Right        *DecisionTreeNode `json:"right"`
	IsLeaf       bool              `json:"is_leaf"`
}

DecisionTreeNode represents a node in a decision tree

type LabelEncoder

type LabelEncoder struct {
	LabelToInt map[string]int
	IntToLabel map[int]string
	NClasses   int
}

LabelEncoder represents a label encoder

func (*LabelEncoder) Decode

func (le *LabelEncoder) Decode(value int) string

Decode converts an integer to string label

func (*LabelEncoder) Encode

func (le *LabelEncoder) Encode(label string) int

Encode converts a string label to integer

type LinearRegression

type LinearRegression struct {
	Weights      []float64              `json:"weights"`
	Bias         float64                `json:"bias"`
	FeatureNames []string               `json:"feature_names"`
	Version      string                 `json:"version"`
	Metadata     map[string]interface{} `json:"metadata"`
}

LinearRegression implements a lightweight linear regression model

func NewLinearRegression

func NewLinearRegression(featureNames []string) *LinearRegression

NewLinearRegression creates a new linear regression model

func (*LinearRegression) GetFeatureNames

func (lr *LinearRegression) GetFeatureNames() []string

GetFeatureNames returns the feature names

func (*LinearRegression) GetModelType

func (lr *LinearRegression) GetModelType() string

GetModelType returns the model type

func (*LinearRegression) GetVersion

func (lr *LinearRegression) GetVersion() string

GetVersion returns the model version

func (*LinearRegression) Predict

func (lr *LinearRegression) Predict(features []float64) float64

Predict returns the prediction for given features

type LogisticRegressionModel

type LogisticRegressionModel struct {
	Coefficients []float64
	Intercept    float64
	Classes      []int
	NFeatures    int
}

LogisticRegressionModel represents a logistic regression model

func (*LogisticRegressionModel) Predict

func (lr *LogisticRegressionModel) Predict(features []float64) float64

Predict using logistic regression

type MLModel

type MLModel interface {
	Predict(features []float64) float64
	GetFeatureNames() []string
	GetModelType() string
	GetVersion() string
}

MLModel represents a lightweight ML model for confidence scoring

func LoadModelFromJSON

func LoadModelFromJSON(data []byte, modelType string) (MLModel, error)

LoadModelFromJSON loads a model from JSON data

type MLModels

type MLModels struct {
	FalsePositiveDetector *LogisticRegressionModel
	ConfidencePredictor   *RandomForestModel
	SeverityClassifier    *RandomForestModel
	Encoders              map[string]*LabelEncoder
}

MLModels contains all converted ML models

func NewMLModels

func NewMLModels() *MLModels

NewMLModels creates a new MLModels instance with trained models

type Message

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

Message represents a chat message

type ModelRegistry

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

ModelRegistry manages multiple ML models

func NewModelRegistry

func NewModelRegistry() *ModelRegistry

NewModelRegistry creates a new model registry

func (*ModelRegistry) GetModel

func (mr *ModelRegistry) GetModel(name string) (MLModel, bool)

GetModel retrieves a model by name

func (*ModelRegistry) ListModels

func (mr *ModelRegistry) ListModels() []string

ListModels returns all registered model names

func (*ModelRegistry) PredictWithFallback

func (mr *ModelRegistry) PredictWithFallback(features []float64, modelNames ...string) float64

PredictWithFallback tries multiple models with fallback

func (*ModelRegistry) RegisterModel

func (mr *ModelRegistry) RegisterModel(name string, model MLModel)

RegisterModel registers a model with a name

type OpenAIClient

type OpenAIClient struct {
	APIKey     string
	BaseURL    string
	Model      string
	HTTPClient *http.Client
}

OpenAIClient handles communication with OpenAI API

func NewOpenAIClient

func NewOpenAIClient(apiKey string) *OpenAIClient

NewOpenAIClient creates a new OpenAI client

func (*OpenAIClient) EvaluateFindings

func (client *OpenAIClient) EvaluateFindings(findings []types.Finding, batchSize int) ([]AIEvaluation, *AIEvaluationSummary, error)

EvaluateFindings sends findings to AI for evaluation in batches

func (*OpenAIClient) EvaluateFindingsWithProgress

func (client *OpenAIClient) EvaluateFindingsWithProgress(findings []types.Finding, batchSize int, progressCallback func(int, int)) ([]AIEvaluation, *AIEvaluationSummary, error)

EvaluateFindingsWithProgress sends findings to AI for evaluation in batches with progress callback

type RandomForest

type RandomForest struct {
	Trees        []*DecisionTree        `json:"trees"`
	FeatureNames []string               `json:"feature_names"`
	Version      string                 `json:"version"`
	Metadata     map[string]interface{} `json:"metadata"`
}

RandomForest implements a lightweight random forest model

func NewRandomForest

func NewRandomForest(featureNames []string) *RandomForest

NewRandomForest creates a new random forest model

func (*RandomForest) AddTree

func (rf *RandomForest) AddTree(tree *DecisionTree)

AddTree adds a decision tree to the forest

func (*RandomForest) GetFeatureNames

func (rf *RandomForest) GetFeatureNames() []string

GetFeatureNames returns the feature names

func (*RandomForest) GetModelType

func (rf *RandomForest) GetModelType() string

GetModelType returns the model type

func (*RandomForest) GetVersion

func (rf *RandomForest) GetVersion() string

GetVersion returns the model version

func (*RandomForest) Predict

func (rf *RandomForest) Predict(features []float64) float64

Predict returns the average prediction from all trees

type RandomForestModel

type RandomForestModel struct {
	FeatureImportances []float64
	Classes            []string
	NFeatures          int
	NEstimators        int
}

RandomForestModel represents a simplified random forest model

func (*RandomForestModel) PredictClass

func (rf *RandomForestModel) PredictClass(features []float64) string

PredictClass using random forest (simplified)

type Usage

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

Usage represents token usage information

Jump to

Keyboard shortcuts

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