agent

package
v0.0.0-...-5ed4770 Latest Latest
Warning

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

Go to latest
Published: Jun 30, 2026 License: MIT Imports: 43 Imported by: 0

Documentation

Overview

Package agent provides the AI agent implementation for Sequify, including system prompt loading, ChatModel creation, and tool definitions.

Package agent provides the agent layer for LLM-driven intent processing.

Package agent provides the AI agent implementation for Sequify, including system prompt loading, ChatModel creation, and tool definitions.

Package agent provides the AI agent implementation for Sequify. This file implements the eino callback handler for statistics collection.

Index

Constants

This section is empty.

Variables

View Source
var (
	// ErrJSONParse indicates a JSON parsing error during streaming.
	ErrJSONParse = errors.New("json parse error")

	// ErrTransientNetwork indicates a transient network error during streaming.
	ErrTransientNetwork = errors.New("transient network error")
)

Streaming error types for improved error classification. These allow type-safe error handling instead of string matching.

View Source
var ErrConnectionNotFound = errors.New("connection not found")

ErrConnectionNotFound indicates that a connection lookup failed. Used by capability tools for device failover logic.

View Source
var ErrInterrupted = errors.New("agent interrupted")

ErrInterrupted is returned by Agent.Query when the agent is interrupted (e.g., askUserQuestion). The caller should handle the interrupt and start the typing loop. This is NOT a framework error.

View Source
var ErrPendingResultOwnership = fmt.Errorf("result submitted by unauthorized connection")

ErrPendingResultOwnership is returned when a result is submitted by a different connection than the one that registered it.

Functions

func ExtractRetryAfter

func ExtractRetryAfter(err error) time.Duration

ExtractRetryAfter extracts the Retry-After duration from a rate limit error. Returns 0 if the error does not carry retry-after information.

func IsRateLimitError

func IsRateLimitError(err error) bool

IsRateLimitError returns true if the error indicates a rate limit condition. It checks for the rateLimitDetector interface or common HTTP 429 patterns.

func NewCapabilityTools

func NewCapabilityTools(connLookup ConnLookup, reqSender ClientRequestSender, toolTimeout time.Duration, opts ...CapabilityOption) []tool.InvokableTool

NewCapabilityTools creates the two capability tools (request_capabilities + call_capability). The connLookup is used for WebSocket communication with the client. The toolTimeout configures how long to wait for client responses.

func NewChatModel

NewChatModel creates an eino ChatModel based on the provider configuration. Supported providers:

  • openai, deepseek, moonshot, ollama: OpenAI-compatible (via BaseURL for non-OpenAI)
  • claude: Anthropic Claude
  • mock: Mock LLM for testing (cycles through canned responses)

Returns model.ToolCallingChatModel for concurrency-safe tool binding via WithTools.

func PatchToolCallsMiddleware

func PatchToolCallsMiddleware() adk.ChatModelAgentMiddleware

PatchToolCallsMiddleware returns a ChatModelAgentMiddleware that patches dangling tool calls before the model sees them.

func SystemPrompt

func SystemPrompt(conversationID string) (string, error)

SystemPrompt loads the system prompt from the embedded .md file and renders it with the given conversation ID. The current time is set to UTC ISO 8601 format. Returns the rendered prompt string or an error if loading/parsing fails.

func WithConnID

func WithConnID(ctx context.Context, connID string) context.Context

WithConnID returns a context with the connection ID set. Used by the agent runner to inject connID before tool execution.

func WithConnectionLookup

func WithConnectionLookup(cl ConnLookup) func(*IntentHandler)

WithConnectionLookup returns a functional option that attaches a ConnLookup to the IntentHandler. When set, the handler will refresh stale connIDs from the Asynq payload by looking up live connections for the conversation owner.

func WithConversationID

func WithConversationID(ctx context.Context, conversationID string) context.Context

WithConversationID returns a context with the conversation ID set.

func WithMessageID

func WithMessageID(ctx context.Context, messageID int64) context.Context

WithMessageID returns a context with the message ID set.

func WithMetrics

func WithMetrics(m *telemetry.Metrics) func(*IntentHandler)

WithMetrics returns a functional option that attaches Prometheus metrics to the IntentHandler. When set, intent counts and durations are recorded.

func WrapRecoverableStreamingError

func WrapRecoverableStreamingError(err error) error

WrapRecoverableStreamingError wraps a streaming error with appropriate typed error markers if it matches known recoverable error patterns. This enables type-safe error handling with errors.Is downstream.

Returns the original error if it doesn't match any known patterns.

Types

type Agent

type Agent struct {
	ChatModel model.ToolCallingChatModel
	Tools     []tool.InvokableTool
	Config    config.AgentConfig
	Sender    UpdateDispatcher // Update delivery abstraction
	// contains filtered or unexported fields
}

Agent holds the LLM ChatModel, tools, and Runner for processing user intents.

func NewAgent

func NewAgent(
	ctx context.Context,
	chatModel model.ToolCallingChatModel,
	tools []tool.InvokableTool,
	cfg config.AgentConfig,
	sender UpdateDispatcher,
	systemPrompt string,
	toolStore store.Store,
	stats StatisticsStore,
	opts ...Option,
) (*Agent, error)

NewAgent creates an Agent with the given ChatModel, tools, config, sender, and systemPrompt. It creates the eino ChatModelAgent + Runner internally. If chatModel is nil, the Agent is created without a Runner (for testing or deferred initialization). If toolStore is provided, tool execution results are recorded to the database via middleware. If stats is provided, LLM/token/tool statistics are recorded via eino callbacks. If metrics is provided, Prometheus metrics are recorded for LLM calls and tool executions.

func (*Agent) LastCompression

func (a *Agent) LastCompression() *CompressionInfo

LastCompression returns and clears the last compression result. Called by IntentHandler after Agent.Query returns.

func (*Agent) PendingResults

func (a *Agent) PendingResults() *PendingResults

PendingResults returns the pending tool results registry. Used by the session result handler to route client results to waiting tools.

func (*Agent) Query

func (a *Agent) Query(ctx context.Context, input RunInput) error

Query executes the agent with the given input and pushes results via Sender. It consumes Runner events, streams text chunks, and sends persistent Updates. Execution continues even if the context is cancelled (fire-and-forget). Returns error only for framework-level failures (not LLM/tool errors).

Phase 33: LLM errors are retried up to Config.LLMMaxRetries times with exponential backoff. A non-persistent "retry" Update is sent before each retry. On final failure, a persistent "error" Update is sent.

func (*Agent) QueryStats

func (a *Agent) QueryStats() (iterations, tools int64)

QueryStats returns the iteration and tool counts from the last Query call.

type CallCapInput

type CallCapInput struct {
	CapabilityName string                 `json:"capability_name"`
	Parameters     map[string]interface{} `json:"parameters"`
}

CallCapInput is the input type for the call_capability tool.

type CancelRegistry

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

CancelRegistry manages per-conversation context cancellation. It stores cancel functions keyed by conversation ID, allowing external callers (e.g., cancel handler) to abort in-progress Agent execution.

func NewCancelRegistry

func NewCancelRegistry() *CancelRegistry

NewCancelRegistry creates an empty CancelRegistry.

func (*CancelRegistry) Cancel

func (r *CancelRegistry) Cancel(convID string) bool

Cancel invokes and removes the cancel function for the given conversation ID. Returns true if a cancel func was found and invoked, false otherwise.

func (*CancelRegistry) Register

func (r *CancelRegistry) Register(convID string, cancel context.CancelFunc)

Register stores a cancel function for the given conversation ID. If a cancel func already exists for the ID, it is overwritten.

func (*CancelRegistry) Unregister

func (r *CancelRegistry) Unregister(convID string)

Unregister removes the cancel function for the given conversation ID without invoking it. Used for cleanup when Agent execution completes.

type CapabilityOption

type CapabilityOption func(*capabilityTools)

CapabilityOption configures a capabilityTools instance.

func WithToolMetrics

func WithToolMetrics(m *telemetry.Metrics) CapabilityOption

WithToolMetrics sets the Prometheus metrics for tool execution recording.

func WithUserResolver

func WithUserResolver(fn func(ctx context.Context, convID string) (string, error)) CapabilityOption

WithUserResolver sets a callback that resolves a conversation ID to its owner's user ID. Used as a fallback when the originating connection is already gone and the user ID cannot be read from it directly.

type ClientRequestSender

type ClientRequestSender interface {
	SendRequestToClient(ctx context.Context, connID, method string, payload []byte, timeout time.Duration) (*ClientResponse, error)
}

ClientRequestSender abstracts sending requests to clients and awaiting responses.

type ClientResponse

type ClientResponse struct {
	ID      string          `json:"id"`
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Payload json.RawMessage `json:"payload,omitempty"`
}

ClientResponse represents a response from a client to a server-initiated request. Mirrors websocket.ClientResponse without importing the websocket package.

type CompressionInfo

type CompressionInfo struct {
	SummaryMessage   string   // The generated summary text
	CompressedMsgIDs []uint64 // Message IDs that were compressed (for DB update)
}

CompressionInfo records the result of a summarization compression.

type ConnInfo

type ConnInfo interface {
	ConnID() string
	UserID() string
}

ConnInfo abstracts connection identity, decoupling agent from websocket.Connection.

type ConnLookup

type ConnLookup interface {
	Get(connID string) (ConnInfo, bool)
	GetByUser(userID string) []ConnInfo
}

ConnLookup abstracts connection lookup, decoupling agent from websocket.ConnectionManager.

type EditHandler

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

EditHandler processes edit_message requests by deleting messages after the edited message, inserting the new content, and re-running the agent.

func NewEditHandler

func NewEditHandler(store store.Store, agent *Agent, sender UpdateDispatcher, cancels *CancelRegistry, logger *slog.Logger) *EditHandler

NewEditHandler creates an EditHandler with the given dependencies.

func (*EditHandler) Handle

func (h *EditHandler) Handle(ctx context.Context, connID, conversationID string, messageID uint64, newContent string) error

Handle processes an edit_message request. connID is the WebSocket connection ID for sending streaming updates.

Flow:

  1. Cancel any in-progress intent processing for this conversation
  2. Validate the target message exists and belongs to the user
  3. Delete all messages after the target messageID
  4. Update the target message's content
  5. Roll back last_processed_message_id to messageID - 1
  6. Re-run the agent with the updated history

type IntentHandler

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

IntentHandler processes Asynq intent tasks by loading pending messages from the store, calling Agent.Query with a batch of user message IDs, and updating last_processed_message_id on success.

func NewIntentHandler

func NewIntentHandler(store store.Store, agent *Agent, sender UpdateDispatcher, cancels *CancelRegistry, logger *slog.Logger, stats StatisticsStore, opts ...func(*IntentHandler)) *IntentHandler

NewIntentHandler creates an IntentHandler with the given dependencies.

func (*IntentHandler) Handle

func (h *IntentHandler) Handle(ctx context.Context, payload []byte) error

Handle processes an Asynq intent task payload. It loads pending user messages, calls Agent.Query with the batch, and updates last_processed_message_id only on success. Returns nil for context.Canceled (user cancelled, no retry). Returns asynq.SkipRetry for invalid payloads.

type InterruptRegistry

type InterruptRegistry interface {
	Store(conversationID string, req interruptRequest)
	LoadAndDelete(conversationID string) (interruptRequest, bool)
	Peek(conversationID string) (interruptRequest, bool)
}

InterruptRegistry manages pending interrupt requests between the askUserQuestion tool and the interrupt middleware.

func NewInterruptRegistry

func NewInterruptRegistry() InterruptRegistry

NewInterruptRegistry creates a new default InterruptRegistry.

type Message

type Message struct {
	Role    string
	Content string
}

Message is a lightweight message for agent context (mirrors store.Message fields).

type Option

type Option func(*Agent)

Option configures an Agent.

func WithAgentMetrics

func WithAgentMetrics(m *telemetry.Metrics) Option

WithAgentMetrics sets the Prometheus metrics for the Agent. When set, LLM/token/tool metrics are recorded via eino callbacks.

type PendingResults

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

PendingResults is a thread-safe registry for pending tool executions. Tools register before sending request to client, and the result handler writes to the channel when the client responds.

func NewPendingResults

func NewPendingResults() *PendingResults

NewPendingResults creates a new PendingResults registry.

func (*PendingResults) Complete

func (pr *PendingResults) Complete(requestID string, result task.ToolResult, connID, userID string) (bool, error)

Complete writes a result to the pending tool's channel and removes it from registry. Returns false if no pending tool found for the requestID. Returns ErrPendingResultOwnership if the submitting connID/userID don't match the registered ones.

func (*PendingResults) Register

func (pr *PendingResults) Register(requestID, toolName, connID, userID string) <-chan task.ToolResult

Register adds a pending tool result and returns the result channel. The connID and userID are stored for ownership validation when the result is submitted.

type PendingToolResult

type PendingToolResult struct {
	RequestID string
	ToolName  string
	ConnID    string // Connection ID that registered this pending result
	UserID    string // User ID that registered this pending result
	Ch        chan task.ToolResult
}

PendingToolResult represents a pending tool execution waiting for client result.

type PersistentUpdate

type PersistentUpdate struct {
	Kind    string          // reply / tool_call / tool_result / error
	Payload json.RawMessage // Update payload
}

PersistentUpdate represents a persistent Update to be stored and pushed.

type ResumeHandler

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

ResumeHandler processes Asynq resume tasks by loading checkpoint data, resuming the eino Runner, and re-enqueuing skipped intents.

func NewResumeHandler

func NewResumeHandler(store store.Store, agent *Agent, taskQueue task.TaskQueue, cancels *CancelRegistry, logger *slog.Logger) *ResumeHandler

NewResumeHandler creates a ResumeHandler with the given dependencies.

func (*ResumeHandler) Handle

func (h *ResumeHandler) Handle(ctx context.Context, payload []byte) error

Handle processes an Asynq resume task payload.

type RunInput

type RunInput struct {
	ConversationID string
	ConnID         string
	Messages       []Message // conversation history (loaded by caller from Store)
	UserMessageIDs []int64   // batch of user message IDs being processed
}

RunInput holds the parameters for an Agent.Query call.

type StatisticsStore

type StatisticsStore interface {
	RecordTokenUsage(ctx context.Context, stat *store.TokenStat) error
	RecordIntentStat(ctx context.Context, stat *store.IntentStat) error
	RecordToolStat(ctx context.Context, stat *store.ToolStat) error
	RecordLLMStat(ctx context.Context, stat *store.LLMStat) error
}

StatisticsStore defines the persistence contract for usage and performance statistics. Write failures must not affect the main business flow.

type StreamingChunk

type StreamingChunk struct {
	Delta string // Incremental text since last chunk
	Done  bool   // true for the final chunk
}

StreamingChunk represents a non-persistent streaming text chunk.

type SystemPromptData

type SystemPromptData struct {
	CurrentTime    string
	ConversationID string
}

SystemPromptData holds the template variables for system prompt rendering.

type UpdateDispatcher

type UpdateDispatcher interface {
	// SendStreaming sends a non-persistent streaming chunk (seq=0).
	SendStreaming(ctx context.Context, conversationID string, chunk StreamingChunk) error

	// SendPersistent sends a persistent Update (allocated seq, written to DB).
	SendPersistent(ctx context.Context, conversationID string, update PersistentUpdate) error
}

UpdateDispatcher defines the interface for sending real-time updates (replies, streaming chunks, tool calls, errors) to connected clients. Implementations handle the write-fan-out pattern for multi-device support. Agent calls these methods; the implementation handles seq allocation and persistence.

type WillRetryError

type WillRetryError struct {
	Attempt int
	Delay   time.Duration
	Err     error
}

WillRetryError wraps an LLM error with retry metadata. It is sent as a non-persistent Update before each retry attempt.

func (*WillRetryError) Error

func (e *WillRetryError) Error() string

func (*WillRetryError) Unwrap

func (e *WillRetryError) Unwrap() error

Directories

Path Synopsis
Package mocks provides test doubles for the agent package.
Package mocks provides test doubles for the agent package.

Jump to

Keyboard shortcuts

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