mirror of
https://github.com/sst/opencode.git
synced 2025-07-07 16:14:59 +00:00
805 lines
24 KiB
Go
805 lines
24 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"log/slog"
|
|
|
|
"github.com/sst/opencode/internal/config"
|
|
"github.com/sst/opencode/internal/llm/models"
|
|
"github.com/sst/opencode/internal/llm/prompt"
|
|
"github.com/sst/opencode/internal/llm/provider"
|
|
"github.com/sst/opencode/internal/llm/tools"
|
|
"github.com/sst/opencode/internal/logging"
|
|
"github.com/sst/opencode/internal/message"
|
|
"github.com/sst/opencode/internal/permission"
|
|
"github.com/sst/opencode/internal/session"
|
|
"github.com/sst/opencode/internal/status"
|
|
)
|
|
|
|
// Common errors
|
|
var (
|
|
ErrRequestCancelled = errors.New("request cancelled by user")
|
|
ErrSessionBusy = errors.New("session is currently processing another request")
|
|
)
|
|
|
|
type AgentEvent struct {
|
|
message message.Message
|
|
err error
|
|
}
|
|
|
|
func (e *AgentEvent) Err() error {
|
|
return e.err
|
|
}
|
|
|
|
func (e *AgentEvent) Response() message.Message {
|
|
return e.message
|
|
}
|
|
|
|
type Service interface {
|
|
Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error)
|
|
Cancel(sessionID string)
|
|
IsSessionBusy(sessionID string) bool
|
|
IsBusy() bool
|
|
Update(agentName config.AgentName, modelID models.ModelID) (models.Model, error)
|
|
CompactSession(ctx context.Context, sessionID string, force bool) error
|
|
GetUsage(ctx context.Context, sessionID string) (*int64, error)
|
|
EstimateContextWindowUsage(ctx context.Context, sessionID string) (float64, bool, error)
|
|
}
|
|
|
|
type agent struct {
|
|
sessions session.Service
|
|
messages message.Service
|
|
|
|
tools []tools.BaseTool
|
|
provider provider.Provider
|
|
|
|
titleProvider provider.Provider
|
|
|
|
activeRequests sync.Map
|
|
}
|
|
|
|
func NewAgent(
|
|
agentName config.AgentName,
|
|
sessions session.Service,
|
|
messages message.Service,
|
|
agentTools []tools.BaseTool,
|
|
) (Service, error) {
|
|
agentProvider, err := createAgentProvider(agentName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var titleProvider provider.Provider
|
|
// Only generate titles for the primary agent
|
|
if agentName == config.AgentPrimary {
|
|
titleProvider, err = createAgentProvider(config.AgentTitle)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
agent := &agent{
|
|
provider: agentProvider,
|
|
messages: messages,
|
|
sessions: sessions,
|
|
tools: agentTools,
|
|
titleProvider: titleProvider,
|
|
activeRequests: sync.Map{},
|
|
}
|
|
|
|
return agent, nil
|
|
}
|
|
|
|
func (a *agent) Cancel(sessionID string) {
|
|
if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {
|
|
if cancel, ok := cancelFunc.(context.CancelFunc); ok {
|
|
status.Info(fmt.Sprintf("Request cancellation initiated for session: %s", sessionID))
|
|
cancel()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (a *agent) IsBusy() bool {
|
|
busy := false
|
|
a.activeRequests.Range(func(key, value interface{}) bool {
|
|
if cancelFunc, ok := value.(context.CancelFunc); ok {
|
|
if cancelFunc != nil {
|
|
busy = true
|
|
return false // Stop iterating
|
|
}
|
|
}
|
|
return true // Continue iterating
|
|
})
|
|
return busy
|
|
}
|
|
|
|
func (a *agent) IsSessionBusy(sessionID string) bool {
|
|
_, busy := a.activeRequests.Load(sessionID)
|
|
return busy
|
|
}
|
|
|
|
func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {
|
|
if content == "" {
|
|
return nil
|
|
}
|
|
if a.titleProvider == nil {
|
|
return nil
|
|
}
|
|
session, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
parts := []message.ContentPart{message.TextContent{Text: content}}
|
|
response, err := a.titleProvider.SendMessages(
|
|
ctx,
|
|
[]message.Message{
|
|
{
|
|
Role: message.User,
|
|
Parts: parts,
|
|
},
|
|
},
|
|
make([]tools.BaseTool, 0),
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
title := strings.TrimSpace(strings.ReplaceAll(response.Content, "\n", " "))
|
|
if title == "" {
|
|
return nil
|
|
}
|
|
|
|
session.Title = title
|
|
_, err = a.sessions.Update(ctx, session)
|
|
return err
|
|
}
|
|
|
|
func (a *agent) err(err error) AgentEvent {
|
|
return AgentEvent{
|
|
err: err,
|
|
}
|
|
}
|
|
|
|
func (a *agent) Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error) {
|
|
if !a.provider.Model().SupportsAttachments && attachments != nil {
|
|
attachments = nil
|
|
}
|
|
events := make(chan AgentEvent)
|
|
if a.IsSessionBusy(sessionID) {
|
|
return nil, ErrSessionBusy
|
|
}
|
|
|
|
genCtx, cancel := context.WithCancel(ctx)
|
|
|
|
a.activeRequests.Store(sessionID, cancel)
|
|
go func() {
|
|
slog.Debug("Request started", "sessionID", sessionID)
|
|
defer logging.RecoverPanic("agent.Run", func() {
|
|
events <- a.err(fmt.Errorf("panic while running the agent"))
|
|
})
|
|
var attachmentParts []message.ContentPart
|
|
for _, attachment := range attachments {
|
|
attachmentParts = append(attachmentParts, message.BinaryContent{Path: attachment.FilePath, MIMEType: attachment.MimeType, Data: attachment.Content})
|
|
}
|
|
result := a.processGeneration(genCtx, sessionID, content, attachmentParts)
|
|
if result.Err() != nil && !errors.Is(result.Err(), ErrRequestCancelled) && !errors.Is(result.Err(), context.Canceled) {
|
|
status.Error(result.Err().Error())
|
|
}
|
|
slog.Debug("Request completed", "sessionID", sessionID)
|
|
a.activeRequests.Delete(sessionID)
|
|
cancel()
|
|
events <- result
|
|
close(events)
|
|
}()
|
|
|
|
return events, nil
|
|
}
|
|
|
|
func (a *agent) prepareMessageHistory(ctx context.Context, sessionID string) (session.Session, []message.Message, error) {
|
|
currentSession, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return currentSession, nil, fmt.Errorf("failed to get session: %w", err)
|
|
}
|
|
|
|
var sessionMessages []message.Message
|
|
if currentSession.Summary != "" && !currentSession.SummarizedAt.IsZero() {
|
|
// If summary exists, only fetch messages after the summarization timestamp
|
|
sessionMessages, err = a.messages.ListAfter(ctx, sessionID, currentSession.SummarizedAt)
|
|
if err != nil {
|
|
return currentSession, nil, fmt.Errorf("failed to list messages after summary: %w", err)
|
|
}
|
|
} else {
|
|
// If no summary, fetch all messages
|
|
sessionMessages, err = a.messages.List(ctx, sessionID)
|
|
if err != nil {
|
|
return currentSession, nil, fmt.Errorf("failed to list messages: %w", err)
|
|
}
|
|
}
|
|
|
|
var messages []message.Message
|
|
if currentSession.Summary != "" && !currentSession.SummarizedAt.IsZero() {
|
|
// If summary exists, create a temporary message for the summary
|
|
summaryMessage := message.Message{
|
|
Role: message.Assistant,
|
|
Parts: []message.ContentPart{
|
|
message.TextContent{Text: currentSession.Summary},
|
|
},
|
|
}
|
|
// Start with the summary, then add messages after the summary timestamp
|
|
messages = append([]message.Message{summaryMessage}, sessionMessages...)
|
|
} else {
|
|
// If no summary, just use all messages
|
|
messages = sessionMessages
|
|
}
|
|
|
|
return currentSession, messages, nil
|
|
}
|
|
|
|
func (a *agent) triggerTitleGeneration(sessionID string, content string) {
|
|
go func() {
|
|
defer logging.RecoverPanic("agent.Run", func() {
|
|
status.Error("panic while generating title")
|
|
})
|
|
titleErr := a.generateTitle(context.Background(), sessionID, content)
|
|
if titleErr != nil {
|
|
status.Error(fmt.Sprintf("failed to generate title: %v", titleErr))
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (a *agent) processGeneration(ctx context.Context, sessionID, content string, attachmentParts []message.ContentPart) AgentEvent {
|
|
currentSession, sessionMessages, err := a.prepareMessageHistory(ctx, sessionID)
|
|
if err != nil {
|
|
return a.err(err)
|
|
}
|
|
|
|
// If this is a new session, start title generation asynchronously
|
|
if len(sessionMessages) == 0 && currentSession.Summary == "" {
|
|
a.triggerTitleGeneration(sessionID, content)
|
|
}
|
|
|
|
userMsg, err := a.createUserMessage(ctx, sessionID, content, attachmentParts)
|
|
if err != nil {
|
|
return a.err(fmt.Errorf("failed to create user message: %w", err))
|
|
}
|
|
|
|
messages := append(sessionMessages, userMsg)
|
|
|
|
for {
|
|
// Check for cancellation before each iteration
|
|
select {
|
|
case <-ctx.Done():
|
|
return a.err(ctx.Err())
|
|
default:
|
|
// Continue processing
|
|
}
|
|
|
|
// Check if auto-compaction is needed before calling the provider
|
|
usagePercentage, needsCompaction, errEstimate := a.EstimateContextWindowUsage(ctx, sessionID)
|
|
if errEstimate != nil {
|
|
slog.Warn("Failed to estimate context window usage for auto-compaction", "error", errEstimate, "sessionID", sessionID)
|
|
} else if needsCompaction {
|
|
status.Info(fmt.Sprintf("Context window usage is at %.2f%%. Auto-compacting conversation...", usagePercentage))
|
|
|
|
// Run compaction synchronously
|
|
compactCtx, cancelCompact := context.WithTimeout(ctx, 30*time.Second) // Use appropriate context
|
|
errCompact := a.CompactSession(compactCtx, sessionID, true)
|
|
cancelCompact()
|
|
|
|
if errCompact != nil {
|
|
status.Warn(fmt.Sprintf("Auto-compaction failed: %v. Context window usage may continue to grow.", errCompact))
|
|
} else {
|
|
status.Info("Auto-compaction completed successfully.")
|
|
// After compaction, message history needs to be re-prepared.
|
|
// The 'messages' slice needs to be updated with the new summary and subsequent messages,
|
|
// ensuring the latest user message is correctly appended.
|
|
_, sessionMessagesFromCompact, errPrepare := a.prepareMessageHistory(ctx, sessionID)
|
|
if errPrepare != nil {
|
|
return a.err(fmt.Errorf("failed to re-prepare message history after compaction: %w", errPrepare))
|
|
}
|
|
messages = sessionMessagesFromCompact
|
|
|
|
// Ensure the user message that triggered this cycle is the last one.
|
|
// 'userMsg' was created before this loop using a.createUserMessage.
|
|
// It should be appended to the 'messages' slice if it's not already the last element.
|
|
if len(messages) == 0 || (len(messages) > 0 && messages[len(messages)-1].ID != userMsg.ID) {
|
|
messages = append(messages, userMsg)
|
|
}
|
|
}
|
|
}
|
|
|
|
agentMessage, toolResults, err := a.streamAndHandleEvents(ctx, sessionID, messages)
|
|
if err != nil {
|
|
if errors.Is(err, context.Canceled) {
|
|
agentMessage.AddFinish(message.FinishReasonCanceled)
|
|
a.messages.Update(context.Background(), agentMessage)
|
|
return a.err(ErrRequestCancelled)
|
|
}
|
|
return a.err(fmt.Errorf("failed to process events: %w", err))
|
|
}
|
|
slog.Info("Result", "message", agentMessage.FinishReason(), "toolResults", toolResults)
|
|
if (agentMessage.FinishReason() == message.FinishReasonToolUse) && toolResults != nil {
|
|
// We are not done, we need to respond with the tool response
|
|
messages = append(messages, agentMessage, *toolResults)
|
|
continue
|
|
}
|
|
return AgentEvent{
|
|
message: agentMessage,
|
|
}
|
|
}
|
|
}
|
|
|
|
func (a *agent) createUserMessage(ctx context.Context, sessionID, content string, attachmentParts []message.ContentPart) (message.Message, error) {
|
|
parts := []message.ContentPart{message.TextContent{Text: content}}
|
|
parts = append(parts, attachmentParts...)
|
|
return a.messages.Create(ctx, sessionID, message.CreateMessageParams{
|
|
Role: message.User,
|
|
Parts: parts,
|
|
})
|
|
}
|
|
|
|
func (a *agent) createToolResponseMessage(ctx context.Context, sessionID string, toolResults []message.ToolResult) (*message.Message, error) {
|
|
if len(toolResults) == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
parts := make([]message.ContentPart, 0, len(toolResults))
|
|
for _, tr := range toolResults {
|
|
parts = append(parts, tr)
|
|
}
|
|
|
|
msg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
|
|
Role: message.Tool,
|
|
Parts: parts,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create tool response message: %w", err)
|
|
}
|
|
|
|
return &msg, nil
|
|
}
|
|
|
|
func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {
|
|
eventChan := a.provider.StreamResponse(ctx, msgHistory, a.tools)
|
|
|
|
assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
|
|
Role: message.Assistant,
|
|
Parts: []message.ContentPart{},
|
|
Model: a.provider.Model().ID,
|
|
})
|
|
if err != nil {
|
|
return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)
|
|
}
|
|
|
|
// Add the session and message ID into the context if needed by tools.
|
|
ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
|
|
ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
|
|
|
|
// Process each event in the stream.
|
|
for event := range eventChan {
|
|
if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
|
|
a.finishMessage(ctx, &assistantMsg, message.FinishReasonCanceled)
|
|
return assistantMsg, nil, processErr
|
|
}
|
|
if ctx.Err() != nil {
|
|
a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled)
|
|
return assistantMsg, nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
// If the assistant wants to use tools, execute them
|
|
if assistantMsg.FinishReason() == message.FinishReasonToolUse {
|
|
toolCalls := assistantMsg.ToolCalls()
|
|
if len(toolCalls) > 0 {
|
|
toolResults, err := a.executeToolCalls(ctx, toolCalls)
|
|
if err != nil {
|
|
if errors.Is(err, context.Canceled) {
|
|
a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled)
|
|
}
|
|
return assistantMsg, nil, err
|
|
}
|
|
|
|
// Create a message with the tool results
|
|
toolResponseMsg, err := a.createToolResponseMessage(ctx, sessionID, toolResults)
|
|
if err != nil {
|
|
return assistantMsg, nil, err
|
|
}
|
|
|
|
return assistantMsg, toolResponseMsg, nil
|
|
}
|
|
}
|
|
|
|
return assistantMsg, nil, nil
|
|
}
|
|
|
|
func (a *agent) executeToolCalls(ctx context.Context, toolCalls []message.ToolCall) ([]message.ToolResult, error) {
|
|
toolResults := make([]message.ToolResult, len(toolCalls))
|
|
|
|
for i, toolCall := range toolCalls {
|
|
select {
|
|
case <-ctx.Done():
|
|
// Make all future tool calls cancelled
|
|
for j := i; j < len(toolCalls); j++ {
|
|
toolResults[j] = message.ToolResult{
|
|
ToolCallID: toolCalls[j].ID,
|
|
Content: "Tool execution canceled by user",
|
|
IsError: true,
|
|
}
|
|
}
|
|
return toolResults, ctx.Err()
|
|
default:
|
|
// Continue processing
|
|
var tool tools.BaseTool
|
|
for _, availableTools := range a.tools {
|
|
if availableTools.Info().Name == toolCall.Name {
|
|
tool = availableTools
|
|
}
|
|
}
|
|
|
|
// Tool not found
|
|
if tool == nil {
|
|
toolResults[i] = message.ToolResult{
|
|
ToolCallID: toolCall.ID,
|
|
Content: fmt.Sprintf("Tool not found: %s", toolCall.Name),
|
|
IsError: true,
|
|
}
|
|
continue
|
|
}
|
|
|
|
toolResult, toolErr := tool.Run(ctx, tools.ToolCall{
|
|
ID: toolCall.ID,
|
|
Name: toolCall.Name,
|
|
Input: toolCall.Input,
|
|
})
|
|
|
|
if toolErr != nil {
|
|
if errors.Is(toolErr, permission.ErrorPermissionDenied) {
|
|
toolResults[i] = message.ToolResult{
|
|
ToolCallID: toolCall.ID,
|
|
Content: "Permission denied",
|
|
IsError: true,
|
|
}
|
|
// Cancel all remaining tool calls if permission is denied
|
|
for j := i + 1; j < len(toolCalls); j++ {
|
|
toolResults[j] = message.ToolResult{
|
|
ToolCallID: toolCalls[j].ID,
|
|
Content: "Tool execution canceled by user",
|
|
IsError: true,
|
|
}
|
|
}
|
|
return toolResults, nil
|
|
}
|
|
|
|
// Handle other errors
|
|
toolResults[i] = message.ToolResult{
|
|
ToolCallID: toolCall.ID,
|
|
Content: toolErr.Error(),
|
|
IsError: true,
|
|
}
|
|
continue
|
|
}
|
|
|
|
toolResults[i] = message.ToolResult{
|
|
ToolCallID: toolCall.ID,
|
|
Content: toolResult.Content,
|
|
Metadata: toolResult.Metadata,
|
|
IsError: toolResult.IsError,
|
|
}
|
|
}
|
|
}
|
|
|
|
return toolResults, nil
|
|
}
|
|
|
|
func (a *agent) finishMessage(ctx context.Context, msg *message.Message, finishReson message.FinishReason) {
|
|
msg.AddFinish(finishReson)
|
|
_, _ = a.messages.Update(ctx, *msg)
|
|
}
|
|
|
|
func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
// Continue processing
|
|
}
|
|
|
|
switch event.Type {
|
|
case provider.EventThinkingDelta:
|
|
assistantMsg.AppendReasoningContent(event.Content)
|
|
_, err := a.messages.Update(ctx, *assistantMsg)
|
|
return err
|
|
case provider.EventContentDelta:
|
|
assistantMsg.AppendContent(event.Content)
|
|
_, err := a.messages.Update(ctx, *assistantMsg)
|
|
return err
|
|
case provider.EventToolUseStart:
|
|
assistantMsg.AddToolCall(*event.ToolCall)
|
|
_, err := a.messages.Update(ctx, *assistantMsg)
|
|
return err
|
|
case provider.EventToolUseStop:
|
|
assistantMsg.FinishToolCall(event.ToolCall.ID)
|
|
_, err := a.messages.Update(ctx, *assistantMsg)
|
|
return err
|
|
case provider.EventError:
|
|
if errors.Is(event.Error, context.Canceled) {
|
|
status.Info(fmt.Sprintf("Event processing canceled for session: %s", sessionID))
|
|
return context.Canceled
|
|
}
|
|
status.Error(event.Error.Error())
|
|
return event.Error
|
|
case provider.EventComplete:
|
|
assistantMsg.SetToolCalls(event.Response.ToolCalls)
|
|
assistantMsg.AddFinish(event.Response.FinishReason)
|
|
if _, err := a.messages.Update(ctx, *assistantMsg); err != nil {
|
|
return fmt.Errorf("failed to update message: %w", err)
|
|
}
|
|
return a.TrackUsage(ctx, sessionID, a.provider.Model(), event.Response.Usage)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (a *agent) GetUsage(ctx context.Context, sessionID string) (*int64, error) {
|
|
session, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get session: %w", err)
|
|
}
|
|
|
|
usage := session.PromptTokens + session.CompletionTokens
|
|
return &usage, nil
|
|
}
|
|
|
|
func (a *agent) EstimateContextWindowUsage(ctx context.Context, sessionID string) (float64, bool, error) {
|
|
session, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return 0, false, fmt.Errorf("failed to get session: %w", err)
|
|
}
|
|
|
|
// Get the model's context window size
|
|
model := a.provider.Model()
|
|
contextWindow := model.ContextWindow
|
|
if contextWindow <= 0 {
|
|
// Default to a reasonable size if not specified
|
|
contextWindow = 100000
|
|
}
|
|
|
|
// Calculate current token usage
|
|
currentTokens := session.PromptTokens + session.CompletionTokens
|
|
|
|
// Get the max tokens setting for the agent
|
|
maxTokens := a.provider.MaxTokens()
|
|
|
|
// Calculate percentage of context window used
|
|
usagePercentage := float64(currentTokens) / float64(contextWindow)
|
|
|
|
// Check if we need to auto-compact
|
|
// Auto-compact when:
|
|
// 1. Usage exceeds 90% of context window, OR
|
|
// 2. Current usage + maxTokens would exceed 100% of context window
|
|
needsCompaction := usagePercentage >= 0.9 ||
|
|
float64(currentTokens+maxTokens) > float64(contextWindow)
|
|
|
|
return usagePercentage * 100, needsCompaction, nil
|
|
}
|
|
|
|
func (a *agent) TrackUsage(ctx context.Context, sessionID string, model models.Model, usage provider.TokenUsage) error {
|
|
sess, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get session: %w", err)
|
|
}
|
|
|
|
cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
|
|
model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
|
|
model.CostPer1MIn/1e6*float64(usage.InputTokens) +
|
|
model.CostPer1MOut/1e6*float64(usage.OutputTokens)
|
|
|
|
sess.Cost += cost
|
|
sess.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
|
|
sess.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
|
|
|
|
_, err = a.sessions.Update(ctx, sess)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to save session: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *agent) Update(agentName config.AgentName, modelID models.ModelID) (models.Model, error) {
|
|
if a.IsBusy() {
|
|
return models.Model{}, fmt.Errorf("cannot change model while processing requests")
|
|
}
|
|
|
|
if err := config.UpdateAgentModel(agentName, modelID); err != nil {
|
|
return models.Model{}, fmt.Errorf("failed to update config: %w", err)
|
|
}
|
|
|
|
provider, err := createAgentProvider(agentName)
|
|
if err != nil {
|
|
return models.Model{}, fmt.Errorf("failed to create provider for model %s: %w", modelID, err)
|
|
}
|
|
|
|
a.provider = provider
|
|
|
|
return a.provider.Model(), nil
|
|
}
|
|
|
|
func (a *agent) CompactSession(ctx context.Context, sessionID string, force bool) error {
|
|
// Check if the session is busy
|
|
if a.IsSessionBusy(sessionID) && !force {
|
|
return ErrSessionBusy
|
|
}
|
|
|
|
// Create a cancellable context
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
|
|
// Mark the session as busy during compaction
|
|
compactionCancelFunc := func() {}
|
|
a.activeRequests.Store(sessionID+"-compact", compactionCancelFunc)
|
|
defer a.activeRequests.Delete(sessionID + "-compact")
|
|
|
|
// Fetch the session
|
|
session, err := a.sessions.Get(ctx, sessionID)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get session: %w", err)
|
|
}
|
|
|
|
// Fetch all messages for the session
|
|
sessionMessages, err := a.messages.List(ctx, sessionID)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to list messages: %w", err)
|
|
}
|
|
|
|
var existingSummary string
|
|
if session.Summary != "" && !session.SummarizedAt.IsZero() {
|
|
// Filter messages that were created after the last summarization
|
|
var newMessages []message.Message
|
|
for _, msg := range sessionMessages {
|
|
if msg.CreatedAt.After(session.SummarizedAt) {
|
|
newMessages = append(newMessages, msg)
|
|
}
|
|
}
|
|
sessionMessages = newMessages
|
|
existingSummary = session.Summary
|
|
}
|
|
|
|
// If there are no messages to summarize and no existing summary, return early
|
|
if len(sessionMessages) == 0 && existingSummary == "" {
|
|
return nil
|
|
}
|
|
|
|
messages := []message.Message{
|
|
message.Message{
|
|
Role: message.System,
|
|
Parts: []message.ContentPart{
|
|
message.TextContent{
|
|
Text: `You are a helpful AI assistant tasked with summarizing conversations.
|
|
|
|
When asked to summarize, provide a detailed but concise summary of the conversation.
|
|
Focus on information that would be helpful for continuing the conversation, including:
|
|
- What was done
|
|
- What is currently being worked on
|
|
- Which files are being modified
|
|
- What needs to be done next
|
|
|
|
Your summary should be comprehensive enough to provide context but concise enough to be quickly understood.`,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// If there's an existing summary, include it
|
|
if existingSummary != "" {
|
|
messages = append(messages, message.Message{
|
|
Role: message.Assistant,
|
|
Parts: []message.ContentPart{
|
|
message.TextContent{
|
|
Text: existingSummary,
|
|
},
|
|
},
|
|
})
|
|
}
|
|
|
|
// Add all messages since the last summarized message
|
|
messages = append(messages, sessionMessages...)
|
|
|
|
// Add a final user message requesting the summary
|
|
messages = append(messages, message.Message{
|
|
Role: message.User,
|
|
Parts: []message.ContentPart{
|
|
message.TextContent{
|
|
Text: "Provide a detailed but concise summary of our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next.",
|
|
},
|
|
},
|
|
})
|
|
|
|
// Call provider to get the summary
|
|
response, err := a.provider.SendMessages(ctx, messages, a.tools)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get summary from the assistant: %w", err)
|
|
}
|
|
|
|
// Extract the summary text
|
|
summaryText := strings.TrimSpace(response.Content)
|
|
if summaryText == "" {
|
|
return fmt.Errorf("received empty summary from the assistant")
|
|
}
|
|
|
|
// Update the session with the new summary
|
|
session.Summary = summaryText
|
|
session.SummarizedAt = time.Now()
|
|
|
|
// Save the updated session
|
|
_, err = a.sessions.Update(ctx, session)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to save session with summary: %w", err)
|
|
}
|
|
|
|
// Track token usage
|
|
err = a.TrackUsage(ctx, sessionID, a.provider.Model(), response.Usage)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to track usage: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func createAgentProvider(agentName config.AgentName) (provider.Provider, error) {
|
|
cfg := config.Get()
|
|
agentConfig, ok := cfg.Agents[agentName]
|
|
if !ok {
|
|
return nil, fmt.Errorf("agent %s not found", agentName)
|
|
}
|
|
model, ok := models.SupportedModels[agentConfig.Model]
|
|
if !ok {
|
|
return nil, fmt.Errorf("model %s not supported", agentConfig.Model)
|
|
}
|
|
|
|
providerCfg, ok := cfg.Providers[model.Provider]
|
|
if !ok {
|
|
return nil, fmt.Errorf("provider %s not supported", model.Provider)
|
|
}
|
|
if providerCfg.Disabled {
|
|
return nil, fmt.Errorf("provider %s is not enabled", model.Provider)
|
|
}
|
|
maxTokens := model.DefaultMaxTokens
|
|
if agentConfig.MaxTokens > 0 {
|
|
maxTokens = agentConfig.MaxTokens
|
|
}
|
|
opts := []provider.ProviderClientOption{
|
|
provider.WithAPIKey(providerCfg.APIKey),
|
|
provider.WithModel(model),
|
|
provider.WithSystemMessage(prompt.GetAgentPrompt(agentName, model.Provider)),
|
|
provider.WithMaxTokens(maxTokens),
|
|
}
|
|
if model.Provider == models.ProviderOpenAI && model.CanReason {
|
|
opts = append(
|
|
opts,
|
|
provider.WithOpenAIOptions(
|
|
provider.WithReasoningEffort(agentConfig.ReasoningEffort),
|
|
),
|
|
)
|
|
} else if model.Provider == models.ProviderAnthropic && model.CanReason && agentName == config.AgentPrimary {
|
|
opts = append(
|
|
opts,
|
|
provider.WithAnthropicOptions(
|
|
provider.WithAnthropicShouldThinkFn(provider.DefaultShouldThinkFn),
|
|
),
|
|
)
|
|
}
|
|
agentProvider, err := provider.NewProvider(
|
|
model.Provider,
|
|
opts...,
|
|
)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("could not create provider: %v", err)
|
|
}
|
|
|
|
return agentProvider, nil
|
|
}
|