获取简介

main.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"genBrief/db"
 	"genBrief/util"
+	"strings"
 
 	"github.com/PuerkitoBio/goquery"
 	openai "github.com/sashabaranov/go-openai"
@@ -16,6 +17,7 @@ import (
 //go:embed dom.json
 var domJson string
 var domMap map[string]string
+var maxId int
 
 func init() {
 	domMap = map[string]string{}
@@ -24,13 +26,16 @@ func init() {
 		panic(err)
 	}
 	db.Init()
+	db.New("t_config").Attr("value").Where("name", "prev_loop_new_id").GetRow().Scan(&maxId)
+	fmt.Println("startTask at maxId ", maxId)
 }
 
 func main() {
-	maxId := 4000
 	for {
 		newMaxId := loopMaxId(maxId)
-		if maxId == newMaxId {
+		if maxId >= newMaxId {
+			fmt.Println("finishTask at maxId ", maxId)
+			db.Pool().Update("t_config", map[string]interface{}{"value": maxId}, map[string]interface{}{"name": "prev_loop_new_id"})
 			return
 		}
 		maxId = newMaxId
@@ -69,6 +74,12 @@ func loopMaxId(maxId int) int {
 			update["status"] = "Delete"
 		} else {
 			update["status"] = "Picture"
+			tags := genTag(brief)
+			if tags != "" {
+				alltag := do_insert_tag(id, tags)
+				update["tag"] = alltag
+				update["pics"] = get_pics(alltag)
+			}
 		}
 		fmt.Println("finish", id, brief)
 		if _, err := db.Pool().Update("t_news", update, map[string]interface{}{"id": id}); err != nil {
@@ -99,6 +110,62 @@ func genBrief(content string) string {
 	return resp.Choices[0].Message.Content
 }
 
+func genTag(content string) string {
+	client := openai.NewClient("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")
+	resp, err := client.CreateChatCompletion(
+		context.Background(),
+		openai.ChatCompletionRequest{
+			Model: openai.GPT3Dot5Turbo,
+			Messages: []openai.ChatCompletionMessage{
+				{
+					Role:    openai.ChatMessageRoleUser,
+					Content: content + "\r\n以上新闻内容属于哪一类新闻 A居民 B商业 C金融 D建筑 E屋内装饰 \r\n可以选一项或者两项",
+				},
+			},
+		},
+	)
+	if err != nil {
+		fmt.Printf("ChatCompletion error: %v\n", err)
+		return ""
+	}
+	return resp.Choices[0].Message.Content
+}
+
+func get_pics(tags string) string {
+	tag := strings.Split(tags, ",")[0]
+	if tag == "" || tag == "other" {
+		tag = "gpt"
+	}
+	var url string
+	err := db.New("t_news_img").Attr("url").Where("tag", tag).Order("RDER BY RAND()").GetRow().Scan(&url)
+	if err != nil {
+		fmt.Println("get_pics err", tag, err.Error())
+	}
+	return url
+}
+
+func do_insert_tag(id int64, tags string) string {
+	mtag := map[string]string{
+		"A": "residential",
+		"B": "commercial",
+		"C": "financial",
+		"D": "construction",
+		"E": "indoor",
+	}
+	mtags := []string{}
+	for opt, tag := range mtag {
+		if strings.Contains(tags, opt) {
+			record := map[string]interface{}{
+				"new_id": id,
+				"tag":    tag,
+			}
+			db.Pool().Insert("t_news_tag", record)
+			mtags = append(mtags, tag)
+		}
+	}
+	return strings.Join(mtags, ",")
+}
+
 func getContent(body []byte, dom string) string {
 	doc, _ := goquery.NewDocumentFromReader(bytes.NewReader(body))
 	doc.Find("script").Remove()