refactor(stock-data): 重构股票数据获取逻辑

- 移除了不必要的并发请求,简化了代码结构
- 新增 FetchPrice 函数,用于获取股票价格信息
- 优化 SearchStockInfo 函数,提高了搜索效率和准确性
- 新增 SearchStockInfoByCode 函数,用于根据股票代码获取相关信息- 修复了一些潜在的错误和性能问题
This commit is contained in:
spark 2025-02-04 18:12:08 +08:00
parent 64b37b687c
commit b00bddcdec
4 changed files with 93 additions and 35 deletions

View File

@ -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{}{

View File

@ -13,7 +13,7 @@ func TestNewDeepSeekOpenAiConfig(t *testing.T) {
select {
case msg := <-res:
if msg == "" {
return
continue
}
t.Log(msg)
}

View File

@ -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)

View File

@ -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")
}