Files
rag-local/main.go
T
2026-08-05 10:28:44 +08:00

96 lines
2.8 KiB
Go

package main
import (
"context"
"os"
"os/signal"
"path/filepath"
"strings"
"syscall"
"time"
commonHttp "rag-local/common"
"rag-local/kb/controller"
"rag-local/kb/service"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/net/ghttp"
_ "github.com/gogf/gf/contrib/drivers/sqlite/v2"
_ "modernc.org/sqlite/vec"
)
func main() {
// ==================== API 路由(通过 RouteRegister 自动注册) ====================
commonHttp.RouteRegister([]interface{}{
controller.SystemConfig,
controller.ModelConfig,
controller.Dataset,
controller.Document,
controller.Chunk,
controller.ParseTask,
controller.Conversation,
controller.Message,
controller.KgEntity,
controller.KgRelation,
})
// ==================== Workspace 文件服务(源文件访问,路径穿越防护) ====================
commonHttp.Httpserver.BindHandler("/workspace/*", func(r *ghttp.Request) {
relPath := strings.TrimPrefix(r.URL.Path, "/workspace/")
if relPath == "" || strings.Contains(relPath, "..") {
r.Response.WriteStatus(403)
return
}
filePath := filepath.Join("workspace", relPath)
if _, err := os.Stat(filePath); os.IsNotExist(err) {
r.Response.WriteStatus(404)
return
}
r.Response.ServeFile(filePath)
})
// ==================== 前端静态资源服务(SPA,前后端合并部署) ====================
if st, err := os.Stat("ui-src/dist"); err == nil && st.IsDir() {
commonHttp.Httpserver.BindHandler("/*", func(r *ghttp.Request) {
path := strings.TrimPrefix(r.URL.Path, "/")
if strings.Contains(path, "..") {
r.Response.WriteStatus(404)
return
}
if path == "" || path == "index.html" {
path = "index.html"
}
filePath := filepath.Join("ui-src/dist", path)
if st, err := os.Stat(filePath); err == nil && !st.IsDir() {
r.Response.ServeFile(filePath)
return
}
r.Response.WriteStatus(404)
})
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
// ==================== 解析任务轮询(文档流水线消费端) ====================
service.ParseTaskService.StartParsePoller(ctx)
// ==================== 访问令牌(首次启动生成,打印到控制台) ====================
token, err := service.SystemConfigService.EnsureAccessToken(ctx)
if err != nil {
g.Log().Fatal(ctx, "ensure access token failed: %v", err)
}
g.Log().Infof(ctx, "============================================")
g.Log().Infof(ctx, "访问令牌(登录用): %s", token)
g.Log().Infof(ctx, "请在登录页输入上述令牌")
g.Log().Infof(ctx, "============================================")
g.Log().Info(ctx, "service started on :8080")
<-ctx.Done()
g.Log().Info(ctx, "shutting down...")
time.Sleep(3 * time.Second)
g.Log().Info(ctx, "bye")
}