chat

Documentation

Overview

Package chat provides functionality for interacting with Large Language Models (LLMs) through both direct prompts and streaming interfaces.

This package implements core chat functionality for the aigent application, including:

• Direct prompt handling for simple question-answer interactions
• Streaming chat responses for real-time UI updates
• Tool execution capabilities allowing LLMs to perform actions
• Session management for maintaining conversation context

The package is designed around several key components:

• ChatSession: Represents a conversation with an LLM, maintaining message history
• StreamHandler: Interface for handling streaming responses from LLMs
• ToolRegistry: Manages available tools that LLMs can use to perform actions

Typical usage involves creating a ChatSession, configuring options through one of the available Options structs (CoreChatOptions, PromptOptions, or StreamOptions), and then either handling a direct prompt or processing a streaming chat response.

Tool execution is a central feature, allowing LLMs to perform actions like retrieving the current time. The package provides a SetupToolRegistry function to configure default tools, and both direct and streaming interfaces support tool execution.

Example usage for direct prompt handling:

options := chat.PromptOptions{
	CoreChatOptions: chat.CoreChatOptions{
		Client:       client,
		Provider:     "openai",
		ModelID:      "gpt-4",
		Logger:       logger,
		ToolRegistry: chat.SetupToolRegistry(logger),
		Temperature:  0.7,
		MaxTokens:    2000,
	},
	EnableTools: true,
}

err := chat.HandleDirectPrompt(ctx, "What time is it?", options)

Example usage for streaming chat:

session := &chat.ChatSession{
	Messages: []llm.Message{
		{Role: llm.RoleUser, Content: "What time is it?"},
	},
}

handler := &MyStreamHandler{} // Implements StreamHandler interface

options := chat.StreamOptions{
	CoreChatOptions: chat.CoreChatOptions{
		Client:       client,
		Provider:     "openai",
		ModelID:      "gpt-4",
		Logger:       logger,
		ToolRegistry: chat.SetupToolRegistry(logger),
		Temperature:  0.7,
		MaxTokens:    2000,
	},
	Handler: handler,
}

err := chat.ProcessChatStream(ctx, session, options)

Index

• func AddMCPToolsToPersona(logger *slog.Logger, p persona.Persona, mcpManager mcp.Manager)
• func CleanupAllPersonas(logger *slog.Logger, personaRegistry *persona.Registry)
• func GetSystemPromptFromPersona(logger *slog.Logger, personaRegistry *persona.Registry, personaID string) string
• func GetToolRegistryFromPersona(logger *slog.Logger, personaRegistry *persona.Registry, personaID string, ...) *llm.ToolRegistry
• func HandleDirectPrompt(ctx context.Context, prompt string, options PromptOptions) error
• func ProcessChatStream(ctx context.Context, session *ChatSession, options StreamOptions) error
• func SetupPersonaRegistry(logger *slog.Logger) *persona.Registry
• func SetupToolRegistry(logger *slog.Logger) *llm.ToolRegistry
• type ChatSession
• type ConsoleStreamHandler
  • func (h *ConsoleStreamHandler) OnChunk(chunk *llm.ChatResponseChunk)
  • func (h *ConsoleStreamHandler) OnComplete(fullContent string, toolCalls []llm.ToolCall, err error)
• type CoreChatOptions
• type MCPManager
  • func NewMCPManager(logger *slog.Logger) *MCPManager
  • func (m *MCPManager) AddServer(ctx context.Context, cfg mcp.ServerConfig) error
  • func (m *MCPManager) Close() error
  • func (m *MCPManager) GetServer(name string) (*mcp.Server, bool)
  • func (m *MCPManager) Servers() map[string]*mcp.Server
• type MCPServer
• type MCPServerConfig
• type NoopStreamHandler
  • func (h *NoopStreamHandler) OnChunk(chunk *llm.ChatResponseChunk)
  • func (h *NoopStreamHandler) OnComplete(fullContent string, toolCalls []llm.ToolCall, err error)
  • func (h *NoopStreamHandler) OnToolExecution(toolCall llm.ToolCall, result string, err error)
• type PromptOptions
• type StreamHandler
• type StreamOptions

Functions

func AddMCPToolsToPersona

func AddMCPToolsToPersona(logger *slog.Logger, p persona.Persona, mcpManager mcp.Manager)

AddMCPToolsToPersona adds MCP tools to a persona

func CleanupAllPersonas

func CleanupAllPersonas(logger *slog.Logger, personaRegistry *persona.Registry)

CleanupAllPersonas calls the Cleanup method on all registered personas

func GetSystemPromptFromPersona

func GetSystemPromptFromPersona(logger *slog.Logger, personaRegistry *persona.Registry, personaID string) string

GetSystemPromptFromPersona returns the system prompt for the specified persona. If the persona is not found or an error occurs, it returns a default system prompt.

func GetToolRegistryFromPersona

func GetToolRegistryFromPersona(logger *slog.Logger, personaRegistry *persona.Registry, personaID string, mcpManager mcp.Manager) *llm.ToolRegistry

GetToolRegistryFromPersona creates a tool registry based on the specified persona. If the persona is not found or an error occurs, it falls back to the default tool registry.

func HandleDirectPrompt

func HandleDirectPrompt(ctx context.Context, prompt string, options PromptOptions) error

HandleDirectPrompt processes a prompt directly without UI and returns the response. It handles the complete lifecycle of a prompt, including potential tool calls and their execution. The function continues processing until no more tool calls are returned or an error occurs.

func ProcessChatStream

func ProcessChatStream(ctx context.Context, session *ChatSession, options StreamOptions) error

ProcessChatStream handles a streaming chat response, including tool execution. It processes the stream chunk by chunk, accumulating content and tool calls, and executes tools when the stream is complete. The function will recursively call itself to continue the conversation after tool execution.

func SetupPersonaRegistry

func SetupPersonaRegistry(logger *slog.Logger) *persona.Registry

SetupPersonaRegistry creates and configures a persona registry with available personas.

func SetupToolRegistry

func SetupToolRegistry(logger *slog.Logger) *llm.ToolRegistry

SetupToolRegistry creates and configures a tool registry with default tools. It initializes a new tool registry and registers the get_time tool, which allows LLMs to retrieve the current time.

Types

type ChatSession

type ChatSession struct {
	Messages []llm.Message
}

ChatSession represents a single conversation with an LLM. It maintains the history of messages exchanged during the conversation, which provides context for the LLM to generate appropriate responses.

type ConsoleStreamHandler

type ConsoleStreamHandler struct {
	NoopStreamHandler
}

ConsoleStreamHandler handles streaming output to the console

func (*ConsoleStreamHandler) OnChunk

func (h *ConsoleStreamHandler) OnChunk(chunk *llm.ChatResponseChunk)

func (*ConsoleStreamHandler) OnComplete

func (h *ConsoleStreamHandler) OnComplete(fullContent string, toolCalls []llm.ToolCall, err error)

type CoreChatOptions

type CoreChatOptions struct {
	Client       *llm.Client
	Provider     string
	ModelID      string
	Logger       *slog.Logger
	ToolRegistry *llm.ToolRegistry
	Temperature  float32
	MaxTokens    int
	NoStream     bool   // Whether to disable streaming responses
	SystemPrompt string // System prompt to use for the chat
}

CoreChatOptions contains configuration options shared between UI and non-UI chat functionality. It provides the core settings needed for interacting with LLMs, including client configuration, model selection, logging, tool support, and generation parameters.

type MCPManager

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

MCPManager manages multiple MCP servers

func NewMCPManager

func NewMCPManager(logger *slog.Logger) *MCPManager

NewMCPManager creates a new MCP manager

func (*MCPManager) AddServer

func (m *MCPManager) AddServer(ctx context.Context, cfg mcp.ServerConfig) error

AddServer adds and initializes a new MCP server

func (*MCPManager) Close

func (m *MCPManager) Close() error

Close closes all MCP server connections

func (*MCPManager) GetServer

func (m *MCPManager) GetServer(name string) (*mcp.Server, bool)

GetServer returns a specific MCP server by name

func (*MCPManager) Servers

func (m *MCPManager) Servers() map[string]*mcp.Server

Servers returns all registered MCP servers

type MCPServer

type MCPServer struct {
	Client *mcplib.Client
	Tools  []mcplib.ToolRetType
	// contains filtered or unexported fields
}

MCPServer represents a connected MCP server with its tools and the underlying command

type MCPServerConfig

type MCPServerConfig = mcp.ServerConfig

MCPServerConfig is an alias for mcp.ServerConfig for backward compatibility

type NoopStreamHandler

type NoopStreamHandler struct{}

NoopStreamHandler is a stream handler implementation that does nothing. It can be used as a default handler when no specific handling is needed.

func (*NoopStreamHandler) OnChunk

func (h *NoopStreamHandler) OnChunk(chunk *llm.ChatResponseChunk)

func (*NoopStreamHandler) OnComplete

func (h *NoopStreamHandler) OnComplete(fullContent string, toolCalls []llm.ToolCall, err error)

func (*NoopStreamHandler) OnToolExecution

func (h *NoopStreamHandler) OnToolExecution(toolCall llm.ToolCall, result string, err error)

type PromptOptions

type PromptOptions struct {
	CoreChatOptions
	EnableTools        bool
	ToolApprovalPolicy string // Policy for tool call approval: "auto-approve", "always-deny", "smart"
}

PromptOptions contains configuration for handling direct prompts. It extends CoreChatOptions with additional settings specific to direct prompt handling.

type StreamHandler

type StreamHandler interface {
	OnChunk(chunk *llm.ChatResponseChunk)
	OnComplete(fullContent string, toolCalls []llm.ToolCall, err error)
	OnToolExecution(toolCall llm.ToolCall, result string, err error)
}

StreamHandler is an interface for handling streaming chat events. Implementations receive notifications about chunks of content as they arrive, when the stream completes, and when tools are executed.

type StreamOptions

type StreamOptions struct {
	CoreChatOptions
	Handler              StreamHandler
	ToolApprovalPolicy   string                            // Policy for tool call approval: "auto-approve", "always-deny", "smart"
	HumanApprovalHandler toolapproval.HumanApprovalHandler // Optional handler for human approval
}

StreamOptions contains options for streaming chat responses. It extends CoreChatOptions with a handler for processing streaming events.