llm

Documentation

Index

• func CountTokens(text string, messages []llms.ChatMessage) int
• func CreateAugPromptFromCmdResults(cfg *config.Config, prompt string, cmdResults []config.CmdRes) (augPrompt string, err error)
• func GenKubectlCmds(cfg *config.Config, prompt string, userMsgCount int) ([]string, error)
• func GetEmbedder(cfg *config.Config) (embeddings.Embedder, error)
• func HandleLLMRequest(cfg *config.Config, client llms.Model, prompt string, stream bool, ...) (string, error)
• func ManageChatThreadContext(chatMessages []llms.ChatMessage, maxTokens int)
• func PromptExistsInHistory(messages []llms.ChatMessage, prompt string) bool
• func RetrieveRAG(cfg *config.Config, prompt string, lastTextPrompt string, userMsgCount int) (augPrompt string, err error)
• func TrimSectionsWithEmbeddings(ctx context.Context, embedder embeddings.Embedder, sections []string, ...) []string
• type GoogleEmbedder
  • func (g *GoogleEmbedder) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error)
  • func (g *GoogleEmbedder) EmbedQuery(ctx context.Context, text string) ([]float32, error)
• type RequestFunc

Functions

func CountTokens

func CountTokens(text string, messages []llms.ChatMessage) int

CountTokens counts the token usage for a given text string and/or chat messages The function combines both tokens from text and the given messages

func CreateAugPromptFromCmdResults

func CreateAugPromptFromCmdResults(cfg *config.Config, prompt string, cmdResults []config.CmdRes) (augPrompt string, err error)

CreateAugPromptFromCmdResults formats command results for RAG in the same way as retrieveRAG

func GenKubectlCmds

func GenKubectlCmds(cfg *config.Config, prompt string, userMsgCount int) ([]string, error)

GenKubectlCmds retrieves kubectl commands based on the user input

func GetEmbedder

func GetEmbedder(cfg *config.Config) (embeddings.Embedder, error)

GetEmbedder creates an embedder based on the provider configuration

func HandleLLMRequest

func HandleLLMRequest(cfg *config.Config, client llms.Model, prompt string, stream bool, history bool) (string, error)

handleLLMRequestNoHistoryUpdate is a modified version of handleLLMRequest that doesn't add the human message because it's already added in googleRequestWithChat

func ManageChatThreadContext

func ManageChatThreadContext(chatMessages []llms.ChatMessage, maxTokens int)

ManageChatThreadContext manages the context window of the chat thread

func PromptExistsInHistory

func PromptExistsInHistory(messages []llms.ChatMessage, prompt string) bool

Helper function to check if a prompt exists in chat history

func RetrieveRAG

func RetrieveRAG(cfg *config.Config, prompt string, lastTextPrompt string, userMsgCount int) (augPrompt string, err error)

RetrieveRAG retrieves the data for RAG

func TrimSectionsWithEmbeddings

func TrimSectionsWithEmbeddings(ctx context.Context, embedder embeddings.Embedder, sections []string, prompt string, maxTokens int) []string

trimSectionsWithLangChainGo uses langchaingo embeddings to select the most semantically relevant sections within a token budget.

Types

type GoogleEmbedder

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

GoogleEmbedder is a wrapper around the GoogleAI client that provides embedding functionality

func (*GoogleEmbedder) EmbedDocuments

func (g *GoogleEmbedder) EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error)

EmbedDocuments implements the Embedder interface for GoogleEmbedder

func (*GoogleEmbedder) EmbedQuery

func (g *GoogleEmbedder) EmbedQuery(ctx context.Context, text string) ([]float32, error)

EmbedQuery implements the Embedder interface for GoogleEmbedder

type RequestFunc

type RequestFunc func(cfg *config.Config, prompt string, stream bool, history bool) (string, error)

Define RequestFunc type for easier mocking in tests

var Request RequestFunc = func(cfg *config.Config, prompt string, stream bool, history bool) (string, error) {
	truncPrompt := prompt

	if len(truncPrompt) > cfg.MaxTokens*2 {
		truncPrompt = truncPrompt[:cfg.MaxTokens*2] + "..."
	}

	logger.Log("llmIn", "[%s/%s]: %s", cfg.Provider, cfg.Model, truncPrompt)
	logger.Log("llmIn", "History: %v messages, %d tokens", len(cfg.ChatMessages), CountTokens("", cfg.ChatMessages))

	s := spinner.New(spinner.CharSets[11], time.Duration(cfg.SpinnerTimeout)*time.Millisecond)
	s.Suffix = fmt.Sprintf(" Waiting for %s/%s response...", cfg.Provider, cfg.Model)
	s.Color("green", "bold")

	if !stream {
		s.Start()
		defer s.Stop()
	}

	var err error
	var answer string
	switch cfg.Provider {
	case "ollama":
		answer, err = ollamaRequestWithChat(cfg, truncPrompt, stream, history)
	case "openai", "deepseek":
		answer, err = openaiRequestWithChat(cfg, truncPrompt, stream, history)
	case "google":
		answer, err = googleRequestWithChat(cfg, truncPrompt, stream, history)
	case "anthropic":
		answer, err = anthropicRequestWithChat(cfg, truncPrompt, stream, history)
	default:
		return "", fmt.Errorf("unsupported AI provider: %s", cfg.Provider)
	}

	logger.Log("llmOut", "[%s@%s]: %s", cfg.Provider, cfg.Model, answer)
	return answer, err
}

Request sends a request to the LLM provider