82 lines
1.9 KiB
Go
82 lines
1.9 KiB
Go
package support
|
|
|
|
import (
|
|
"ai_scheduler/internal/config"
|
|
"ai_scheduler/internal/data/constants"
|
|
errorcode "ai_scheduler/internal/data/error"
|
|
"ai_scheduler/internal/pkg/util"
|
|
"ai_scheduler/internal/pkg/utils_vllm"
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
"github.com/gofiber/fiber/v2"
|
|
)
|
|
|
|
// HytAddressIngester 货易通地址提取实现
|
|
type HytAddressIngester struct {
|
|
cfg *config.Config
|
|
}
|
|
|
|
func NewHytAddressIngester(cfg *config.Config) *HytAddressIngester {
|
|
return &HytAddressIngester{cfg: cfg}
|
|
}
|
|
|
|
// Auth 鉴权逻辑
|
|
func (s *HytAddressIngester) Auth(c *fiber.Ctx) error {
|
|
// 读取头
|
|
token := strings.TrimSpace(c.Get("X-Source-Key"))
|
|
ts := strings.TrimSpace(c.Get("X-Timestamp"))
|
|
|
|
// 时间窗口校验
|
|
if ts != "" && !util.IsInTimeWindow(ts, 5*time.Minute) {
|
|
return errorcode.AuthNotFound
|
|
}
|
|
// token校验
|
|
if token == "" || token != constants.TokenAddressIngestHyt {
|
|
return errorcode.KeyNotFound
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Ingest 执行提取逻辑
|
|
func (s *HytAddressIngester) Ingest(ctx context.Context, text string) (*AddressIngestResp, error) {
|
|
// 模型调用
|
|
client, cleanup, err := utils_vllm.NewClient(s.cfg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer cleanup()
|
|
|
|
res, err := client.Chat(ctx, []*schema.Message{
|
|
{
|
|
Role: "system",
|
|
Content: constants.SystemPromptAddressIngestHyt,
|
|
},
|
|
{
|
|
Role: "user",
|
|
Content: text,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 解析模型返回结果
|
|
var addr AddressIngestResp
|
|
// 尝试直接解析
|
|
if err := json.Unmarshal([]byte(res.Content), &addr); err != nil {
|
|
// 修复json字符串
|
|
repairedContent, err := util.JSONRepair(res.Content)
|
|
if err != nil {
|
|
return nil, errorcode.ParamErrf("invalid response body: %v", err)
|
|
}
|
|
if err := json.Unmarshal([]byte(repairedContent), &addr); err != nil {
|
|
return nil, errorcode.ParamErrf("invalid response body: %v", err)
|
|
}
|
|
}
|
|
return &addr, nil
|
|
}
|