YogeLiu/doc_ai_backed
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add list api

Browse Source
This commit is contained in:
yoge
2025-12-18 12:39:50 +08:00
parent d06f2d9df1
commit 8a6da5b627
15 changed files with 133 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
api/v1/formula/handler.go
View File

@@ -9,6 +9,7 @@ import (
"gitea.com/bitwsd/document_ai/internal/service"
"gitea.com/bitwsd/document_ai/internal/storage/dao"
"gitea.com/bitwsd/document_ai/pkg/common"
"gitea.com/bitwsd/document_ai/pkg/constant"
"gitea.com/bitwsd/document_ai/pkg/utils"
"github.com/gin-gonic/gin"
@@ -37,11 +38,14 @@ func NewFormulaEndpoint() *FormulaEndpoint {
// @Router /v1/formula/recognition [post]
func (endpoint *FormulaEndpoint) CreateTask(ctx *gin.Context) {
var req formula.CreateFormulaRecognitionRequest
uid := ctx.GetInt64(constant.ContextUserID)
if err := ctx.BindJSON(&req); err != nil {
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeParamError, "Invalid parameters"))
return
}
req.UserID = uid
if !utils.InArray(req.TaskType, []string{string(dao.TaskTypeFormula), string(dao.TaskTypeFormula)}) {
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeParamError, "Invalid task type"))
return

10
api/v1/formula/router.go
View File

@@ -1,12 +1,16 @@
package formula
import (
"gitea.com/bitwsd/document_ai/pkg/common"
"github.com/gin-gonic/gin"
)
func SetupRouter(engine *gin.RouterGroup) {
endpoint := NewFormulaEndpoint()
engine.POST("/formula/recognition", endpoint.CreateTask)
engine.POST("/formula/ai_enhance", endpoint.AIEnhanceRecognition)
engine.GET("/formula/recognition/:task_no", endpoint.GetTaskStatus)
formulaRouter := engine.Group("/formula", common.GetAuthMiddleware())
{
formulaRouter.POST("/recognition", endpoint.CreateTask)
formulaRouter.POST("/ai_enhance", endpoint.AIEnhanceRecognition)
formulaRouter.GET("/recognition/:task_no", endpoint.GetTaskStatus)
}
}

4
api/v1/oss/handler.go
View File

@@ -11,7 +11,6 @@ import (
"gitea.com/bitwsd/document_ai/config"
"gitea.com/bitwsd/document_ai/internal/storage/dao"
"gitea.com/bitwsd/document_ai/pkg/common"
"gitea.com/bitwsd/document_ai/pkg/constant"
"gitea.com/bitwsd/document_ai/pkg/oss"
"gitea.com/bitwsd/document_ai/pkg/utils"
"github.com/gin-gonic/gin"
@@ -38,7 +37,6 @@ func GetPostObjectSignature(ctx *gin.Context) {
// @Failure 200 {object} common.Response "Error response"
// @Router /signature_url [get]
func GetSignatureURL(ctx *gin.Context) {
userID := ctx.GetInt64(constant.ContextUserID)
type Req struct {
FileHash string `json:"file_hash" binding:"required"`
FileName string `json:"file_name" binding:"required"`
@@ -51,7 +49,7 @@ func GetSignatureURL(ctx *gin.Context) {
}
taskDao := dao.NewRecognitionTaskDao()
sess := dao.DB.WithContext(ctx)
task, err := taskDao.GetTaskByFileURL(sess, userID, req.FileHash)
task, err := taskDao.GetTaskByFileURL(sess, req.FileHash)
if err != nil && err != gorm.ErrRecordNotFound {
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeDBError, "failed to get task"))
return

4
api/v1/task/handler.go
View File

@@ -3,10 +3,10 @@ package task
import (
"net/http"
"gitea.com/bitwsd/document_ai/pkg/log"
"gitea.com/bitwsd/document_ai/internal/model/task"
"gitea.com/bitwsd/document_ai/internal/service"
"gitea.com/bitwsd/document_ai/pkg/common"
"gitea.com/bitwsd/document_ai/pkg/log"
"github.com/gin-gonic/gin"
)
@@ -43,6 +43,8 @@ func (h *TaskEndpoint) GetTaskList(c *gin.Context) {
return
}
req.UserID = common.GetUserIDFromContext(c)
if req.Page <= 0 {
req.Page = 1
}

3
api/v1/task/router.go
View File

@@ -1,11 +1,12 @@
package task
import (
"gitea.com/bitwsd/document_ai/pkg/common"
"github.com/gin-gonic/gin"
)
func SetupRouter(engine *gin.RouterGroup) {
endpoint := NewTaskEndpoint()
engine.POST("/task/evaluate", endpoint.EvaluateTask)
engine.GET("/task/list", endpoint.GetTaskList)
engine.GET("/task/list", common.MustAuthMiddleware(), endpoint.GetTaskList)
}

9
api/v1/user/router.go
View File

@@ -9,10 +9,9 @@ func SetupRouter(router *gin.RouterGroup) {
userEndpoint := NewUserEndpoint()
userRouter := router.Group("/user")
{
userRouter.POST("/get/sms", userEndpoint.SendVerificationCode)
userRouter.POST("/login/phone", userEndpoint.LoginByPhoneCode)
userRouter.POST("/register/email", userEndpoint.RegisterByEmail)
userRouter.POST("/login/email", userEndpoint.LoginByEmail)
userRouter.GET("/info", common.GetAuthMiddleware(), userEndpoint.GetUserInfo)
userRouter.POST("/sms", userEndpoint.SendVerificationCode)
userRouter.POST("/register", userEndpoint.RegisterByEmail)
userRouter.POST("/login", userEndpoint.LoginByEmail)
userRouter.GET("/info", common.MustAuthMiddleware(), userEndpoint.GetUserInfo)
}
}

12
config/config_dev.yaml
View File

@@ -5,7 +5,7 @@ server:
database:
driver: mysql
host: mysql
port: 3306 # 容器内部端口，不是宿主机映射的 3006
port: 3306
username: root
password: texpixel#pwd123!
dbname: doc_ai
@@ -13,7 +13,7 @@ database:
max_open: 100
redis:
addr: redis:6379 # 容器内部端口，不是宿主机映射的 6079
addr: redis:6379
password: yoge@123321!
db: 0
@@ -22,7 +22,7 @@ limit:
log:
appName: document_ai
level: info # debug, info, warn, error
level: info
format: console # json, console
outputPath: ./logs/app.log # 日志文件路径
maxSize: 2 # 单个日志文件最大尺寸，单位MB
@@ -41,6 +41,6 @@ aliyun:
oss:
endpoint: oss-cn-beijing.aliyuncs.com
inner_endpoint: oss-cn-beijing-internal.aliyuncs.com
access_key_id: LTAI5tKogxeiBb4gJGWEePWN
access_key_secret: l4oCxtt5iLSQ1DAs40guTzKUfrxXwq
bucket_name: bitwsd-doc-ai
access_key_id: LTAI5t8qXhow6NCdYDtu1saF
access_key_secret: qZ2SwYsNCEBckCVSOszH31yYwXU44A
bucket_name: texpixel-doc

1
internal/model/formula/request.go
View File

@@ -5,6 +5,7 @@ type CreateFormulaRecognitionRequest struct {
FileHash string `json:"file_hash" binding:"required"` // file hash
FileName string `json:"file_name" binding:"required"` // file name
TaskType string `json:"task_type" binding:"required,oneof=FORMULA"` // task type
UserID int64 `json:"user_id"` // user id
}
type GetRecognitionStatusRequest struct {

18
internal/model/task/request.go
View File

@@ -11,12 +11,7 @@ type TaskListRequest struct {
TaskType string `json:"task_type" form:"task_type" binding:"required"`
Page int `json:"page" form:"page"`
PageSize int `json:"page_size" form:"page_size"`
}
type PdfInfo struct {
PageCount int `json:"page_count"`
PageWidth int `json:"page_width"`
PageHeight int `json:"page_height"`
UserID int64 `json:"-"`
}
type TaskListDTO struct {
@@ -26,11 +21,16 @@ type TaskListDTO struct {
Path string `json:"path"`
TaskType string `json:"task_type"`
CreatedAt string `json:"created_at"`
PdfInfo PdfInfo `json:"pdf_info"`
Latex string `json:"latex"`
Markdown string `json:"markdown"`
MathML string `json:"mathml"`
MathMLMW string `json:"mathml_mw"`
ImageBlob string `json:"image_blob"`
DocxURL string `json:"docx_url"`
PDFURL string `json:"pdf_url"`
}
type TaskListResponse struct {
TaskList []*TaskListDTO `json:"task_list"`
HasMore bool `json:"has_more"`
NextPage int `json:"next_page"`
Total int64 `json:"total"`
}

6
internal/service/recognition_service.go
View File

@@ -105,6 +105,7 @@ func (s *RecognitionService) CreateRecognitionTask(ctx context.Context, req *for
sess := dao.DB.WithContext(ctx)
taskDao := dao.NewRecognitionTaskDao()
task := &dao.RecognitionTask{
UserID: req.UserID,
TaskUUID: utils.NewUUID(),
TaskType: dao.TaskType(req.TaskType),
Status: dao.TaskStatusPending,
@@ -166,7 +167,8 @@ func (s *RecognitionService) GetFormualTask(ctx context.Context, taskNo string)
return nil, common.NewError(common.CodeDBError, "查询任务结果失败", err)
}
latex := taskRet.NewContentCodec().GetContent().(string)
return &formula.GetFormulaTaskResponse{TaskNo: taskNo, Latex: latex, Status: int(task.Status)}, nil
markdown := fmt.Sprintf("$$%s$$", latex)
return &formula.GetFormulaTaskResponse{TaskNo: taskNo, Latex: latex, Markdown: markdown, Status: int(task.Status)}, nil
}
func (s *RecognitionService) handleFormulaRecognition(ctx context.Context, taskID int64) error {
@@ -281,7 +283,7 @@ func (s *RecognitionService) processFormulaTask(ctx context.Context, taskID int6
headers := map[string]string{"Content-Type": "application/json", utils.RequestIDHeaderKey: utils.GetRequestIDFromContext(ctx)}
// 发送请求时会使用带超时的context
resp, err := s.httpClient.RequestWithRetry(ctx, http.MethodPost, "http://cloud.texpixel.com:1080/formula/predict", bytes.NewReader(jsonData), headers)
resp, err := s.httpClient.RequestWithRetry(ctx, http.MethodPost, "https://cloud.texpixel.com:10443/formula/predict", bytes.NewReader(jsonData), headers)
if err != nil {
if ctx.Err() == context.DeadlineExceeded {
log.Error(ctx, "func", "processFormulaTask", "msg", "请求超时")

52
internal/service/task.go
View File

@@ -3,25 +3,30 @@ package service
import (
"context"
"errors"
"fmt"
"strings"
"gitea.com/bitwsd/document_ai/pkg/log"
"gitea.com/bitwsd/document_ai/internal/model/task"
"gitea.com/bitwsd/document_ai/internal/storage/dao"
"gorm.io/gorm"
"gitea.com/bitwsd/document_ai/pkg/log"
)
type TaskService struct {
db *gorm.DB
recognitionTaskDao *dao.RecognitionTaskDao
evaluateTaskDao *dao.EvaluateTaskDao
recognitionResultDao *dao.RecognitionResultDao
}
func NewTaskService() *TaskService {
return &TaskService{dao.DB}
return &TaskService{
recognitionTaskDao: dao.NewRecognitionTaskDao(),
evaluateTaskDao: dao.NewEvaluateTaskDao(),
recognitionResultDao: dao.NewRecognitionResultDao(),
}
}
func (svc *TaskService) EvaluateTask(ctx context.Context, req *task.EvaluateTaskRequest) error {
taskDao := dao.NewRecognitionTaskDao()
task, err := taskDao.GetByTaskNo(svc.db.WithContext(ctx), req.TaskNo)
task, err := svc.recognitionTaskDao.GetByTaskNo(dao.DB.WithContext(ctx), req.TaskNo)
if err != nil {
log.Error(ctx, "func", "EvaluateTask", "msg", "get task by task no failed", "error", err)
return err
@@ -36,14 +41,13 @@ func (svc *TaskService) EvaluateTask(ctx context.Context, req *task.EvaluateTask
return errors.New("task not finished")
}
evaluateTaskDao := dao.NewEvaluateTaskDao()
evaluateTask := &dao.EvaluateTask{
TaskID: task.ID,
Satisfied: req.Satisfied,
Feedback: req.Feedback,
Comment: strings.Join(req.Suggestion, ","),
}
err = evaluateTaskDao.Create(svc.db.WithContext(ctx), evaluateTask)
err = svc.evaluateTaskDao.Create(dao.DB.WithContext(ctx), evaluateTask)
if err != nil {
log.Error(ctx, "func", "EvaluateTask", "msg", "create evaluate task failed", "error", err)
return err
@@ -53,19 +57,43 @@ func (svc *TaskService) EvaluateTask(ctx context.Context, req *task.EvaluateTask
}
func (svc *TaskService) GetTaskList(ctx context.Context, req *task.TaskListRequest) (*task.TaskListResponse, error) {
taskDao := dao.NewRecognitionTaskDao()
tasks, err := taskDao.GetTaskList(svc.db.WithContext(ctx), dao.TaskType(req.TaskType), req.Page, req.PageSize)
tasks, total, err := svc.recognitionTaskDao.GetTaskList(dao.DB.WithContext(ctx), req.UserID, dao.TaskType(req.TaskType), req.Page, req.PageSize)
if err != nil {
log.Error(ctx, "func", "GetTaskList", "msg", "get task list failed", "error", err)
return nil, err
}
taskIDs := make([]int64, 0, len(tasks))
for _, item := range tasks {
taskIDs = append(taskIDs, item.ID)
}
recognitionResults, err := svc.recognitionResultDao.GetByTaskIDs(dao.DB.WithContext(ctx), taskIDs)
if err != nil {
log.Error(ctx, "func", "GetTaskList", "msg", "get recognition results failed", "error", err)
return nil, err
}
recognitionResultMap := make(map[int64]*dao.RecognitionResult)
for _, item := range recognitionResults {
recognitionResultMap[item.TaskID] = item
}
resp := &task.TaskListResponse{
TaskList: make([]*task.TaskListDTO, 0, len(tasks)),
HasMore: false,
NextPage: 0,
Total: total,
}
for _, item := range tasks {
var latex string
var markdown string
recognitionResult := recognitionResultMap[item.ID]
if recognitionResult != nil {
latex = recognitionResult.NewContentCodec().GetContent().(string)
markdown = fmt.Sprintf("$$%s$$", latex)
}
resp.TaskList = append(resp.TaskList, &task.TaskListDTO{
Latex: latex,
Markdown: markdown,
TaskID: item.TaskUUID,
FileName: item.FileName,
Status: item.Status.String(),

5
internal/storage/dao/result.go
View File

@@ -84,6 +84,11 @@ func (dao *RecognitionResultDao) GetByTaskID(tx *gorm.DB, taskID int64) (result
return
}
func (dao *RecognitionResultDao) GetByTaskIDs(tx *gorm.DB, taskIDs []int64) (results []*RecognitionResult, err error) {
err = tx.Where("task_id IN (?)", taskIDs).Find(&results).Error
return
}
func (dao *RecognitionResultDao) Update(tx *gorm.DB, id int64, updates map[string]interface{}) error {
return tx.Model(&RecognitionResult{}).Where("id = ?", id).Updates(updates).Error
}

15
internal/storage/dao/task.go
View File

@@ -69,9 +69,9 @@ func (dao *RecognitionTaskDao) GetByTaskNo(tx *gorm.DB, taskUUID string) (task *
return
}
func (dao *RecognitionTaskDao) GetTaskByFileURL(tx *gorm.DB, userID int64, fileHash string) (task *RecognitionTask, err error) {
func (dao *RecognitionTaskDao) GetTaskByFileURL(tx *gorm.DB, fileHash string) (task *RecognitionTask, err error) {
task = &RecognitionTask{}
err = tx.Model(RecognitionTask{}).Where("user_id = ? AND file_hash = ?", userID, fileHash).First(task).Error
err = tx.Model(RecognitionTask{}).Where("file_hash = ?", fileHash).Last(task).Error
return
}
@@ -87,8 +87,13 @@ func (dao *RecognitionTaskDao) GetTaskByID(tx *gorm.DB, id int64) (task *Recogni
return task, nil
}
func (dao *RecognitionTaskDao) GetTaskList(tx *gorm.DB, taskType TaskType, page int, pageSize int) (tasks []*RecognitionTask, err error) {
func (dao *RecognitionTaskDao) GetTaskList(tx *gorm.DB, userID int64, taskType TaskType, page int, pageSize int) (tasks []*RecognitionTask, total int64, err error) {
offset := (page - 1) * pageSize
err = tx.Model(RecognitionTask{}).Where("task_type = ?", taskType).Offset(offset).Limit(pageSize).Order(clause.OrderByColumn{Column: clause.Column{Name: "id"}, Desc: true}).Find(&tasks).Error
return
query := tx.Model(RecognitionTask{}).Where("user_id = ? AND task_type = ?", userID, taskType)
err = query.Count(&total).Error
if err != nil {
return nil, 0, err
}
err = query.Offset(offset).Limit(pageSize).Order(clause.OrderByColumn{Column: clause.Column{Name: "id"}, Desc: true}).Find(&tasks).Error
return tasks, total, err
}

13
pkg/common/middleware.go
View File

@@ -45,6 +45,19 @@ func AuthMiddleware(ctx *gin.Context) {
ctx.Set(constant.ContextUserID, claims.UserId)
}
func MustAuthMiddleware() gin.HandlerFunc {
return func(ctx *gin.Context) {
token := ctx.GetHeader("Authorization")
if token != "" {
token = strings.TrimPrefix(token, "Bearer ")
claims, err := jwt.ParseToken(token)
if err == nil {
ctx.Set(constant.ContextUserID, claims.UserId)
}
}
}
}
func GetAuthMiddleware() gin.HandlerFunc {
return func(ctx *gin.Context) {
token := ctx.GetHeader("Authorization")

22
pkg/cors/cors.go
View File

@@ -19,9 +19,9 @@ type Config struct {
func DefaultConfig() Config {
return Config{
AllowOrigins: []string{"*"},
AllowMethods: []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
AllowHeaders: []string{"Origin", "Content-Type", "Accept"},
ExposeHeaders: []string{"Content-Length"},
AllowMethods: []string{"GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH"},
AllowHeaders: []string{"Origin", "Content-Type", "Accept", "Authorization", "X-Requested-With"},
ExposeHeaders: []string{"Content-Length", "Content-Type"},
AllowCredentials: true,
MaxAge: 86400, // 24 hours
}
@@ -30,16 +30,30 @@ func DefaultConfig() Config {
func Cors(config Config) gin.HandlerFunc {
return func(c *gin.Context) {
origin := c.Request.Header.Get("Origin")
if origin == "" {
c.Next()
return
}
// 检查是否允许该来源
allowOrigin := "*"
allowOrigin := ""
for _, o := range config.AllowOrigins {
if o == "*" {
// 通配符时，回显实际 origin（兼容 credentials）
allowOrigin = origin
break
}
if o == origin {
allowOrigin = origin
break
}
}
if allowOrigin == "" {
c.Next()
return
}
c.Header("Access-Control-Allow-Origin", allowOrigin)
c.Header("Access-Control-Allow-Methods", strings.Join(config.AllowMethods, ","))
c.Header("Access-Control-Allow-Headers", strings.Join(config.AllowHeaders, ","))