140 lines
3.4 KiB
Go
140 lines
3.4 KiB
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)
|
|
}
|
|
|
|
prompt, err := a.companion.BuildPrompt(ctx, &companion.PromptInput{
|
|
UserMessage: content,
|
|
RecentMessages: a.messages,
|
|
Memories: memories,
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("build prompt: %w", err)
|
|
}
|
|
|
|
resp, err := a.llm.Complete(&llama.CompletionRequest{
|
|
Prompt: prompt,
|
|
Temperature: a.cfg.Temperature,
|
|
TopP: a.cfg.TopP,
|
|
MaxTokens: a.cfg.MaxTokens,
|
|
Stream: false,
|
|
CachePrompt: true,
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("llm complete: %w", err)
|
|
}
|
|
|
|
assistantMsg := types.Message{
|
|
ID: fmt.Sprintf("%x", time.Now().UnixNano()),
|
|
Role: types.RoleAssistant,
|
|
Content: resp.Content,
|
|
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 resp.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, tool_call, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
|
msg.ID, a.conversationID, string(msg.Role), msg.Content, toolCallJSON, msg.CreatedAt.Format(time.RFC3339),
|
|
)
|
|
if err != nil {
|
|
slog.Warn("failed to save message", "error", err)
|
|
}
|
|
}
|