GoLoom/internal/model/types.go

package model

import (
	"context"
	"sync"
)

// ============================================================
// Chat 数据类型
// ============================================================

// ChatRole 消息角色
type ChatRole string

const (
	ChatRoleSystem    ChatRole = "system"
	ChatRoleUser      ChatRole = "user"
	ChatRoleAssistant ChatRole = "assistant"
	ChatRoleTool      ChatRole = "tool"
)

// ChatMessage 聊天消息
type ChatMessage struct {
	Role       ChatRole
	Content    string
	ToolCallID string
	Name       string
	ToolCalls  []ChatToolCall
}

// ChatToolCall 工具调用请求
type ChatToolCall struct {
	ID        string
	Name      string
	Arguments string
}

// ChatReply 聊天回复
type ChatReply struct {
	Content   string
	ToolCalls []ChatToolCall
}

// ChatStreamEvent 流式事件
type ChatStreamEvent struct {
	Delta     string
	ToolCalls []ChatToolCall
	Done      bool
}

// ChatContent 聊天输入内容
type ChatContent struct {
	Texts []TextPart
}

// TextPart 文本片段
type TextPart struct {
	Message string
}

// ============================================================
// 核心接口
// ============================================================

// Tool 外部工具接口
type Tool interface {
	Name() string
	Description() string
	Call(ctx context.Context, input string) (string, error)
}

// ChatModel 聊天模型接口，支持同步生成和流式输出
type ChatModel interface {
	Generate(ctx context.Context, messages []ChatMessage) (ChatReply, error)
	Stream(ctx context.Context, messages []ChatMessage) (<-chan ChatStreamEvent, <-chan error)
}

// Agent 智能体接口
type Agent interface {
	Name() string
	Run(ctx context.Context, content ChatContent) (string, error)
	Stream(ctx context.Context, content ChatContent, out chan<- string) error
}

// Runner 运行器接口，管理会话并执行 Agent
type Runner interface {
	CreateSession(userID string) (string, error)
	Run(userID, sessionID string, content ChatContent) ([]string, error)
	Stream(userID, sessionID string, content ChatContent) (<-chan string, <-chan error)
}

// ============================================================
// 注册与存储
// ============================================================

// RegisteredAgent 已注册的 Agent 信息
type RegisteredAgent struct {
	AppName   string
	AgentID   string
	AgentName string
	AgentDesc string
	Runner    Runner
}

// AgentRegistry Agent 注册表接口
type AgentRegistry interface {
	Register(agent RegisteredAgent) error
	Get(agentID string) (RegisteredAgent, bool)
	List() []RegisteredAgent
}

// SessionStore 会话存储接口
type SessionStore interface {
	Get(userID, agentID string) (string, bool)
	Set(userID, agentID, sessionID string) error
}

// ============================================================
// 内存实现
// ============================================================

// InMemoryAgentRegistry 基于内存的 Agent 注册表
type InMemoryAgentRegistry struct {
	mu     sync.RWMutex
	agents map[string]RegisteredAgent
}

func NewInMemoryAgentRegistry() *InMemoryAgentRegistry {
	return &InMemoryAgentRegistry{agents: make(map[string]RegisteredAgent)}
}

func (r *InMemoryAgentRegistry) Register(agent RegisteredAgent) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.agents[agent.AgentID] = agent
	return nil
}

func (r *InMemoryAgentRegistry) Get(agentID string) (RegisteredAgent, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	agent, ok := r.agents[agentID]
	return agent, ok
}

func (r *InMemoryAgentRegistry) List() []RegisteredAgent {
	r.mu.RLock()
	defer r.mu.RUnlock()
	agents := make([]RegisteredAgent, 0, len(r.agents))
	for _, agent := range r.agents {
		agents = append(agents, agent)
	}
	return agents
}

// InMemorySessionStore 基于内存的会话存储
type InMemorySessionStore struct {
	mu       sync.RWMutex
	sessions map[string]string
}

func NewInMemorySessionStore() *InMemorySessionStore {
	return &InMemorySessionStore{sessions: make(map[string]string)}
}

func (s *InMemorySessionStore) Get(userID, agentID string) (string, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sessionID, ok := s.sessions[userID+":"+agentID]
	return sessionID, ok
}

func (s *InMemorySessionStore) Set(userID, agentID, sessionID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[userID+":"+agentID] = sessionID
	return nil
}