mirror of
https://github.com/NoFxAiOS/nofx.git
synced 2026-07-16 09:11:03 +08:00
feat: Add Telegram AI Assistant (moltbot-nofx integration)
- Add assistant package with AI Agent runtime - agent.go: Core agent loop with tool calling - session.go: Conversation memory management - tool.go: Tool interface and base implementation - trading_tools.go: Trading-specific tools (13 tools) - prompts.go: Trading expert system prompts (EN/ZH) - Add telegram package for Telegram bot integration - bot.go: Telegram bot with rate limiting & access control - config.go: Environment-based configuration - Update main.go to initialize Telegram bot on startup - Update .env.example with new configuration options - Add gopkg.in/telebot.v3 dependency Trading tools available: - Query: get_balance, get_positions, list_traders, get_trader_status - Control: start_trader, stop_trader - Trading: get_market_price, open_long, open_short, close_position - Config: list_strategies, list_exchanges, list_ai_models
This commit is contained in:
tinklefund
350
assistant/agent.go
Normal file
350
assistant/agent.go
Normal file
@@ -0,0 +1,350 @@
|
||||
// Package assistant implements the AI Agent runtime with tool calling capabilities
|
||||
// Inspired by moltbot's agent architecture, specialized for trading
|
||||
package assistant
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"nofx/logger"
|
||||
"nofx/mcp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Agent represents an AI assistant with tool-calling capabilities
|
||||
type Agent struct {
|
||||
// AI client for LLM calls
|
||||
aiClient mcp.AIClient
|
||||
|
||||
// Tool registry
|
||||
tools map[string]Tool
|
||||
toolsLock sync.RWMutex
|
||||
|
||||
// Session/memory management
|
||||
sessions map[string]*Session
|
||||
sessionsLock sync.RWMutex
|
||||
|
||||
// Configuration
|
||||
config AgentConfig
|
||||
|
||||
// System prompt
|
||||
systemPrompt string
|
||||
}
|
||||
|
||||
// AgentConfig holds agent configuration
|
||||
type AgentConfig struct {
|
||||
// Max tool calls per turn (prevent infinite loops)
|
||||
MaxToolCalls int `json:"max_tool_calls"`
|
||||
|
||||
// Max conversation history to keep
|
||||
MaxHistoryMessages int `json:"max_history_messages"`
|
||||
|
||||
// Timeout for single AI call
|
||||
AITimeout time.Duration `json:"ai_timeout"`
|
||||
|
||||
// Model to use
|
||||
Model string `json:"model"`
|
||||
}
|
||||
|
||||
// DefaultAgentConfig returns sensible defaults
|
||||
func DefaultAgentConfig() AgentConfig {
|
||||
return AgentConfig{
|
||||
MaxToolCalls: 10,
|
||||
MaxHistoryMessages: 50,
|
||||
AITimeout: 120 * time.Second,
|
||||
Model: "deepseek-chat",
|
||||
}
|
||||
}
|
||||
|
||||
// NewAgent creates a new AI agent
|
||||
func NewAgent(aiClient mcp.AIClient, config AgentConfig) *Agent {
|
||||
agent := &Agent{
|
||||
aiClient: aiClient,
|
||||
tools: make(map[string]Tool),
|
||||
sessions: make(map[string]*Session),
|
||||
config: config,
|
||||
}
|
||||
|
||||
// Set default system prompt
|
||||
agent.systemPrompt = DefaultTradingSystemPrompt()
|
||||
|
||||
return agent
|
||||
}
|
||||
|
||||
// RegisterTool adds a tool to the agent's toolkit
|
||||
func (a *Agent) RegisterTool(tool Tool) {
|
||||
a.toolsLock.Lock()
|
||||
defer a.toolsLock.Unlock()
|
||||
a.tools[tool.Name()] = tool
|
||||
logger.Infof("🔧 Registered tool: %s", tool.Name())
|
||||
}
|
||||
|
||||
// RegisterTools adds multiple tools
|
||||
func (a *Agent) RegisterTools(tools ...Tool) {
|
||||
for _, tool := range tools {
|
||||
a.RegisterTool(tool)
|
||||
}
|
||||
}
|
||||
|
||||
// SetSystemPrompt sets the agent's system prompt
|
||||
func (a *Agent) SetSystemPrompt(prompt string) {
|
||||
a.systemPrompt = prompt
|
||||
}
|
||||
|
||||
// GetSession returns or creates a session for the given ID
|
||||
func (a *Agent) GetSession(sessionID string) *Session {
|
||||
a.sessionsLock.Lock()
|
||||
defer a.sessionsLock.Unlock()
|
||||
|
||||
if session, ok := a.sessions[sessionID]; ok {
|
||||
return session
|
||||
}
|
||||
|
||||
session := NewSession(sessionID, a.config.MaxHistoryMessages)
|
||||
a.sessions[sessionID] = session
|
||||
return session
|
||||
}
|
||||
|
||||
// Chat processes a user message and returns the agent's response
|
||||
// This is the main entry point for the agent loop
|
||||
func (a *Agent) Chat(ctx context.Context, sessionID string, userMessage string) (*AgentResponse, error) {
|
||||
session := a.GetSession(sessionID)
|
||||
|
||||
// Add user message to history
|
||||
session.AddMessage(Message{
|
||||
Role: "user",
|
||||
Content: userMessage,
|
||||
Timestamp: time.Now(),
|
||||
})
|
||||
|
||||
// Build the full prompt with tools
|
||||
systemPrompt := a.buildSystemPromptWithTools()
|
||||
conversationPrompt := a.buildConversationPrompt(session)
|
||||
|
||||
// Agent loop - keep calling AI until it's done or max iterations
|
||||
var finalResponse string
|
||||
toolCallCount := 0
|
||||
|
||||
for {
|
||||
// Check context cancellation
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// Check max tool calls
|
||||
if toolCallCount >= a.config.MaxToolCalls {
|
||||
logger.Warnf("⚠️ Max tool calls reached (%d), stopping agent loop", a.config.MaxToolCalls)
|
||||
break
|
||||
}
|
||||
|
||||
// Call AI
|
||||
response, err := a.aiClient.CallWithMessages(systemPrompt, conversationPrompt)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("AI call failed: %w", err)
|
||||
}
|
||||
|
||||
// Parse response for tool calls
|
||||
toolCalls, textResponse, err := a.parseResponse(response)
|
||||
if err != nil {
|
||||
// If parsing fails, treat entire response as text
|
||||
finalResponse = response
|
||||
break
|
||||
}
|
||||
|
||||
// If no tool calls, we're done
|
||||
if len(toolCalls) == 0 {
|
||||
finalResponse = textResponse
|
||||
break
|
||||
}
|
||||
|
||||
// Execute tool calls
|
||||
toolResults := a.executeToolCalls(ctx, toolCalls)
|
||||
toolCallCount += len(toolCalls)
|
||||
|
||||
// Add tool calls and results to conversation for next iteration
|
||||
conversationPrompt += fmt.Sprintf("\n\nAssistant called tools:\n%s\n\nTool results:\n%s\n\nBased on the tool results, please provide your response to the user:",
|
||||
formatToolCalls(toolCalls),
|
||||
formatToolResults(toolResults))
|
||||
|
||||
// If there's also a text response, capture it
|
||||
if textResponse != "" {
|
||||
finalResponse = textResponse
|
||||
}
|
||||
}
|
||||
|
||||
// Add assistant response to history
|
||||
session.AddMessage(Message{
|
||||
Role: "assistant",
|
||||
Content: finalResponse,
|
||||
Timestamp: time.Now(),
|
||||
})
|
||||
|
||||
return &AgentResponse{
|
||||
Text: finalResponse,
|
||||
SessionID: sessionID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// buildSystemPromptWithTools creates the system prompt including tool definitions
|
||||
func (a *Agent) buildSystemPromptWithTools() string {
|
||||
a.toolsLock.RLock()
|
||||
defer a.toolsLock.RUnlock()
|
||||
|
||||
var toolDefs []string
|
||||
for _, tool := range a.tools {
|
||||
toolDef := fmt.Sprintf(`- **%s**: %s
|
||||
Parameters: %s`, tool.Name(), tool.Description(), tool.ParameterSchema())
|
||||
toolDefs = append(toolDefs, toolDef)
|
||||
}
|
||||
|
||||
toolsSection := ""
|
||||
if len(toolDefs) > 0 {
|
||||
toolsSection = fmt.Sprintf(`
|
||||
|
||||
## Available Tools
|
||||
|
||||
You can call tools by responding with JSON in this format:
|
||||
{"tool_calls": [{"name": "tool_name", "arguments": {"param": "value"}}]}
|
||||
|
||||
After receiving tool results, provide a natural language response to the user.
|
||||
|
||||
Tools:
|
||||
%s
|
||||
`, strings.Join(toolDefs, "\n"))
|
||||
}
|
||||
|
||||
return a.systemPrompt + toolsSection
|
||||
}
|
||||
|
||||
// buildConversationPrompt builds the conversation history as a prompt
|
||||
func (a *Agent) buildConversationPrompt(session *Session) string {
|
||||
messages := session.GetMessages()
|
||||
var parts []string
|
||||
|
||||
for _, msg := range messages {
|
||||
parts = append(parts, fmt.Sprintf("%s: %s", strings.Title(msg.Role), msg.Content))
|
||||
}
|
||||
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// parseResponse extracts tool calls and text from AI response
|
||||
func (a *Agent) parseResponse(response string) ([]ToolCall, string, error) {
|
||||
// Try to find JSON tool calls in response
|
||||
// Look for {"tool_calls": [...]} pattern
|
||||
|
||||
var toolCalls []ToolCall
|
||||
textResponse := response
|
||||
|
||||
// Try to parse as JSON
|
||||
if strings.Contains(response, "tool_calls") {
|
||||
// Find JSON block
|
||||
start := strings.Index(response, "{")
|
||||
end := strings.LastIndex(response, "}")
|
||||
|
||||
if start >= 0 && end > start {
|
||||
jsonStr := response[start : end+1]
|
||||
|
||||
var parsed struct {
|
||||
ToolCalls []struct {
|
||||
Name string `json:"name"`
|
||||
Arguments json.RawMessage `json:"arguments"`
|
||||
} `json:"tool_calls"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal([]byte(jsonStr), &parsed); err == nil {
|
||||
for _, tc := range parsed.ToolCalls {
|
||||
toolCalls = append(toolCalls, ToolCall{
|
||||
Name: tc.Name,
|
||||
Arguments: tc.Arguments,
|
||||
})
|
||||
}
|
||||
|
||||
// Extract text before/after JSON
|
||||
textResponse = strings.TrimSpace(response[:start] + response[end+1:])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return toolCalls, textResponse, nil
|
||||
}
|
||||
|
||||
// executeToolCalls runs the requested tools
|
||||
func (a *Agent) executeToolCalls(ctx context.Context, calls []ToolCall) []ToolResult {
|
||||
a.toolsLock.RLock()
|
||||
defer a.toolsLock.RUnlock()
|
||||
|
||||
var results []ToolResult
|
||||
|
||||
for _, call := range calls {
|
||||
tool, ok := a.tools[call.Name]
|
||||
if !ok {
|
||||
results = append(results, ToolResult{
|
||||
Name: call.Name,
|
||||
Error: fmt.Sprintf("unknown tool: %s", call.Name),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("🔧 Executing tool: %s", call.Name)
|
||||
|
||||
result, err := tool.Execute(ctx, call.Arguments)
|
||||
if err != nil {
|
||||
logger.Errorf("❌ Tool %s failed: %v", call.Name, err)
|
||||
results = append(results, ToolResult{
|
||||
Name: call.Name,
|
||||
Error: err.Error(),
|
||||
})
|
||||
} else {
|
||||
logger.Infof("✅ Tool %s completed", call.Name)
|
||||
results = append(results, ToolResult{
|
||||
Name: call.Name,
|
||||
Result: result,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// ToolCall represents a tool invocation request from the AI
|
||||
type ToolCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments json.RawMessage `json:"arguments"`
|
||||
}
|
||||
|
||||
// ToolResult represents the result of a tool execution
|
||||
type ToolResult struct {
|
||||
Name string `json:"name"`
|
||||
Result interface{} `json:"result,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// AgentResponse is the final response from the agent
|
||||
type AgentResponse struct {
|
||||
Text string `json:"text"`
|
||||
SessionID string `json:"session_id"`
|
||||
}
|
||||
|
||||
func formatToolCalls(calls []ToolCall) string {
|
||||
var parts []string
|
||||
for _, c := range calls {
|
||||
parts = append(parts, fmt.Sprintf("- %s(%s)", c.Name, string(c.Arguments)))
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func formatToolResults(results []ToolResult) string {
|
||||
var parts []string
|
||||
for _, r := range results {
|
||||
if r.Error != "" {
|
||||
parts = append(parts, fmt.Sprintf("- %s: ERROR: %s", r.Name, r.Error))
|
||||
} else {
|
||||
resultJSON, _ := json.Marshal(r.Result)
|
||||
parts = append(parts, fmt.Sprintf("- %s: %s", r.Name, string(resultJSON)))
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
117
assistant/prompts.go
Normal file
117
assistant/prompts.go
Normal file
@@ -0,0 +1,117 @@
|
||||
package assistant
|
||||
|
||||
// DefaultTradingSystemPrompt returns the default system prompt for trading assistant
|
||||
func DefaultTradingSystemPrompt() string {
|
||||
return `# NOFX Trading Assistant
|
||||
|
||||
You are an expert AI trading assistant powered by NOFX - an advanced AI-powered trading system.
|
||||
|
||||
## Your Capabilities
|
||||
|
||||
1. **Account Management**
|
||||
- Check balances across multiple exchanges
|
||||
- View current positions and P&L
|
||||
- Monitor portfolio performance
|
||||
|
||||
2. **Trading Operations**
|
||||
- Execute trades (open/close positions)
|
||||
- Manage stop-loss and take-profit orders
|
||||
- Adjust leverage and margin settings
|
||||
|
||||
3. **AI Traders Management**
|
||||
- Start/stop AI traders
|
||||
- Monitor AI trader performance
|
||||
- Configure trading strategies
|
||||
|
||||
4. **Strategy & Analysis**
|
||||
- Create and modify trading strategies
|
||||
- Initiate AI debate sessions for market analysis
|
||||
- Backtest strategies on historical data
|
||||
|
||||
5. **Market Intelligence**
|
||||
- Get real-time prices and market data
|
||||
- Analyze market conditions
|
||||
- Track open interest and funding rates
|
||||
|
||||
## Guidelines
|
||||
|
||||
1. **Safety First**: Always confirm with the user before executing trades or making significant changes
|
||||
2. **Be Precise**: When dealing with numbers, be exact - trading involves real money
|
||||
3. **Explain Reasoning**: Help users understand your analysis and recommendations
|
||||
4. **Risk Awareness**: Always remind users about the risks involved in trading
|
||||
5. **Proactive Monitoring**: Alert users to important position changes or market movements
|
||||
|
||||
## Response Style
|
||||
|
||||
- Be concise but thorough
|
||||
- Use tables for data when appropriate
|
||||
- Include relevant metrics (P&L, ROI, etc.)
|
||||
- Provide actionable insights, not just data dumps
|
||||
- Support both English and Chinese (respond in the user's language)
|
||||
|
||||
## Important Notes
|
||||
|
||||
- Never share API keys or sensitive credentials
|
||||
- Always use proper position sizing based on user's risk tolerance
|
||||
- Warn users about high-risk operations (high leverage, large positions)
|
||||
|
||||
Remember: You are a professional trading assistant. Users trust you with their trading operations. Be accurate, be helpful, and be responsible.`
|
||||
}
|
||||
|
||||
// ChineseSystemPrompt returns Chinese version of the system prompt
|
||||
func ChineseSystemPrompt() string {
|
||||
return `# NOFX 交易助手
|
||||
|
||||
你是一个由 NOFX 驱动的专业 AI 交易助手 - 一个先进的 AI 驱动交易系统。
|
||||
|
||||
## 你的能力
|
||||
|
||||
1. **账户管理**
|
||||
- 查询多交易所余额
|
||||
- 查看当前持仓和盈亏
|
||||
- 监控投资组合表现
|
||||
|
||||
2. **交易操作**
|
||||
- 执行交易（开仓/平仓）
|
||||
- 管理止损止盈订单
|
||||
- 调整杠杆和保证金设置
|
||||
|
||||
3. **AI 交易员管理**
|
||||
- 启动/停止 AI 交易员
|
||||
- 监控 AI 交易员表现
|
||||
- 配置交易策略
|
||||
|
||||
4. **策略与分析**
|
||||
- 创建和修改交易策略
|
||||
- 发起 AI 辩论会议进行市场分析
|
||||
- 回测历史数据
|
||||
|
||||
5. **市场情报**
|
||||
- 获取实时价格和市场数据
|
||||
- 分析市场状况
|
||||
- 跟踪持仓量和资金费率
|
||||
|
||||
## 行为准则
|
||||
|
||||
1. **安全第一**：执行交易或重大操作前，务必与用户确认
|
||||
2. **精确无误**：涉及数字时必须精确 - 交易涉及真金白银
|
||||
3. **解释逻辑**：帮助用户理解你的分析和建议
|
||||
4. **风险意识**：始终提醒用户交易风险
|
||||
5. **主动监控**：及时提醒用户重要的仓位变化或市场波动
|
||||
|
||||
## 回复风格
|
||||
|
||||
- 简洁但全面
|
||||
- 适当使用表格展示数据
|
||||
- 包含相关指标（盈亏、收益率等）
|
||||
- 提供可操作的见解，而非单纯的数据罗列
|
||||
- 支持中英文（根据用户使用的语言回复）
|
||||
|
||||
## 重要提示
|
||||
|
||||
- 永远不要分享 API 密钥或敏感凭证
|
||||
- 根据用户的风险承受能力进行合理的仓位管理
|
||||
- 对高风险操作（高杠杆、大仓位）发出警告
|
||||
|
||||
记住：你是专业的交易助手。用户将交易操作托付于你。准确、有用、负责。`
|
||||
}
|
||||
122
assistant/session.go
Normal file
122
assistant/session.go
Normal file
@@ -0,0 +1,122 @@
|
||||
package assistant
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Message represents a single message in conversation
|
||||
type Message struct {
|
||||
Role string `json:"role"` // "user", "assistant", "system", "tool"
|
||||
Content string `json:"content"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
|
||||
// For tool messages
|
||||
ToolName string `json:"tool_name,omitempty"`
|
||||
ToolResult interface{} `json:"tool_result,omitempty"`
|
||||
}
|
||||
|
||||
// Session represents a conversation session with memory
|
||||
type Session struct {
|
||||
ID string `json:"id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
|
||||
// User info
|
||||
UserID string `json:"user_id"`
|
||||
UserName string `json:"user_name"`
|
||||
Platform string `json:"platform"` // "telegram", "web", etc.
|
||||
|
||||
// Conversation history
|
||||
messages []Message
|
||||
maxMessages int
|
||||
mu sync.RWMutex
|
||||
|
||||
// Custom metadata
|
||||
Metadata map[string]interface{} `json:"metadata"`
|
||||
}
|
||||
|
||||
// NewSession creates a new conversation session
|
||||
func NewSession(id string, maxMessages int) *Session {
|
||||
return &Session{
|
||||
ID: id,
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
messages: make([]Message, 0),
|
||||
maxMessages: maxMessages,
|
||||
Metadata: make(map[string]interface{}),
|
||||
}
|
||||
}
|
||||
|
||||
// AddMessage adds a message to the session
|
||||
func (s *Session) AddMessage(msg Message) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
s.messages = append(s.messages, msg)
|
||||
s.UpdatedAt = time.Now()
|
||||
|
||||
// Trim old messages if exceeding max
|
||||
if len(s.messages) > s.maxMessages {
|
||||
// Keep the most recent messages
|
||||
s.messages = s.messages[len(s.messages)-s.maxMessages:]
|
||||
}
|
||||
}
|
||||
|
||||
// GetMessages returns a copy of all messages
|
||||
func (s *Session) GetMessages() []Message {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
result := make([]Message, len(s.messages))
|
||||
copy(result, s.messages)
|
||||
return result
|
||||
}
|
||||
|
||||
// GetRecentMessages returns the N most recent messages
|
||||
func (s *Session) GetRecentMessages(n int) []Message {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
if n >= len(s.messages) {
|
||||
result := make([]Message, len(s.messages))
|
||||
copy(result, s.messages)
|
||||
return result
|
||||
}
|
||||
|
||||
result := make([]Message, n)
|
||||
copy(result, s.messages[len(s.messages)-n:])
|
||||
return result
|
||||
}
|
||||
|
||||
// Clear removes all messages from the session
|
||||
func (s *Session) Clear() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.messages = make([]Message, 0)
|
||||
s.UpdatedAt = time.Now()
|
||||
}
|
||||
|
||||
// SetUserInfo sets user information
|
||||
func (s *Session) SetUserInfo(userID, userName, platform string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.UserID = userID
|
||||
s.UserName = userName
|
||||
s.Platform = platform
|
||||
}
|
||||
|
||||
// SetMetadata sets a metadata value
|
||||
func (s *Session) SetMetadata(key string, value interface{}) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.Metadata[key] = value
|
||||
}
|
||||
|
||||
// GetMetadata gets a metadata value
|
||||
func (s *Session) GetMetadata(key string) (interface{}, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
v, ok := s.Metadata[key]
|
||||
return v, ok
|
||||
}
|
||||
47
assistant/tool.go
Normal file
47
assistant/tool.go
Normal file
@@ -0,0 +1,47 @@
|
||||
package assistant
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
)
|
||||
|
||||
// Tool represents a callable tool that the AI agent can use
|
||||
type Tool interface {
|
||||
// Name returns the tool's unique identifier
|
||||
Name() string
|
||||
|
||||
// Description returns a human-readable description for the AI
|
||||
Description() string
|
||||
|
||||
// ParameterSchema returns JSON schema for the tool's parameters
|
||||
ParameterSchema() string
|
||||
|
||||
// Execute runs the tool with the given arguments
|
||||
Execute(ctx context.Context, args json.RawMessage) (interface{}, error)
|
||||
}
|
||||
|
||||
// BaseTool provides common functionality for tools
|
||||
type BaseTool struct {
|
||||
ToolName string
|
||||
ToolDescription string
|
||||
ToolSchema string
|
||||
ExecuteFunc func(ctx context.Context, args json.RawMessage) (interface{}, error)
|
||||
}
|
||||
|
||||
func (t *BaseTool) Name() string { return t.ToolName }
|
||||
func (t *BaseTool) Description() string { return t.ToolDescription }
|
||||
func (t *BaseTool) ParameterSchema() string { return t.ToolSchema }
|
||||
|
||||
func (t *BaseTool) Execute(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
return t.ExecuteFunc(ctx, args)
|
||||
}
|
||||
|
||||
// NewTool creates a simple tool from a function
|
||||
func NewTool(name, description, schema string, fn func(ctx context.Context, args json.RawMessage) (interface{}, error)) Tool {
|
||||
return &BaseTool{
|
||||
ToolName: name,
|
||||
ToolDescription: description,
|
||||
ToolSchema: schema,
|
||||
ExecuteFunc: fn,
|
||||
}
|
||||
}
|
||||
530
assistant/trading_tools.go
Normal file
530
assistant/trading_tools.go
Normal file
@@ -0,0 +1,530 @@
|
||||
package assistant
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"nofx/logger"
|
||||
"nofx/manager"
|
||||
"nofx/store"
|
||||
)
|
||||
|
||||
// TradingTools provides all trading-related tools for the AI agent
|
||||
type TradingTools struct {
|
||||
traderManager *manager.TraderManager
|
||||
store *store.Store
|
||||
}
|
||||
|
||||
// NewTradingTools creates trading tools with access to NOFX core
|
||||
func NewTradingTools(tm *manager.TraderManager, st *store.Store) *TradingTools {
|
||||
return &TradingTools{
|
||||
traderManager: tm,
|
||||
store: st,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAllTools returns all trading tools
|
||||
func (t *TradingTools) GetAllTools() []Tool {
|
||||
return []Tool{
|
||||
t.GetBalanceTool(),
|
||||
t.GetPositionsTool(),
|
||||
t.ListTradersTool(),
|
||||
t.GetTraderStatusTool(),
|
||||
t.StartTraderTool(),
|
||||
t.StopTraderTool(),
|
||||
t.GetMarketPriceTool(),
|
||||
t.OpenLongTool(),
|
||||
t.OpenShortTool(),
|
||||
t.ClosePositionTool(),
|
||||
t.ListStrategiesTool(),
|
||||
t.ListExchangesTool(),
|
||||
t.ListAIModelsTool(),
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Query Tools ====================
|
||||
|
||||
// GetBalanceTool returns the get_balance tool
|
||||
func (t *TradingTools) GetBalanceTool() Tool {
|
||||
return NewTool(
|
||||
"get_balance",
|
||||
"Get account balance for a trader. Returns available balance, total equity, and margin info.",
|
||||
`{"trader_id": "string (required) - The trader ID to query"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
trader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
balance, err := trader.GetAccountInfo()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get balance: %w", err)
|
||||
}
|
||||
|
||||
return balance, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// GetPositionsTool returns the get_positions tool
|
||||
func (t *TradingTools) GetPositionsTool() Tool {
|
||||
return NewTool(
|
||||
"get_positions",
|
||||
"Get all open positions for a trader. Returns symbol, side, size, entry price, unrealized P&L.",
|
||||
`{"trader_id": "string (required) - The trader ID to query"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
trader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
positions, err := trader.GetPositions()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get positions: %w", err)
|
||||
}
|
||||
|
||||
return positions, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ListTradersTool returns the list_traders tool
|
||||
func (t *TradingTools) ListTradersTool() Tool {
|
||||
return NewTool(
|
||||
"list_traders",
|
||||
"List all configured AI traders with their status (running/stopped), exchange, AI model, and performance.",
|
||||
`{}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
traders, err := t.store.Trader().List("default")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list traders: %w", err)
|
||||
}
|
||||
|
||||
var result []map[string]interface{}
|
||||
for _, tr := range traders {
|
||||
traderInfo := map[string]interface{}{
|
||||
"id": tr.ID,
|
||||
"name": tr.Name,
|
||||
"is_running": tr.IsRunning,
|
||||
"ai_model_id": tr.AIModelID,
|
||||
"exchange_id": tr.ExchangeID,
|
||||
"strategy_id": tr.StrategyID,
|
||||
"created_at": tr.CreatedAt,
|
||||
}
|
||||
|
||||
// Try to get live status if trader is running
|
||||
if liveTrader, err := t.traderManager.GetTrader(tr.ID); err == nil {
|
||||
status := liveTrader.GetStatus()
|
||||
traderInfo["live_status"] = status
|
||||
}
|
||||
|
||||
result = append(result, traderInfo)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// GetTraderStatusTool returns detailed status of a specific trader
|
||||
func (t *TradingTools) GetTraderStatusTool() Tool {
|
||||
return NewTool(
|
||||
"get_trader_status",
|
||||
"Get detailed status of a specific trader including current positions, recent trades, and performance metrics.",
|
||||
`{"trader_id": "string (required) - The trader ID to query"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
// Get trader config from store
|
||||
traderConfig, err := t.store.Trader().GetByID(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
result := map[string]interface{}{
|
||||
"id": traderConfig.ID,
|
||||
"name": traderConfig.Name,
|
||||
"is_running": traderConfig.IsRunning,
|
||||
"ai_model_id": traderConfig.AIModelID,
|
||||
"exchange_id": traderConfig.ExchangeID,
|
||||
"strategy_id": traderConfig.StrategyID,
|
||||
}
|
||||
|
||||
// If trader is running, get live data
|
||||
trader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err == nil && trader != nil {
|
||||
result["live_status"] = trader.GetStatus()
|
||||
if balance, err := trader.GetAccountInfo(); err == nil {
|
||||
result["balance"] = balance
|
||||
}
|
||||
if positions, err := trader.GetPositions(); err == nil {
|
||||
result["positions"] = positions
|
||||
}
|
||||
}
|
||||
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ==================== Control Tools ====================
|
||||
|
||||
// StartTraderTool starts an AI trader
|
||||
func (t *TradingTools) StartTraderTool() Tool {
|
||||
return NewTool(
|
||||
"start_trader",
|
||||
"Start an AI trader to begin automated trading. The trader will execute trades based on its configured strategy.",
|
||||
`{"trader_id": "string (required) - The trader ID to start"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
// Check if already running
|
||||
existingTrader, _ := t.traderManager.GetTrader(params.TraderID)
|
||||
if existingTrader != nil {
|
||||
status := existingTrader.GetStatus()
|
||||
if isRunning, ok := status["is_running"].(bool); ok && isRunning {
|
||||
return nil, fmt.Errorf("trader is already running")
|
||||
}
|
||||
// Remove from memory to reload
|
||||
t.traderManager.RemoveTrader(params.TraderID)
|
||||
}
|
||||
|
||||
// Load and start trader
|
||||
if err := t.traderManager.LoadUserTradersFromStore(t.store, "default"); err != nil {
|
||||
return nil, fmt.Errorf("failed to load trader: %w", err)
|
||||
}
|
||||
|
||||
trader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get trader after load: %w", err)
|
||||
}
|
||||
|
||||
// Start the trader in a goroutine
|
||||
go func() {
|
||||
if err := trader.Run(); err != nil {
|
||||
logger.Errorf("Trader %s error: %v", params.TraderID, err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Update status in database
|
||||
if err := t.store.Trader().UpdateStatus("default", params.TraderID, true); err != nil {
|
||||
logger.Warnf("Failed to update trader status in DB: %v", err)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"success": true,
|
||||
"trader_id": params.TraderID,
|
||||
"message": "Trader started successfully",
|
||||
}, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// StopTraderTool stops an AI trader
|
||||
func (t *TradingTools) StopTraderTool() Tool {
|
||||
return NewTool(
|
||||
"stop_trader",
|
||||
"Stop an AI trader. This will halt automated trading but keep existing positions open.",
|
||||
`{"trader_id": "string (required) - The trader ID to stop"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
trader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
// Check if running
|
||||
status := trader.GetStatus()
|
||||
if isRunning, ok := status["is_running"].(bool); ok && !isRunning {
|
||||
return nil, fmt.Errorf("trader is already stopped")
|
||||
}
|
||||
|
||||
// Stop the trader
|
||||
trader.Stop()
|
||||
|
||||
// Update status in database
|
||||
if err := t.store.Trader().UpdateStatus("default", params.TraderID, false); err != nil {
|
||||
logger.Warnf("Failed to update trader status in DB: %v", err)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"success": true,
|
||||
"trader_id": params.TraderID,
|
||||
"message": "Trader stopped successfully",
|
||||
}, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ==================== Trading Tools ====================
|
||||
|
||||
// GetMarketPriceTool gets current market price
|
||||
func (t *TradingTools) GetMarketPriceTool() Tool {
|
||||
return NewTool(
|
||||
"get_market_price",
|
||||
"Get current market price for a trading pair from a specific trader's exchange.",
|
||||
`{"trader_id": "string (required)", "symbol": "string (required) - e.g., BTCUSDT"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
Symbol string `json:"symbol"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
autoTrader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
// Get the underlying trader interface
|
||||
underlyingTrader := autoTrader.GetUnderlyingTrader()
|
||||
if underlyingTrader == nil {
|
||||
return nil, fmt.Errorf("underlying trader not available")
|
||||
}
|
||||
|
||||
price, err := underlyingTrader.GetMarketPrice(params.Symbol)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get price: %w", err)
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"symbol": params.Symbol,
|
||||
"price": price,
|
||||
}, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// OpenLongTool opens a long position
|
||||
func (t *TradingTools) OpenLongTool() Tool {
|
||||
return NewTool(
|
||||
"open_long",
|
||||
"Open a long (buy) position. WARNING: This will execute a real trade!",
|
||||
`{"trader_id": "string (required)", "symbol": "string (required)", "quantity": "number (required)", "leverage": "number (optional, default 1)"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
Symbol string `json:"symbol"`
|
||||
Quantity float64 `json:"quantity"`
|
||||
Leverage int `json:"leverage"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
if params.Leverage == 0 {
|
||||
params.Leverage = 1
|
||||
}
|
||||
|
||||
autoTrader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
underlyingTrader := autoTrader.GetUnderlyingTrader()
|
||||
if underlyingTrader == nil {
|
||||
return nil, fmt.Errorf("underlying trader not available")
|
||||
}
|
||||
|
||||
result, err := underlyingTrader.OpenLong(params.Symbol, params.Quantity, params.Leverage)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to open long: %w", err)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// OpenShortTool opens a short position
|
||||
func (t *TradingTools) OpenShortTool() Tool {
|
||||
return NewTool(
|
||||
"open_short",
|
||||
"Open a short (sell) position. WARNING: This will execute a real trade!",
|
||||
`{"trader_id": "string (required)", "symbol": "string (required)", "quantity": "number (required)", "leverage": "number (optional, default 1)"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
Symbol string `json:"symbol"`
|
||||
Quantity float64 `json:"quantity"`
|
||||
Leverage int `json:"leverage"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
if params.Leverage == 0 {
|
||||
params.Leverage = 1
|
||||
}
|
||||
|
||||
autoTrader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
underlyingTrader := autoTrader.GetUnderlyingTrader()
|
||||
if underlyingTrader == nil {
|
||||
return nil, fmt.Errorf("underlying trader not available")
|
||||
}
|
||||
|
||||
result, err := underlyingTrader.OpenShort(params.Symbol, params.Quantity, params.Leverage)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to open short: %w", err)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ClosePositionTool closes a position
|
||||
func (t *TradingTools) ClosePositionTool() Tool {
|
||||
return NewTool(
|
||||
"close_position",
|
||||
"Close an existing position (long or short). WARNING: This will execute a real trade!",
|
||||
`{"trader_id": "string (required)", "symbol": "string (required)", "side": "string (required) - 'long' or 'short'", "quantity": "number (optional) - leave empty to close all"}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
var params struct {
|
||||
TraderID string `json:"trader_id"`
|
||||
Symbol string `json:"symbol"`
|
||||
Side string `json:"side"`
|
||||
Quantity float64 `json:"quantity"`
|
||||
}
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("invalid arguments: %w", err)
|
||||
}
|
||||
|
||||
autoTrader, err := t.traderManager.GetTrader(params.TraderID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("trader not found: %w", err)
|
||||
}
|
||||
|
||||
underlyingTrader := autoTrader.GetUnderlyingTrader()
|
||||
if underlyingTrader == nil {
|
||||
return nil, fmt.Errorf("underlying trader not available")
|
||||
}
|
||||
|
||||
var result map[string]interface{}
|
||||
|
||||
if params.Side == "long" {
|
||||
result, err = underlyingTrader.CloseLong(params.Symbol, params.Quantity)
|
||||
} else if params.Side == "short" {
|
||||
result, err = underlyingTrader.CloseShort(params.Symbol, params.Quantity)
|
||||
} else {
|
||||
return nil, fmt.Errorf("invalid side: %s (must be 'long' or 'short')", params.Side)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to close position: %w", err)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ==================== Config Tools ====================
|
||||
|
||||
// ListStrategiesTool lists all strategies
|
||||
func (t *TradingTools) ListStrategiesTool() Tool {
|
||||
return NewTool(
|
||||
"list_strategies",
|
||||
"List all trading strategies configured in the system.",
|
||||
`{}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
strategies, err := t.store.Strategy().List("default")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list strategies: %w", err)
|
||||
}
|
||||
return strategies, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ListExchangesTool lists all exchange configurations
|
||||
func (t *TradingTools) ListExchangesTool() Tool {
|
||||
return NewTool(
|
||||
"list_exchanges",
|
||||
"List all configured exchanges (without showing API keys).",
|
||||
`{}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
exchanges, err := t.store.Exchange().List("default")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list exchanges: %w", err)
|
||||
}
|
||||
|
||||
// Remove sensitive data
|
||||
var result []map[string]interface{}
|
||||
for _, ex := range exchanges {
|
||||
result = append(result, map[string]interface{}{
|
||||
"id": ex.ID,
|
||||
"name": ex.Name,
|
||||
"exchange_type": ex.ExchangeType,
|
||||
"type": ex.Type,
|
||||
"enabled": ex.Enabled,
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// ListAIModelsTool lists all AI model configurations
|
||||
func (t *TradingTools) ListAIModelsTool() Tool {
|
||||
return NewTool(
|
||||
"list_ai_models",
|
||||
"List all configured AI models (without showing API keys).",
|
||||
`{}`,
|
||||
func(ctx context.Context, args json.RawMessage) (interface{}, error) {
|
||||
models, err := t.store.AIModel().List("default")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list AI models: %w", err)
|
||||
}
|
||||
|
||||
// Remove sensitive data
|
||||
var result []map[string]interface{}
|
||||
for _, m := range models {
|
||||
result = append(result, map[string]interface{}{
|
||||
"id": m.ID,
|
||||
"name": m.Name,
|
||||
"provider": m.Provider,
|
||||
"custom_model": m.CustomModelName,
|
||||
"enabled": m.Enabled,
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
},
|
||||
)
|
||||
}
|
||||
Reference in New Issue
Block a user