ai

Documentation

Index

• func EmbeddingHashMatches(embedding *embeddinglib.Embedding, hash embeddinglib.Sha256Hash) bool
• func MarshalEmbedding(embedding *embeddinglib.Embedding) ([]byte, error)
• func UnmarshalEmbedding(bytes []byte) (*embeddinglib.Embedding, error)
• type BatchReducer
  • func NewBatchReducer[T, V any](processFn func(ctx context.Context, data []T) (V, error), batchSize int) *BatchReducer[T, V]
  • func (b *BatchReducer[T, V]) Add(ctx context.Context, data T) (V, error)
  • func (b *BatchReducer[T, V]) Finalize(ctx context.Context) (V, error)
• type Chat
  • func (chat *Chat) Clear()
  • func (chat *Chat) Complete(ctx context.Context, userInput string, ...) (any, *tokens.TokenCount, error)
  • func (chat *Chat) GetMessages() []openai.ChatCompletionMessage
  • func (chat *Chat) Insert(role string, content string) int
  • func (chat *Chat) Reply(ctx context.Context, userInput string, ...) (any, *tokens.TokenCount, error)
• type Client
  • func NewClient(authToken string) *Client
  • func NewClientFromConfig(config openai.ClientConfig) *Client
  • func (client *Client) ClassifyMessage(ctx context.Context, message string, classes map[string]string) (string, error)
  • func (client *Client) CommandSummary(ctx context.Context, messages []openai.ChatCompletionMessage, ...) (string, *tokens.TokenCount, error)
  • func (client *Client) ComputeEmbeddings(ctx context.Context, input []string) ([]embedding.Vector64, error)
  • func (client *Client) NewAuditQuery(username string) *Chat
  • func (client *Client) NewChat(toolContext *modeltools.ToolContext) *Chat
  • func (client *Client) NewCommand(username string) *Chat
  • func (client *Client) RunTool(ctx context.Context, toolContext *modeltools.ToolContext, ...) (any, *tokens.TokenCount, error)
  • func (client *Client) Summary(ctx context.Context, message string) (string, error)
• type Document
• type EmbeddingProcessor
  • func NewEmbeddingProcessor(cfg *EmbeddingProcessorConfig) *EmbeddingProcessor
  • func (e *EmbeddingProcessor) Run(ctx context.Context, initialDelay, period time.Duration) error
• type EmbeddingProcessorConfig
• type Embeddings
• type FilterFn
• type MockEmbedder
  • func (m *MockEmbedder) ComputeEmbeddings(_ context.Context, input []string) ([]embedding.Vector64, error)
• type SimpleRetriever
  • func NewSimpleRetriever() *SimpleRetriever
  • func (r *SimpleRetriever) GetRelevant(query *embedding.Embedding, k int, filter FilterFn) []*Document
  • func (r *SimpleRetriever) Insert(id string, embedding *embedding.Embedding) bool
  • func (r *SimpleRetriever) Remove(id string)
  • func (r *SimpleRetriever) Swap(s *SimpleRetriever)

Functions

func EmbeddingHashMatches

func EmbeddingHashMatches(embedding *embeddinglib.Embedding, hash embeddinglib.Sha256Hash) bool

EmbeddingHashMatches returns true if the hash of the embedding matches the given hash.

func MarshalEmbedding

func MarshalEmbedding(embedding *embeddinglib.Embedding) ([]byte, error)

MarshalEmbedding marshals the ai.Embedding resource to binary ProtoBuf.

func UnmarshalEmbedding

func UnmarshalEmbedding(bytes []byte) (*embeddinglib.Embedding, error)

UnmarshalEmbedding unmarshals binary ProtoBuf into an ai.Embedding resource.

Types

type BatchReducer

type BatchReducer[T, V any] struct {
	// contains filtered or unexported fields
}

BatchReducer is a helper that processes data in batches.

func NewBatchReducer

func NewBatchReducer[T, V any](processFn func(ctx context.Context, data []T) (V, error), batchSize int) *BatchReducer[T, V]

NewBatchReducer is a BatchReducer constructor.

func (*BatchReducer[T, V]) Add

func (b *BatchReducer[T, V]) Add(ctx context.Context, data T) (V, error)

Add adds a new item to the batch. If the batch is full, it will be processed and the result will be returned. Otherwise, a zero value will be returned. Finalize must be called to process the remaining data in the batch.

func (*BatchReducer[T, V]) Finalize

func (b *BatchReducer[T, V]) Finalize(ctx context.Context) (V, error)

Finalize processes the remaining data in the batch and returns the result.

type Chat

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

Chat represents a conversation between a user and an assistant with context memory.

func (*Chat) Clear

func (chat *Chat) Clear()

Clear clears the conversation.

func (*Chat) Complete

func (chat *Chat) Complete(ctx context.Context, userInput string, progressUpdates func(*model.AgentAction)) (any, *tokens.TokenCount, error)

Complete completes the conversation with a message from the assistant based on the current context and user input. On success, it returns the message. Returned types: - message: one of the message types below - error: an error if one occurred Message types: - CompletionCommand: a command from the assistant - Message: a text message from the assistant - AccessRequest: an access request suggestion from the assistant

func (*Chat) GetMessages

func (chat *Chat) GetMessages() []openai.ChatCompletionMessage

GetMessages returns the messages in the conversation.

func (*Chat) Insert

func (chat *Chat) Insert(role string, content string) int

Insert inserts a message into the conversation. Returns the index of the message.

func (*Chat) Reply

func (chat *Chat) Reply(ctx context.Context, userInput string, progressUpdates func(*model.AgentAction)) (any, *tokens.TokenCount, error)

Reply replies to the user input with a message from the assistant based on the current context.

type Client

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

Client is a client for OpenAI API.

func NewClient

func NewClient(authToken string) *Client

NewClient creates a new client for OpenAI API.

func NewClientFromConfig

func NewClientFromConfig(config openai.ClientConfig) *Client

NewClientFromConfig creates a new client for OpenAI API from config.

func (*Client) ClassifyMessage

func (client *Client) ClassifyMessage(ctx context.Context, message string, classes map[string]string) (string, error)

ClassifyMessage takes a user message, a list of categories, and uses the AI mode as a zero-shot classifier.

func (*Client) CommandSummary

func (client *Client) CommandSummary(ctx context.Context, messages []openai.ChatCompletionMessage, output map[string][]byte) (string, *tokens.TokenCount, error)

CommandSummary creates a command summary based on the command output. The message history is also passed to the model in order to keep context and extract relevant information from the output.

func (*Client) ComputeEmbeddings

func (client *Client) ComputeEmbeddings(ctx context.Context, input []string) ([]embedding.Vector64, error)

ComputeEmbeddings takes a map of nodes and calls openAI to generate embeddings for those nodes. ComputeEmbeddings is responsible for implementing a retry mechanism if the embedding computation is flaky.

func (*Client) NewAuditQuery

func (client *Client) NewAuditQuery(username string) *Chat

func (*Client) NewChat

func (client *Client) NewChat(toolContext *modeltools.ToolContext) *Chat

NewChat creates a new chat. The username is set in the conversation context, so that the AI can use it to personalize the conversation. embeddingServiceClient is used to get the embeddings from the Auth Server.

func (*Client) NewCommand

func (client *Client) NewCommand(username string) *Chat

func (*Client) RunTool

func (client *Client) RunTool(ctx context.Context, toolContext *modeltools.ToolContext, toolName, toolInput string) (any, *tokens.TokenCount, error)

func (*Client) Summary

func (client *Client) Summary(ctx context.Context, message string) (string, error)

Summary creates a short summary for the given input.

type Document

type Document struct {
	*embedding.Embedding
	SimilarityScore float64
}

Document is a embedding enriched with similarity score

type EmbeddingProcessor

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

EmbeddingProcessor is responsible for processing nodes, generating embeddings and storing their embeddings in the backend.

func NewEmbeddingProcessor

func NewEmbeddingProcessor(cfg *EmbeddingProcessorConfig) *EmbeddingProcessor

NewEmbeddingProcessor returns a new EmbeddingProcessor.

func (*EmbeddingProcessor) Run

func (e *EmbeddingProcessor) Run(ctx context.Context, initialDelay, period time.Duration) error

Run runs the EmbeddingProcessor.

type EmbeddingProcessorConfig

type EmbeddingProcessorConfig struct {
	AIClient            embeddinglib.Embedder
	EmbeddingSrv        Embeddings
	EmbeddingsRetriever *SimpleRetriever
	NodeSrv             *services.UnifiedResourceCache
	Log                 logrus.FieldLogger
	Jitter              retryutils.Jitter
}

EmbeddingProcessorConfig is the configuration for EmbeddingProcessor.

type Embeddings

type Embeddings interface {
	// GetAllEmbeddings returns all embeddings.
	GetAllEmbeddings(ctx context.Context) stream.Stream[*embeddinglib.Embedding]

	// UpsertEmbedding creates or update a single ai.Embedding in the backend.
	UpsertEmbedding(ctx context.Context, embedding *embeddinglib.Embedding) (*embeddinglib.Embedding, error)
}

Embeddings implements the minimal interface used by the Embedding processor.

type FilterFn

type FilterFn func(id string, embedding *embedding.Embedding) bool

FilterFn is a function that filters out embeddings. If the function returns false, the embedding is filtered out.

type MockEmbedder

type MockEmbedder struct {
	TimesCalled map[string]int
	// contains filtered or unexported fields
}

MockEmbedder returns embeddings based on the sha256 hash function. Those embeddings have no semantic meaning but ensure different embedded content provides different embeddings.

func (*MockEmbedder) ComputeEmbeddings

func (m *MockEmbedder) ComputeEmbeddings(_ context.Context, input []string) ([]embedding.Vector64, error)

type SimpleRetriever

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

SimpleRetriever is a simple implementation of embeddings retriever. It stores all the embeddings in memory and retrieves the k nearest neighbors by iterating over all the embeddings. Do not use for large datasets.

func NewSimpleRetriever

func NewSimpleRetriever() *SimpleRetriever

func (*SimpleRetriever) GetRelevant

func (r *SimpleRetriever) GetRelevant(query *embedding.Embedding, k int, filter FilterFn) []*Document

GetRelevant returns the k nearest neighbors to the query embedding. If a filter is provided, only the embeddings that pass the filter are considered.

func (*SimpleRetriever) Insert

func (r *SimpleRetriever) Insert(id string, embedding *embedding.Embedding) bool

Insert adds the embedding to the retriever. If the retriever is full, the embedding is not added and false is returned.

func (*SimpleRetriever) Remove

func (r *SimpleRetriever) Remove(id string)

Remove removes the embedding from the retriever by ID.

func (*SimpleRetriever) Swap

func (r *SimpleRetriever) Swap(s *SimpleRetriever)

Swap replaces the embeddings in the retriever with the embeddings from the provided retriever. The mutex is acquired for the receiver, but not for the provided retriever.