rag_eta_report_abstract.go 8.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317
  1. package elastic
  2. import (
  3. "context"
  4. "encoding/json"
  5. "eta/eta_api/models/rag"
  6. "eta/eta_api/utils"
  7. "fmt"
  8. "github.com/olivere/elastic/v7"
  9. "strings"
  10. "time"
  11. )
  12. // 摘要索引
  13. var EsRagEtaReportAbstractName = utils.EsRagEtaReportAbstractName
  14. type RagEtaReportAbstractItem struct {
  15. RagEtaReportAbstractId int `gorm:"primaryKey;column:rag_eta_report_abstract_id" description:"-"`
  16. RagEtaReportId int `gorm:"column:rag_eta_report_id" description:"ETA报告id"`
  17. Abstract string `gorm:"column:abstract;type:longtext;comment:摘要内容;" description:"摘要内容"`
  18. QuestionId int `gorm:"column:question_id" description:"提示词Id"`
  19. Version int `gorm:"column:version" description:"版本号"`
  20. VectorKey string `gorm:"column:vector_key" description:"向量key标识"`
  21. ModifyTime time.Time `gorm:"column:modify_time" description:"modifyTime"`
  22. CreateTime time.Time `gorm:"column:create_time" description:"createTime"`
  23. Title string `gorm:"column:title;type:varchar(255);comment:标题;" description:"标题"`
  24. TagIdList []int `description:"品种id列表"`
  25. TagNameList []string `description:"品种名称列表"`
  26. }
  27. func (m *RagEtaReportAbstractItem) ToView() rag.RagEtaReportAbstractView {
  28. var modifyTime, createTime string
  29. if !m.CreateTime.IsZero() {
  30. createTime = m.CreateTime.Format(utils.FormatDateTime)
  31. }
  32. if !m.ModifyTime.IsZero() {
  33. modifyTime = m.ModifyTime.Format(utils.FormatDateTime)
  34. }
  35. //tagId := 0
  36. //if len(m.TagIdList) > 0 {
  37. // tagId = m.TagIdList[0]
  38. //}
  39. tagsName := ``
  40. if len(m.TagNameList) > 0 {
  41. tagsName = strings.Join(m.TagNameList, `,`)
  42. }
  43. return rag.RagEtaReportAbstractView{
  44. RagEtaReportAbstractId: m.RagEtaReportAbstractId,
  45. RagEtaReportId: m.RagEtaReportId,
  46. Abstract: m.Abstract,
  47. QuestionId: m.QuestionId,
  48. //QuestionContent: m.,
  49. Version: m.Version,
  50. //Tags: m.Ta,
  51. VectorKey: m.VectorKey,
  52. ModifyTime: modifyTime,
  53. CreateTime: createTime,
  54. Title: m.Title,
  55. TagsName: tagsName,
  56. }
  57. }
  58. func (m *RagEtaReportAbstractItem) ToViewList(list []*RagEtaReportAbstractItem) (wechatArticleViewList []rag.RagEtaReportAbstractView) {
  59. wechatArticleViewList = make([]rag.RagEtaReportAbstractView, 0)
  60. for _, v := range list {
  61. wechatArticleViewList = append(wechatArticleViewList, v.ToView())
  62. }
  63. return
  64. }
  65. // RagEtaReportEsAddOrEdit
  66. // @Description: 新增/编辑微信文章
  67. // @author: Roc
  68. // @datetime 2025-03-13 10:24:05
  69. // @param docId string
  70. // @param item RagEtaReportAndPlatform
  71. // @return err error
  72. func RagEtaReportAbstractEsAddOrEdit(docId string, item RagEtaReportAbstractItem) (err error) {
  73. if docId == "" {
  74. return
  75. }
  76. if EsRagEtaReportAbstractName == `` {
  77. return
  78. }
  79. defer func() {
  80. if err != nil {
  81. fmt.Println("RagEtaReportEsAddOrEdit Err:", err.Error())
  82. }
  83. }()
  84. client := utils.EsClient
  85. resp, err := client.Index().Index(EsRagEtaReportAbstractName).Id(docId).BodyJson(item).Refresh("true").Do(context.Background())
  86. if err != nil {
  87. fmt.Println("新增失败:", err.Error())
  88. return err
  89. }
  90. if resp.Status == 0 {
  91. fmt.Println("新增成功", resp.Result)
  92. err = nil
  93. } else {
  94. fmt.Println("RagEtaReportEsAddOrEdit", resp.Status, resp.Result)
  95. }
  96. return
  97. }
  98. // RagEtaReportEsDel
  99. // @Description: 删除微信文章
  100. // @author: Roc
  101. // @datetime 2025-03-13 10:23:55
  102. // @param docId string
  103. // @return err error
  104. func RagEtaReportAbstractEsDel(docId string) (err error) {
  105. if docId == "" {
  106. return
  107. }
  108. if EsRagEtaReportAbstractName == `` {
  109. return
  110. }
  111. defer func() {
  112. if err != nil {
  113. fmt.Println("EsDeleteEdbInfoData Err:", err.Error())
  114. }
  115. }()
  116. client := utils.EsClient
  117. resp, err := client.Delete().Index(EsRagEtaReportAbstractName).Id(docId).Refresh(`true`).Do(context.Background())
  118. if err != nil {
  119. return
  120. }
  121. if resp.Status == 0 {
  122. fmt.Println("删除成功")
  123. } else {
  124. fmt.Println("RagEtaReportEsDel", resp.Status, resp.Result)
  125. }
  126. return
  127. }
  128. // RagEtaReportAbstractEsSearch
  129. // @Description: 搜索
  130. // @author: Roc
  131. // @datetime 2025-03-13 19:54:54
  132. // @param keywordStr string
  133. // @param tagIdList []int
  134. // @param from int
  135. // @param size int
  136. // @param sortMap map[string]string
  137. // @return total int64
  138. // @return list []*RagEtaReportAbstractItem
  139. // @return err error
  140. func RagEtaReportAbstractEsSearch(keywordStr string, tagIdList []int, questionId, from, size int, sortMap map[string]string) (total int64, list []*RagEtaReportAbstractItem, err error) {
  141. indexName := EsRagEtaReportAbstractName
  142. list = make([]*RagEtaReportAbstractItem, 0)
  143. defer func() {
  144. if err != nil {
  145. fmt.Println("RagEtaReportAbstractEsSearch Err:", err.Error())
  146. }
  147. }()
  148. query := elastic.NewBoolQuery()
  149. if len(tagIdList) > 0 {
  150. termsList := make([]interface{}, 0)
  151. for _, v := range tagIdList {
  152. termsList = append(termsList, v)
  153. }
  154. query = query.Must(elastic.NewTermsQuery("TagIdList", termsList...))
  155. }
  156. // 提示词id
  157. if questionId > 0 {
  158. query = query.Must(elastic.NewTermsQuery("QuestionId", questionId))
  159. }
  160. // 名字匹配
  161. if keywordStr != `` {
  162. query = query.Must(elastic.NewMultiMatchQuery(keywordStr, "Abstract"))
  163. }
  164. // 排序
  165. sortList := make([]*elastic.FieldSort, 0)
  166. // 如果没有关键字,那么就走指标id倒序
  167. for orderKey, orderType := range sortMap {
  168. switch orderType {
  169. case "asc":
  170. sortList = append(sortList, elastic.NewFieldSort(orderKey).Asc())
  171. case "desc":
  172. sortList = append(sortList, elastic.NewFieldSort(orderKey).Desc())
  173. }
  174. }
  175. return searchRagEtaReportAbstract(indexName, query, sortList, from, size)
  176. }
  177. // searchEdbInfoDataV2 查询es中的数据
  178. func searchRagEtaReportAbstract(indexName string, query elastic.Query, sortList []*elastic.FieldSort, from, size int) (total int64, list []*RagEtaReportAbstractItem, err error) {
  179. total, err = searchRagEtaReportAbstractTotal(indexName, query)
  180. if err != nil {
  181. return
  182. }
  183. // 获取列表数据
  184. list, err = searchRagEtaReportAbstractList(indexName, query, sortList, from, size)
  185. if err != nil {
  186. return
  187. }
  188. return
  189. }
  190. // searchEdbInfoDataTotal
  191. // @Description: 查询es中的数量
  192. // @author: Roc
  193. // @datetime 2024-12-23 11:19:04
  194. // @param indexName string
  195. // @param query elastic.Query
  196. // @return total int64
  197. // @return err error
  198. func searchRagEtaReportAbstractTotal(indexName string, query elastic.Query) (total int64, err error) {
  199. defer func() {
  200. if err != nil {
  201. fmt.Println("searchRagEtaReportAbstractTotal Err:", err.Error())
  202. }
  203. }()
  204. client := utils.EsClient
  205. //根据条件数量统计
  206. requestTotalHits := client.Count(indexName).Query(query)
  207. total, err = requestTotalHits.Do(context.Background())
  208. if err != nil {
  209. return
  210. }
  211. return
  212. }
  213. // searchEdbInfoDataList
  214. // @Description: 查询es中的明细数据
  215. // @author: Roc
  216. // @datetime 2024-12-23 11:18:48
  217. // @param indexName string
  218. // @param query elastic.Query
  219. // @param sortList []*elastic.FieldSort
  220. // @param from int
  221. // @param size int
  222. // @return list []*data_manage.EdbInfoList
  223. // @return err error
  224. func searchRagEtaReportAbstractList(indexName string, query elastic.Query, sortList []*elastic.FieldSort, from, size int) (list []*RagEtaReportAbstractItem, err error) {
  225. list = make([]*RagEtaReportAbstractItem, 0)
  226. defer func() {
  227. if err != nil {
  228. fmt.Println("searchRagEtaReportAbstractList Err:", err.Error())
  229. }
  230. }()
  231. client := utils.EsClient
  232. // 高亮
  233. highlight := elastic.NewHighlight()
  234. highlight = highlight.Fields(elastic.NewHighlighterField("Content"))
  235. highlight = highlight.PreTags("<font color='red'>").PostTags("</font>")
  236. //request := client.Search(indexName).Highlight(highlight).From(from).Size(size) // sets the JSON request
  237. request := client.Search(indexName).From(from).Size(size) // sets the JSON request
  238. // 如果有指定排序,那么就按照排序来
  239. if len(sortList) > 0 {
  240. for _, v := range sortList {
  241. request = request.SortBy(v)
  242. }
  243. }
  244. searchMap := make(map[string]string)
  245. searchResp, err := request.Query(query).Do(context.Background())
  246. if err != nil {
  247. return
  248. }
  249. //fmt.Println(searchResp)
  250. //fmt.Println(searchResp.Status)
  251. if searchResp.Status != 0 {
  252. return
  253. }
  254. //total = searchResp.TotalHits()
  255. if searchResp.Hits != nil {
  256. for _, v := range searchResp.Hits.Hits {
  257. if _, ok := searchMap[v.Id]; !ok {
  258. itemJson, tmpErr := v.Source.MarshalJSON()
  259. if tmpErr != nil {
  260. err = tmpErr
  261. fmt.Println("movieJson err:", err)
  262. return
  263. }
  264. item := new(RagEtaReportAbstractItem)
  265. tmpErr = json.Unmarshal(itemJson, &item)
  266. if tmpErr != nil {
  267. fmt.Println("json.Unmarshal movieJson err:", tmpErr)
  268. err = tmpErr
  269. return
  270. }
  271. if len(v.Highlight["Content"]) > 0 {
  272. item.Abstract = v.Highlight["Content"][0]
  273. }
  274. list = append(list, item)
  275. searchMap[v.Id] = v.Id
  276. }
  277. }
  278. }
  279. return
  280. }