DivaCode/internal/agent/agent.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"git.db123.ir/db123/divacode/internal/companion"
	"git.db123.ir/db123/divacode/internal/config"
	"git.db123.ir/db123/divacode/internal/llama"
	"git.db123.ir/db123/divacode/internal/memory"
	"git.db123.ir/db123/divacode/internal/storage"
	"git.db123.ir/db123/divacode/internal/tools"
	"git.db123.ir/db123/divacode/pkg/types"
)

type Agent struct {
	mu        sync.Mutex
	cfg       *config.Config
	llm       *llama.Client
	mem       *memory.Store
	store     *storage.DB
	companion *companion.Engine
	tools     *tools.Registry

	conversationID string
	messages       []types.Message
}

func New(cfg *config.Config, llm *llama.Client, store *storage.DB, mem *memory.Store, comp *companion.Engine, toolReg *tools.Registry) *Agent {
	return &Agent{
		cfg:       cfg,
		llm:       llm,
		store:     store,
		mem:       mem,
		companion: comp,
		tools:     toolReg,
	}
}

func (a *Agent) NewConversation() error {
	a.mu.Lock()
	defer a.mu.Unlock()
	id := fmt.Sprintf("%x", time.Now().UnixNano())
	a.conversationID = id
	a.messages = nil
	if err := a.store.Exec("INSERT INTO conversations (id) VALUES (?)", id); err != nil {
		return fmt.Errorf("new conversation: %w", err)
	}
	slog.Info("new conversation", "id", id)
	return nil
}

func (a *Agent) HandleMessage(ctx context.Context, content string) (string, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	userMsg := types.Message{
		ID:        fmt.Sprintf("%x", time.Now().UnixNano()),
		Role:      types.RoleUser,
		Content:   content,
		CreatedAt: time.Now(),
	}
	a.messages = append(a.messages, userMsg)
	a.saveMessage(userMsg)

	memories, err := a.mem.Search(ctx, content, a.cfg.MemoryTopK)
	if err != nil {
		slog.Warn("memory search failed", "error", err)
	}

	buildResult, err := a.companion.BuildPrompt(ctx, &companion.PromptInput{
		UserMessage:    content,
		RecentMessages: a.messages,
		Memories:       memories,
	})
	if err != nil {
		return "", fmt.Errorf("build prompt: %w", err)
	}

	segments := []types.PromptSegment{
		{Type: types.SegSystem, Content: buildResult.SystemContent},
	}
	start := 0
	if len(a.messages) > 10 {
		start = len(a.messages) - 10
	}
	for _, msg := range a.messages[start:] {
		segType := types.SegUser
		if msg.Role == types.RoleAssistant {
			segType = types.SegAssistant
		}
		segments = append(segments, types.PromptSegment{Type: segType, Content: msg.Content})
	}

	tmpl := llama.NewTemplateEngine(a.cfg.ChatTemplate)
	prompt := tmpl.Build(segments)

	resp, err := a.llm.Complete(&llama.CompletionRequest{
		Prompt:      prompt,
		Temperature: a.cfg.Temperature,
		TopP:        a.cfg.TopP,
		MaxTokens:   a.cfg.MaxTokens,
		Stop:        tmpl.StopTokens(),
		Stream:      false,
		CachePrompt: true,
	})
	if err != nil {
		return "", fmt.Errorf("llm complete: %w", err)
	}

	parsed := ParseResponse(resp.Content)

	assistantMsg := types.Message{
		ID:        fmt.Sprintf("%x", time.Now().UnixNano()),
		Role:      types.RoleAssistant,
		Content:   parsed.Content,
		Reasoning: parsed.Reasoning,
		CreatedAt: time.Now(),
	}
	a.messages = append(a.messages, assistantMsg)
	a.saveMessage(assistantMsg)

	go func() {
		if err := a.mem.Store(context.Background(), userMsg.Content, "conversation"); err != nil {
			slog.Warn("memory store failed", "error", err)
		}
	}()

	return parsed.Content, nil
}

func (a *Agent) Messages() []types.Message {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.messages
}

func (a *Agent) ConversationID() string {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.conversationID
}

func (a *Agent) saveMessage(msg types.Message) {
	toolCallJSON := ""
	if msg.ToolCall != nil {
		b, _ := json.Marshal(msg.ToolCall)
		toolCallJSON = string(b)
	}
	err := a.store.Exec(
		"INSERT INTO messages (id, conversation_id, role, content, reasoning, tool_call, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
		msg.ID, a.conversationID, string(msg.Role), msg.Content, msg.Reasoning, toolCallJSON, msg.CreatedAt.Format(time.RFC3339),
	)
	if err != nil {
		slog.Warn("failed to save message", "error", err)
	}
}