From b00bddcdecdda6f46c9a5f3e91c28c7307c711a9 Mon Sep 17 00:00:00 2001 From: spark Date: Tue, 4 Feb 2025 18:12:08 +0800 Subject: [PATCH] =?UTF-8?q?refactor(stock-data):=20=E9=87=8D=E6=9E=84?= =?UTF-8?q?=E8=82=A1=E7=A5=A8=E6=95=B0=E6=8D=AE=E8=8E=B7=E5=8F=96=E9=80=BB?= =?UTF-8?q?=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 移除了不必要的并发请求，简化了代码结构 - 新增 FetchPrice 函数，用于获取股票价格信息 - 优化 SearchStockInfo 函数，提高了搜索效率和准确性 - 新增 SearchStockInfoByCode 函数，用于根据股票代码获取相关信息- 修复了一些潜在的错误和性能问题 --- backend/data/openai_api.go | 42 +++++++-------- backend/data/openai_api_test.go | 2 +- backend/data/stock_data_api.go | 80 ++++++++++++++++++++++++----- backend/data/stock_data_api_test.go | 4 ++ 4 files changed, 93 insertions(+), 35 deletions(-) diff --git a/backend/data/openai_api.go b/backend/data/openai_api.go index f641880..241d9f2 100644 --- a/backend/data/openai_api.go +++ b/backend/data/openai_api.go @@ -142,7 +142,7 @@ func (o OpenAi) NewChatStream(stock, stockCode string) <-chan string { } wg := &sync.WaitGroup{} - wg.Add(4) + wg.Add(2) go func() { defer wg.Done() @@ -168,26 +168,26 @@ func (o OpenAi) NewChatStream(stock, stockCode string) <-chan string { } }() - go func() { - defer wg.Done() - messages := SearchStockInfo(stock, "depth") - for _, message := range *messages { - msg = append(msg, map[string]interface{}{ - "role": "assistant", - "content": message, - }) - } - }() - go func() { - defer wg.Done() - messages := SearchStockInfo(stock, "telegram") - for _, message := range *messages { - msg = append(msg, map[string]interface{}{ - "role": "assistant", - "content": message, - }) - } - }() + //go func() { + // defer wg.Done() + // messages := SearchStockInfo(stock, "depth") + // for _, message := range *messages { + // msg = append(msg, map[string]interface{}{ + // "role": "assistant", + // "content": message, + // }) + // } + //}() + //go func() { + // defer wg.Done() + // messages := SearchStockInfo(stock, "telegram") + // for _, message := range *messages { + // msg = append(msg, map[string]interface{}{ + // "role": "assistant", + // "content": message, + // }) + // } + //}() wg.Wait() msg = append(msg, map[string]interface{}{ diff --git a/backend/data/openai_api_test.go b/backend/data/openai_api_test.go index 48e084f..f339cc9 100644 --- a/backend/data/openai_api_test.go +++ b/backend/data/openai_api_test.go @@ -13,7 +13,7 @@ func TestNewDeepSeekOpenAiConfig(t *testing.T) { select { case msg := <-res: if msg == "" { - return + continue } t.Log(msg) } diff --git a/backend/data/stock_data_api.go b/backend/data/stock_data_api.go index db424ba..6ff9ab4 100644 --- a/backend/data/stock_data_api.go +++ b/backend/data/stock_data_api.go @@ -524,16 +524,8 @@ func SearchStockPriceInfo(stockCode string) *[]string { tasks = append(tasks, chromedp.Navigate(url)) tasks = append(tasks, chromedp.WaitVisible("div.quote-change-box", chromedp.ByQuery)) tasks = append(tasks, chromedp.ActionFunc(func(ctx context.Context) error { - chromedp.WaitVisible("span.quote-price", chromedp.ByQuery) - price := "" - for { - chromedp.Text("span.quote-price", &price, chromedp.BySearch).Do(ctx) - logger.SugaredLogger.Infof("price:%s", price) - if price != "" && validator.IsNumberStr(price) { - break - } - } - + price, _ := FetchPrice(ctx) + logger.SugaredLogger.Infof("price:%s", price) return nil })) tasks = append(tasks, chromedp.OuterHTML("html", &htmlContent, chromedp.ByQuery)) @@ -557,6 +549,29 @@ func SearchStockPriceInfo(stockCode string) *[]string { }) return &messages } +func FetchPrice(ctx context.Context) (string, error) { + var price string + timeout := time.After(10 * time.Second) // 设置超时时间为10秒 + ticker := time.NewTicker(1 * time.Second) // 每秒尝试一次 + defer ticker.Stop() + + for { + select { + case <-timeout: + return "", fmt.Errorf("timeout reached while fetching price") + case <-ticker.C: + err := chromedp.Run(ctx, chromedp.Text("span.quote-price", &price, chromedp.BySearch)) + if err != nil { + logger.SugaredLogger.Errorf("failed to fetch price: %v", err) + continue + } + logger.SugaredLogger.Infof("price:%s", price) + if price != "" && validator.IsNumberStr(price) { + return price, nil + } + } + } +} func SearchStockInfo(stock, msgType string) *[]string { // 创建一个 chromedp 上下文 ctx, cancel := chromedp.NewContext( @@ -570,8 +585,8 @@ func SearchStockInfo(stock, msgType string) *[]string { err := chromedp.Run(ctx, chromedp.Navigate(url), // 等待页面加载完成，可以根据需要调整等待时间 - //chromedp.Sleep(3*time.Second), - chromedp.WaitVisible("div.search-content,a.search-content", chromedp.ByQuery), + chromedp.Sleep(3*time.Second), + //chromedp.WaitVisible("a.search-content", chromedp.ByQuery), chromedp.OuterHTML("html", &htmlContent, chromedp.ByQuery), ) if err != nil { @@ -584,7 +599,46 @@ func SearchStockInfo(stock, msgType string) *[]string { return &[]string{} } var messages []string - document.Find("div.search-telegraph-list,a.search-content").Each(func(i int, selection *goquery.Selection) { + document.Find("a.search-content").Each(func(i int, selection *goquery.Selection) { + text := strutil.RemoveNonPrintable(selection.Text()) + if strings.Contains(text, stock) { + messages = append(messages, text) + logger.SugaredLogger.Infof("搜索到消息: %s", text) + } + }) + return &messages +} + +func SearchStockInfoByCode(stock string) *[]string { + // 创建一个 chromedp 上下文 + ctx, cancel := chromedp.NewContext( + context.Background(), + chromedp.WithLogf(logger.SugaredLogger.Infof), + chromedp.WithErrorf(logger.SugaredLogger.Errorf), + ) + defer cancel() + var htmlContent string + stock = strings.ReplaceAll(stock, "sh", "") + stock = strings.ReplaceAll(stock, "sz", "") + url := fmt.Sprintf("https://gushitong.baidu.com/stock/ab-%s", stock) + err := chromedp.Run(ctx, + chromedp.Navigate(url), + // 等待页面加载完成，可以根据需要调整等待时间 + //chromedp.Sleep(3*time.Second), + chromedp.WaitVisible("a.news-item-link", chromedp.ByQuery), + chromedp.OuterHTML("html", &htmlContent, chromedp.ByQuery), + ) + if err != nil { + logger.SugaredLogger.Error(err.Error()) + return &[]string{} + } + document, err := goquery.NewDocumentFromReader(strings.NewReader(htmlContent)) + if err != nil { + logger.SugaredLogger.Error(err.Error()) + return &[]string{} + } + var messages []string + document.Find("a.news-item-link").Each(func(i int, selection *goquery.Selection) { text := strutil.RemoveNonPrintable(selection.Text()) if strings.Contains(text, stock) { messages = append(messages, text) diff --git a/backend/data/stock_data_api_test.go b/backend/data/stock_data_api_test.go index 506d1e7..62dd016 100644 --- a/backend/data/stock_data_api_test.go +++ b/backend/data/stock_data_api_test.go @@ -50,6 +50,10 @@ func TestGetTelegraphSearch(t *testing.T) { //https://www.cls.cn/stock?code=sh600745 } +func TestSearchStockInfoByCode(t *testing.T) { + SearchStockInfoByCode("sh600745") +} + func TestSearchStockPriceInfo(t *testing.T) { SearchStockPriceInfo("sh600745") }