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

feat: add new path for recognize

Browse Source
This commit is contained in:
yogeliu
2025-12-31 17:53:12 +08:00
parent be1047618e
commit 2e158d3fee
2 changed files with 30 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
internal/model/formula/response.go
View File

@@ -22,3 +22,10 @@ type GetFormulaTaskResponse struct {
type FormulaRecognitionResponse struct {
Result string `json:"result"`
}
// ImageOCRResponse 图片OCR接口返回的响应
type ImageOCRResponse struct {
Latex string `json:"latex"` // LaTeX 格式内容
Markdown string `json:"markdown"` // Markdown 格式内容
MathML string `json:"mathml"` // MathML 格式（无公式时为空）
}

23
internal/service/recognition_service.go
View File

@@ -22,6 +22,7 @@ import (
"gitea.com/bitwsd/document_ai/pkg/constant"
"gitea.com/bitwsd/document_ai/pkg/httpclient"
"gitea.com/bitwsd/document_ai/pkg/oss"
"gitea.com/bitwsd/document_ai/pkg/requestid"
"gitea.com/bitwsd/document_ai/pkg/utils"
"gorm.io/gorm"
)
@@ -511,8 +512,8 @@ func (s *RecognitionService) processFormulaTask(ctx context.Context, taskID int6
// 设置Content-Type头为application/json
headers := map[string]string{"Content-Type": "application/json", utils.RequestIDHeaderKey: utils.GetRequestIDFromContext(ctx)}
// 发送请求时会使用带超时的context
resp, err := s.httpClient.RequestWithRetry(ctx, http.MethodPost, "https://cloud.texpixel.com:10443/vlm/formula/predict", bytes.NewReader(jsonData), headers)
// 发送请求到新的 OCR 接口
resp, err := s.httpClient.RequestWithRetry(ctx, http.MethodPost, "https://cloud.texpixel.com:10443/doc_process/v1/image/ocr", bytes.NewReader(jsonData), headers)
if err != nil {
if ctx.Err() == context.DeadlineExceeded {
log.Error(ctx, "func", "processFormulaTask", "msg", "请求超时")
@@ -532,12 +533,18 @@ func (s *RecognitionService) processFormulaTask(ctx context.Context, taskID int6
log.Info(ctx, "func", "processFormulaTask", "msg", "响应内容", "body", body.String())
// 解析 JSON 响应
var formulaResp formula.FormulaRecognitionResponse
if err := json.Unmarshal(body.Bytes(), &formulaResp); err != nil {
var ocrResp formula.ImageOCRResponse
if err := json.Unmarshal(body.Bytes(), &ocrResp); err != nil {
log.Error(ctx, "func", "processFormulaTask", "msg", "解析响应JSON失败", "error", err)
return err
}
err = resultDao.Create(tx, dao.RecognitionResult{TaskID: taskID, TaskType: dao.TaskTypeFormula, Latex: formulaResp.Result})
err = resultDao.Create(tx, dao.RecognitionResult{
TaskID: taskID,
TaskType: dao.TaskTypeFormula,
Latex: ocrResp.Latex,
Markdown: ocrResp.Markdown,
MathML: ocrResp.MathML,
})
if err != nil {
log.Error(ctx, "func", "processFormulaTask", "msg", "保存任务结果失败", "error", err)
return err
@@ -705,15 +712,19 @@ func (s *RecognitionService) processOneTask(ctx context.Context) {
}
ctx = context.WithValue(ctx, utils.RequestIDKey, task.TaskUUID)
// 使用 gls 设置 request_id，确保在整个任务处理过程中可用
requestid.SetRequestID(task.TaskUUID, func() {
log.Info(ctx, "func", "processFormulaQueue", "msg", "获取任务成功", "task_id", taskID)
err = s.processBaiduOCRTask(ctx, taskID, task.FileURL)
err = s.processFormulaTask(ctx, taskID, task.FileURL)
if err != nil {
log.Error(ctx, "func", "processFormulaQueue", "msg", "处理任务失败", "error", err)
return
}
log.Info(ctx, "func", "processFormulaQueue", "msg", "处理任务成功", "task_id", taskID)
})
}
// processMathpixTask 使用 Mathpix API 处理公式识别任务（用于增强识别）