diff --git a/agent/agent.go b/agent/agent.go index 9a5dbc3f..6dace1b4 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -294,6 +294,7 @@ func (a *Agent) buildSystemPrompt(lang string) string { ## 工具使用 你可以调用以下工具来执行操作： +- **search_stock** — 搜索股票（支持中文名、英文名、代码）。当用户提到你不认识的股票时，先用这个工具搜索。 - **execute_trade** — 下单交易（做多/做空/平多/平空）。调用后会创建待确认订单，用户需回复"确认 trade_xxx"才会真正执行。 - **get_positions** — 查看当前所有持仓 - **get_balance** — 查看账户余额 @@ -338,6 +339,7 @@ func (a *Agent) buildSystemPrompt(lang string) string { ## Tools You can call these tools to take action: +- **search_stock** — Search for stocks by name, ticker, or code. Covers A-share, HK, and US markets. Use when the user mentions an unknown stock. - **execute_trade** — Place a trade order (open_long/open_short/close_long/close_short). Creates a pending order that requires user confirmation. - **get_positions** — View all current open positions - **get_balance** — View account balance and equity @@ -376,8 +378,8 @@ func (a *Agent) gatherContext(text string) string { } } - // A-share / stocks — try Sina Finance - stockCode, stockName := resolveStockCode(text) + // A-share / stocks — try Sina Finance (dynamic search as fallback) + stockCode, stockName := resolveStockCodeDynamic(text) if stockCode != "" { quote, err := fetchStockQuote(stockCode) if err == nil && quote.Price > 0 { diff --git a/agent/stock.go b/agent/stock.go index 580ebfcf..32c194bd 100644 --- a/agent/stock.go +++ b/agent/stock.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "net/http" + "net/url" "strconv" "strings" "time" @@ -100,6 +101,176 @@ func resolveStockCode(text string) (string, string) { return "", "" } +// SearchResult represents a stock search result from Sina suggest API. +type SearchResult struct { + Name string // Display name + Code string // Sina-style code (e.g. sz300750, hk00700, gb_tsla) + Ticker string // Raw ticker (e.g. 300750, 00700, tsla) + Type string // Market type code: 11=A股, 31=港股, 41=美股 + Market string // "A股", "港股", "美股" +} + +// searchStock queries Sina's suggest API for dynamic stock search. +// Returns matching stocks across A-share, HK, and US markets. +func searchStock(keyword string) ([]SearchResult, error) { + // type=11 (A股), 31 (港股), 41 (美股) + u := fmt.Sprintf("https://suggest3.sinajs.cn/suggest/type=11,31,41&key=%s&name=suggestdata", + url.QueryEscape(keyword)) + + req, _ := http.NewRequest("GET", u, nil) + req.Header.Set("Referer", "https://finance.sina.com.cn") + + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + reader := transform.NewReader(resp.Body, simplifiedchinese.GBK.NewDecoder()) + body, err := io.ReadAll(reader) + if err != nil { + return nil, err + } + + line := string(body) + // Parse: var suggestdata="item1;item2;..." + start := strings.Index(line, "\"") + end := strings.LastIndex(line, "\"") + if start == -1 || end <= start { + return nil, fmt.Errorf("invalid suggest response") + } + data := line[start+1 : end] + if data == "" { + return nil, nil // no results + } + + var results []SearchResult + items := strings.Split(data, ";") + for _, item := range items { + item = strings.TrimSpace(item) + if item == "" { + continue + } + fields := strings.Split(item, ",") + if len(fields) < 5 { + continue + } + // fields: [0]=name, [1]=type, [2]=ticker, [3]=sinaCode, [4]=displayName + typeCode := fields[1] + ticker := fields[2] + sinaCode := fields[3] + displayName := fields[4] + if displayName == "" { + displayName = fields[0] + } + + var mkt, code string + switch typeCode { + case "11": // A股 + mkt = "A股" + code = sinaCode // already like sz300750, sh600519 + if code == "" { + // Build from ticker + prefix := "sz" + if len(ticker) == 6 && (ticker[0] == '6' || ticker[0] == '9') { + prefix = "sh" + } + code = prefix + ticker + } + case "31": // 港股 + mkt = "港股" + code = "hk" + ticker + case "41": // 美股 + mkt = "美股" + code = "gb_" + ticker + default: + continue // skip funds (201), indices, etc. + } + + results = append(results, SearchResult{ + Name: displayName, + Code: code, + Ticker: ticker, + Type: typeCode, + Market: mkt, + }) + } + + return results, nil +} + +// resolveStockCodeDynamic tries local map first, then falls back to Sina search API. +func resolveStockCodeDynamic(text string) (string, string) { + // First try the static map + code, name := resolveStockCode(text) + if code != "" { + return code, name + } + + // Fall back to Sina search API + // Extract a meaningful search keyword from the text + keyword := extractStockKeyword(text) + if keyword == "" { + return "", "" + } + + results, err := searchStock(keyword) + if err != nil || len(results) == 0 { + return "", "" + } + + // Return the first (best) result + return results[0].Code, results[0].Name +} + +// extractStockKeyword extracts a likely stock name/ticker from user text. +func extractStockKeyword(text string) string { + // Remove common prefixes/suffixes that aren't stock names + text = strings.TrimSpace(text) + + // If the text itself is short enough, use it directly + // (e.g. "中远海控" or "AAPL") + if len([]rune(text)) <= 10 { + return text + } + + // Try to extract quoted terms first: 「xxx」 or "xxx" + quotePairs := [][2]string{ + {"「", "」"}, + {"\u201c", "\u201d"}, + {"\u2018", "\u2019"}, + {"\"", "\""}, + } + for _, pair := range quotePairs { + if s := strings.Index(text, pair[0]); s >= 0 { + if e := strings.Index(text[s+len(pair[0]):], pair[1]); e >= 0 { + return text[s+len(pair[0]) : s+len(pair[0])+e] + } + } + } + + // Look for patterns like "查 XXX", "搜索 XXX", "查一下 XXX" + for _, prefix := range []string{"查一下", "搜索", "查询", "看看", "搜一下", "查", "看", "search ", "find "} { + if idx := strings.Index(text, prefix); idx >= 0 { + rest := strings.TrimSpace(text[idx+len(prefix):]) + // Take the first "word" (either Chinese characters or English word) + words := strings.Fields(rest) + if len(words) > 0 { + return words[0] + } + } + } + + // Last resort: use first few words + words := strings.Fields(text) + if len(words) > 0 { + return words[0] + } + + return "" +} + func fetchStockQuote(code string) (*StockQuote, error) { url := fmt.Sprintf("https://hq.sinajs.cn/list=%s", code) req, _ := http.NewRequest("GET", url, nil) diff --git a/agent/tools.go b/agent/tools.go index c12996fd..57e490a9 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -13,6 +13,23 @@ import ( // agentTools defines the tools available to the LLM for autonomous action. func agentTools() []mcp.Tool { return []mcp.Tool{ + { + Type: "function", + Function: mcp.FunctionDef{ + Name: "search_stock", + Description: "Search for a stock by name, ticker symbol, or keyword. Searches across A-share (沪深), Hong Kong, and US markets. Returns a list of matching stocks with their codes. Use this when the user asks about a stock not in your known list, or when you need to find the exact code for a stock.", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "keyword": map[string]any{ + "type": "string", + "description": "Search keyword: stock name (e.g. '宁德时代', '腾讯'), ticker (e.g. 'TSLA', 'AAPL'), or stock code (e.g. '300750')", + }, + }, + "required": []string{"keyword"}, + }, + }, + }, { Type: "function", Function: mcp.FunctionDef{ @@ -82,6 +99,8 @@ func agentTools() []mcp.Tool { // handleToolCall processes a single tool call from the LLM and returns the result. func (a *Agent) handleToolCall(ctx context.Context, userID int64, lang string, tc mcp.ToolCall) string { switch tc.Function.Name { + case "search_stock": + return a.toolSearchStock(tc.Function.Arguments) case "execute_trade": return a.toolExecuteTrade(ctx, userID, lang, tc.Function.Arguments) case "get_positions": @@ -95,6 +114,60 @@ func (a *Agent) handleToolCall(ctx context.Context, userID int64, lang string, t } } +func (a *Agent) toolSearchStock(argsJSON string) string { + var args struct { + Keyword string `json:"keyword"` + } + if err := json.Unmarshal([]byte(argsJSON), &args); err != nil { + return fmt.Sprintf(`{"error": "invalid arguments: %s"}`, err) + } + + if args.Keyword == "" { + return `{"error": "keyword is required"}` + } + + results, err := searchStock(args.Keyword) + if err != nil { + return fmt.Sprintf(`{"error": "search failed: %s"}`, err) + } + + if len(results) == 0 { + return fmt.Sprintf(`{"results": [], "message": "no stocks found for '%s'"}`, args.Keyword) + } + + // Limit to top 10 results + if len(results) > 10 { + results = results[:10] + } + + // Also fetch real-time quotes for the top results (up to 3) + type enrichedResult struct { + Name string `json:"name"` + Code string `json:"code"` + Market string `json:"market"` + Quote *StockQuote `json:"quote,omitempty"` + } + + var enriched []enrichedResult + for i, r := range results { + er := enrichedResult{Name: r.Name, Code: r.Code, Market: r.Market} + if i < 3 { + q, qErr := fetchStockQuote(r.Code) + if qErr == nil && q.Price > 0 { + er.Quote = q + } + } + enriched = append(enriched, er) + } + + result, _ := json.Marshal(map[string]any{ + "keyword": args.Keyword, + "count": len(enriched), + "results": enriched, + }) + return string(result) +} + func (a *Agent) toolExecuteTrade(_ context.Context, userID int64, lang, argsJSON string) string { var args struct { Action string `json:"action"`