feat(data):增加对美股数据的支持

- 新增 getUSStockPriceInfo 函数用于获取美股实时行情信息
- 修改 SearchStockPriceInfo 函数,支持美股代码查询
- 更新 Tushare 数据接口,增加对美股每日数据的支持
- 优化股票代码处理逻辑,兼容不同市场代码格式
This commit is contained in:
ArvinLovegood 2025-02-28 17:38:48 +08:00
parent dbc25ca582
commit b2b0300aa1
5 changed files with 95 additions and 9 deletions

View File

@ -152,7 +152,11 @@ func (o OpenAi) NewChatStream(stock, stockCode, userQuestion string) <-chan map[
defer wg.Done() defer wg.Done()
endDate := time.Now().Format("20060102") endDate := time.Now().Format("20060102")
startDate := time.Now().Add(-time.Hour * time.Duration(24*o.KDays)).Format("20060102") startDate := time.Now().Add(-time.Hour * time.Duration(24*o.KDays)).Format("20060102")
K := NewTushareApi(getConfig()).GetDaily(ConvertStockCodeToTushareCode(stockCode), startDate, endDate, o.CrawlTimeOut) code := stockCode
if strutil.HasPrefixAny(stockCode, []string{"hk", "sz", "sh"}) {
code = ConvertStockCodeToTushareCode(stockCode)
}
K := NewTushareApi(getConfig()).GetDaily(code, startDate, endDate, o.CrawlTimeOut)
msg = append(msg, map[string]interface{}{ msg = append(msg, map[string]interface{}{
"role": "assistant", "role": "assistant",
"content": stock + "日K数据如下\n" + K, "content": stock + "日K数据如下\n" + K,

View File

@ -534,7 +534,7 @@ func ParseUSStockData(datas []string) (map[string]string, error) {
result["当前价格"] = parts[1] result["当前价格"] = parts[1]
result["日期"] = strutil.SplitAndTrim(parts[3], " ", "")[0] result["日期"] = strutil.SplitAndTrim(parts[3], " ", "")[0]
result["时间"] = strutil.SplitAndTrim(parts[3], " ", "")[1] result["时间"] = strutil.SplitAndTrim(parts[3], " ", "")[1]
logger.SugaredLogger.Infof("美股股票数据解析完成: %v", result) //logger.SugaredLogger.Infof("美股股票数据解析完成: %v", result)
return result, nil return result, nil
} }
@ -576,7 +576,7 @@ func ParseHKStockData(datas []string) (map[string]string, error) {
result["当前价格"] = parts[6] result["当前价格"] = parts[6]
result["日期"] = strings.ReplaceAll(parts[17], "/", "-") result["日期"] = strings.ReplaceAll(parts[17], "/", "-")
result["时间"] = strings.ReplaceAll(parts[18], "\";", ":00") result["时间"] = strings.ReplaceAll(parts[18], "\";", ":00")
logger.SugaredLogger.Infof("股票数据解析完成: %v", result) //logger.SugaredLogger.Infof("股票数据解析完成: %v", result)
return result, nil return result, nil
} }
@ -712,15 +712,70 @@ func GetRealTimeStockPriceInfo(ctx context.Context, stockCode string) (price, pr
func SearchStockPriceInfo(stockCode string, crawlTimeOut int64) *[]string { func SearchStockPriceInfo(stockCode string, crawlTimeOut int64) *[]string {
if strutil.HasPrefixAny(stockCode, []string{"HK", "hk"}) {
return getHKStockPriceInfo(stockCode, crawlTimeOut)
}
if strutil.HasPrefixAny(stockCode, []string{"SZ", "SH", "sh", "sz"}) { if strutil.HasPrefixAny(stockCode, []string{"SZ", "SH", "sh", "sz"}) {
return getSHSZStockPriceInfo(stockCode, crawlTimeOut) return getSHSZStockPriceInfo(stockCode, crawlTimeOut)
} }
if strutil.HasPrefixAny(stockCode, []string{"HK", "hk"}) {
return getHKStockPriceInfo(stockCode, crawlTimeOut)
}
if strutil.HasPrefixAny(stockCode, []string{"US", "us", "gb_"}) {
return getUSStockPriceInfo(stockCode, crawlTimeOut)
}
return &[]string{} return &[]string{}
} }
func getUSStockPriceInfo(stockCode string, crawlTimeOut int64) *[]string {
var messages []string
crawlerAPI := CrawlerApi{}
crawlerBaseInfo := CrawlerBaseInfo{
Name: "SinaCrawler",
Description: "SinaCrawler Crawler Description",
BaseUrl: "https://stock.finance.sina.com.cn",
Headers: map[string]string{"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 Edg/133.0.0.0"},
}
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(crawlTimeOut)*time.Second)
defer cancel()
crawlerAPI = crawlerAPI.NewCrawler(ctx, crawlerBaseInfo)
url := fmt.Sprintf("https://stock.finance.sina.com.cn/usstock/quotes/%s.html", strings.ReplaceAll(stockCode, "gb_", ""))
htmlContent, ok := crawlerAPI.GetHtml(url, "div#hqPrice", true)
if !ok {
return &[]string{}
}
document, err := goquery.NewDocumentFromReader(strings.NewReader(htmlContent))
if err != nil {
logger.SugaredLogger.Error(err.Error())
}
stockName := ""
stockPrice := ""
stockPriceTime := ""
document.Find("div.hq_title >h1").Each(func(i int, selection *goquery.Selection) {
stockName = strutil.RemoveNonPrintable(selection.Text())
logger.SugaredLogger.Infof("股票名称-:%s", stockName)
})
document.Find("#hqPrice").Each(func(i int, selection *goquery.Selection) {
stockPrice = strutil.RemoveNonPrintable(selection.Text())
logger.SugaredLogger.Infof("现价: %s", stockPrice)
})
document.Find("div.hq_time").Each(func(i int, selection *goquery.Selection) {
stockPriceTime = strutil.RemoveNonPrintable(selection.Text())
logger.SugaredLogger.Infof("时间: %s", stockPriceTime)
})
messages = append(messages, fmt.Sprintf("%s:%s现价%s", stockPriceTime, stockName, stockPrice))
logger.SugaredLogger.Infof("股票: %s", messages)
document.Find("div#hqDetails >table tbody tr").Each(func(i int, selection *goquery.Selection) {
text := strutil.RemoveNonPrintable(selection.Text())
logger.SugaredLogger.Infof("股票名称-%s: %s", stockName, text)
messages = append(messages, text)
})
return &messages
}
func getHKStockPriceInfo(stockCode string, crawlTimeOut int64) *[]string { func getHKStockPriceInfo(stockCode string, crawlTimeOut int64) *[]string {
var messages []string var messages []string
crawlerAPI := CrawlerApi{} crawlerAPI := CrawlerApi{}

View File

@ -47,7 +47,9 @@ func TestSearchStockInfoByCode(t *testing.T) {
func TestSearchStockPriceInfo(t *testing.T) { func TestSearchStockPriceInfo(t *testing.T) {
//SearchStockPriceInfo("hk06030", 30) //SearchStockPriceInfo("hk06030", 30)
SearchStockPriceInfo("sh600171", 30) //SearchStockPriceInfo("sh600171", 30)
SearchStockPriceInfo("gb_aapl", 30)
} }
func TestGetRealTimeStockPriceInfo(t *testing.T) { func TestGetRealTimeStockPriceInfo(t *testing.T) {

View File

@ -6,6 +6,7 @@ import (
"github.com/duke-git/lancet/v2/strutil" "github.com/duke-git/lancet/v2/strutil"
"github.com/go-resty/resty/v2" "github.com/go-resty/resty/v2"
"go-stock/backend/logger" "go-stock/backend/logger"
"strings"
"time" "time"
) )
@ -32,14 +33,15 @@ func (receiver TushareApi) GetDaily(tsCode, startDate, endDate string, crawlTime
fields := "ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount" fields := "ts_code,trade_date,open,high,low,close,pre_close,change,pct_chg,vol,amount"
resp := &TushareStockBasicResponse{} resp := &TushareStockBasicResponse{}
stockType := getStockType(tsCode) stockType := getStockType(tsCode)
logger.SugaredLogger.Debugf("tushare daily request: %s", stockType) tsCodeNEW := getTsCode(tsCode)
logger.SugaredLogger.Debugf("tushare daily request: %s,tsCode:%s,tsCodeNEW:%s", stockType, tsCode, tsCodeNEW)
_, err := receiver.client.SetTimeout(time.Duration(crawlTimeOut)*time.Second).R(). _, err := receiver.client.SetTimeout(time.Duration(crawlTimeOut)*time.Second).R().
SetHeader("content-type", "application/json"). SetHeader("content-type", "application/json").
SetBody(&TushareRequest{ SetBody(&TushareRequest{
ApiName: stockType, ApiName: stockType,
Token: receiver.config.TushareToken, Token: receiver.config.TushareToken,
Params: map[string]any{ Params: map[string]any{
"ts_code": tsCode, "ts_code": tsCodeNEW,
"start_date": startDate, "start_date": startDate,
"end_date": endDate, "end_date": endDate,
}, },
@ -68,6 +70,15 @@ func (receiver TushareApi) GetDaily(tsCode, startDate, endDate string, crawlTime
return res return res
} }
func getTsCode(code string) any {
if strutil.HasPrefixAny(code, []string{"US", "us", "gb_"}) {
code = strings.Replace(code, "gb_", "", 1)
code = strings.Replace(code, "us", "", 1)
return code
}
return code
}
func getStockType(code string) string { func getStockType(code string) string {
if strutil.HasSuffixAny(code, []string{"SZ", "SH", "sh", "sz"}) { if strutil.HasSuffixAny(code, []string{"SZ", "SH", "sh", "sz"}) {
return "daily" return "daily"
@ -75,5 +86,8 @@ func getStockType(code string) string {
if strutil.HasSuffixAny(code, []string{"HK", "hk"}) { if strutil.HasSuffixAny(code, []string{"HK", "hk"}) {
return "hk_daily" return "hk_daily"
} }
if strutil.HasPrefixAny(code, []string{"US", "us", "gb_"}) {
return "us_daily"
}
return "" return ""
} }

View File

@ -16,3 +16,14 @@ func TestGetDaily(t *testing.T) {
t.Log(res) t.Log(res)
} }
func TestGetUSDaily(t *testing.T) {
db.Init("../../data/stock.db")
tushareApi := NewTushareApi(getConfig())
res := tushareApi.GetDaily("gb_AAPL", "20250101", "20250217", 30)
t.Log(res)
//
}