mirror of
https://github.com/NoFxAiOS/nofx.git
synced 2026-07-15 00:36:56 +08:00
feat(telegram): add AI agent bot with streaming and account context
- Add Telegram bot with long-polling and AI agent loop (api_call tool)
- SSE streaming with real-time message editing and ⏳ placeholder
- Account state injection at conversation start (models, exchanges,
strategies, traders, per-trader PnL and statistics)
- Lane semaphore per chat serializes concurrent messages (60s timeout)
- Idle timeout watchdog (60s) prevents hung streaming connections
- Look-ahead buffer prevents partial <api_call> tag leaking to user
- Fix PUT /strategies/:id to merge config (read-then-merge pattern)
- Add route registry with full API schema for LLM documentation
- Add TelegramConfig store and Web UI config modal
- Add GetAnyEnabled to AIModel store for bot LLM client selection
This commit is contained in:
183
telegram/agent/agent_test.go
Normal file
183
telegram/agent/agent_test.go
Normal file
@@ -0,0 +1,183 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"nofx/mcp"
|
||||
)
|
||||
|
||||
type mockLLM struct {
|
||||
responses []string
|
||||
calls int
|
||||
lastMsgs []mcp.Message
|
||||
}
|
||||
|
||||
func (m *mockLLM) SetAPIKey(_, _, _ string) {}
|
||||
func (m *mockLLM) SetTimeout(_ time.Duration) {}
|
||||
func (m *mockLLM) CallWithMessages(_, _ string) (string, error) { return m.next() }
|
||||
func (m *mockLLM) CallWithRequest(req *mcp.Request) (string, error) {
|
||||
m.lastMsgs = req.Messages
|
||||
return m.next()
|
||||
}
|
||||
func (m *mockLLM) CallWithRequestStream(req *mcp.Request, onChunk func(string)) (string, error) {
|
||||
m.lastMsgs = req.Messages
|
||||
r, err := m.next()
|
||||
if onChunk != nil {
|
||||
onChunk(r)
|
||||
}
|
||||
return r, err
|
||||
}
|
||||
func (m *mockLLM) next() (string, error) {
|
||||
if m.calls < len(m.responses) {
|
||||
r := m.responses[m.calls]
|
||||
m.calls++
|
||||
return r, nil
|
||||
}
|
||||
return "OK", nil
|
||||
}
|
||||
|
||||
func mockGetLLM(llm *mockLLM) func() mcp.AIClient {
|
||||
return func() mcp.AIClient { return llm }
|
||||
}
|
||||
|
||||
const testPrompt = "You are a test assistant."
|
||||
|
||||
// TestAgentDirectReply: LLM replies without api_call — one call, direct reply.
|
||||
func TestAgentDirectReply(t *testing.T) {
|
||||
llm := &mockLLM{responses: []string{"Hello! How can I help you?"}}
|
||||
a := New(8080, "tok", "test-user", mockGetLLM(llm), testPrompt)
|
||||
|
||||
reply := a.Run("hello", nil)
|
||||
|
||||
if reply != "Hello! How can I help you?" {
|
||||
t.Fatalf("unexpected reply: %q", reply)
|
||||
}
|
||||
if llm.calls != 1 {
|
||||
t.Fatalf("expected 1 LLM call, got %d", llm.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAgentAPICall: LLM calls API, gets result, gives final reply — two LLM calls.
|
||||
func TestAgentAPICall(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/api/my-traders" {
|
||||
w.Write([]byte(`[{"id":"t1","name":"BTC Strategy"}]`)) //nolint:errcheck
|
||||
return
|
||||
}
|
||||
w.WriteHeader(404)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
var port int
|
||||
fmt.Sscanf(srv.Listener.Addr().String(), "127.0.0.1:%d", &port)
|
||||
|
||||
llm := &mockLLM{responses: []string{
|
||||
`Let me check.<api_call>{"method":"GET","path":"/api/my-traders","body":{}}</api_call>`,
|
||||
"You have one trader: BTC Strategy.",
|
||||
}}
|
||||
a := New(port, "tok", "test-user", mockGetLLM(llm), testPrompt)
|
||||
|
||||
reply := a.Run("list my traders", nil)
|
||||
|
||||
if reply != "You have one trader: BTC Strategy." {
|
||||
t.Fatalf("unexpected reply: %q", reply)
|
||||
}
|
||||
if llm.calls != 2 {
|
||||
t.Fatalf("expected 2 LLM calls, got %d", llm.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAgentMultiStep: LLM chains two API calls before final reply — three LLM calls.
|
||||
func TestAgentMultiStep(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte(`{"ok":true}`)) //nolint:errcheck
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
var port int
|
||||
fmt.Sscanf(srv.Listener.Addr().String(), "127.0.0.1:%d", &port)
|
||||
|
||||
llm := &mockLLM{responses: []string{
|
||||
`Checking account.<api_call>{"method":"GET","path":"/api/account","body":{}}</api_call>`,
|
||||
`Now checking positions.<api_call>{"method":"GET","path":"/api/positions","body":{}}</api_call>`,
|
||||
"Account looks healthy and no open positions.",
|
||||
}}
|
||||
a := New(port, "tok", "test-user", mockGetLLM(llm), testPrompt)
|
||||
|
||||
reply := a.Run("show me account status", nil)
|
||||
|
||||
if llm.calls != 3 {
|
||||
t.Fatalf("expected 3 LLM calls (2 api + 1 final), got %d", llm.calls)
|
||||
}
|
||||
if reply != "Account looks healthy and no open positions." {
|
||||
t.Fatalf("unexpected final reply: %q", reply)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAgentAPIResultInContext: API result must appear in next LLM message.
|
||||
func TestAgentAPIResultInContext(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte(`{"balance":1234.56}`)) //nolint:errcheck
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
var port int
|
||||
fmt.Sscanf(srv.Listener.Addr().String(), "127.0.0.1:%d", &port)
|
||||
|
||||
llm := &mockLLM{responses: []string{
|
||||
`<api_call>{"method":"GET","path":"/api/account","body":{}}</api_call>`,
|
||||
"Balance is 1234.56 USDT.",
|
||||
}}
|
||||
a := New(port, "tok", "test-user", mockGetLLM(llm), testPrompt)
|
||||
a.Run("show balance", nil)
|
||||
|
||||
found := false
|
||||
for _, msg := range llm.lastMsgs {
|
||||
if strings.Contains(msg.Content, "API result") || strings.Contains(msg.Content, "balance") {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("API result not found in subsequent LLM context")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParseAPICall: unit tests for the XML tag parser.
|
||||
func TestParseAPICall(t *testing.T) {
|
||||
t.Run("valid call", func(t *testing.T) {
|
||||
resp := `Stopping trader.<api_call>{"method":"POST","path":"/api/traders/t1/stop","body":{}}</api_call>`
|
||||
req, text := parseAPICall(resp)
|
||||
if req == nil {
|
||||
t.Fatal("expected api_call, got nil")
|
||||
}
|
||||
if req.Method != "POST" || req.Path != "/api/traders/t1/stop" {
|
||||
t.Fatalf("unexpected req: %+v", req)
|
||||
}
|
||||
if text != "Stopping trader." {
|
||||
t.Fatalf("unexpected text before tag: %q", text)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no call tag", func(t *testing.T) {
|
||||
req, text := parseAPICall("Just a reply.")
|
||||
if req != nil {
|
||||
t.Fatal("expected nil api_call")
|
||||
}
|
||||
if text != "Just a reply." {
|
||||
t.Fatalf("expected original text, got %q", text)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("malformed JSON", func(t *testing.T) {
|
||||
req, _ := parseAPICall(`<api_call>NOT JSON</api_call>`)
|
||||
if req != nil {
|
||||
t.Fatal("expected nil for malformed JSON")
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user