diff --git a/config/config.yaml b/config/config.yaml index 4baff3f..c85ac02 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -47,4 +47,6 @@ tools: base_url: "https://gateway.dev.cdlsxd.cn/zltx_api/admin/direct/ai/search/" enabled: true api_key: "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJ1c2VyQ2VudGVyIiwiZXhwIjoxNzU2MTgyNTM1LCJuYmYiOjE3NTYxODA3MzUsImp0aSI6IjEiLCJQaG9uZSI6IjE4MDAwMDAwMDAwIiwiVXNlck5hbWUiOiJsc3hkIiwiUmVhbE5hbWUiOiLotoXnuqfnrqHnkIblkZgiLCJBY2NvdW50VHlwZSI6MSwiR3JvdXBDb2RlcyI6IlZDTF9DQVNISUVSLFZDTF9PUEVSQVRFLFZDTF9BRE1JTixWQ0xfQUFBLFZDTF9WQ0xfT1BFUkFULFZDTF9JTlZPSUNFLENSTV9BRE1JTixMSUFOTElBTl9BRE1JTixNQVJLRVRNQUcyX0FETUlOLFBIT05FQklMTF9BRE1JTixRSUFOWkhVX1NVUFBFUl9BRE0sTUFSS0VUSU5HU0FBU19TVVBFUkFETUlOLENBUkRfQ09ERSxDQVJEX1BST0NVUkVNRU5ULE1BUktFVElOR1NZU1RFTV9TVVBFUixTVEFUSVNUSUNBTFNZU1RFTV9BRE1JTixaTFRYX0FETUlOLFpMVFhfT1BFUkFURSIsIkRpbmdVc2VySWQiOiIxNjIwMjYxMjMwMjg5MzM4MzQifQ.N1xv1PYbcO8_jR5adaczc16YzGsr4z101gwEZdulkRaREBJNYTOnFrvRxTFx3RJTooXsqTqroE1MR84v_1WPX6BS6kKonA-kC1Jgot6yrt5rFWhGNGb2Cpr9rKIFCCQYmiGd3AUgDazEeaQ0_sodv3E-EXg9VfE1SX8nMcck9Yjnc8NCy7RTWaBIaSeOdZcEl-JfCD0S6GSx3oErp_hk-U9FKGwf60wAuDGTY1R0BP4BYpcEqS-C2LSnsSGyURi54Cuk5xH8r1WuF0Dm5bwAj5d7Hvs77-N_sUF-C5ONqyZJRAEhYLgcmN9RX_WQZfizdQJxizlTczdpzYfy-v-1eQ" - + knowledge: + base_url: "http://117.175.169.61:10000" + enabled: true diff --git a/go.mod b/go.mod index 7768c99..4b59108 100644 --- a/go.mod +++ b/go.mod @@ -6,13 +6,12 @@ toolchain go1.24.7 require ( gitea.cdlsxd.cn/self-tools/l_request v1.0.8 - github.com/bytedance/sonic v1.11.6 github.com/emirpasic/gods v1.18.1 github.com/go-kratos/kratos/v2 v2.9.1 github.com/gofiber/fiber/v2 v2.52.9 github.com/gofiber/websocket/v2 v2.2.1 github.com/google/wire v0.7.0 - github.com/ollama/ollama v0.11.11 + github.com/ollama/ollama v0.12.0 github.com/redis/go-redis/v9 v9.14.0 github.com/spf13/viper v1.17.0 github.com/tmc/langchaingo v0.1.13 @@ -27,10 +26,7 @@ require ( require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/andybalholm/brotli v1.1.0 // indirect - github.com/bytedance/sonic/loader v0.1.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/cloudwego/base64x v0.1.4 // indirect - github.com/cloudwego/iasm v0.2.0 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dlclark/regexp2 v1.11.4 // indirect github.com/fasthttp/websocket v1.5.3 // indirect @@ -41,7 +37,6 @@ require ( github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/klauspost/compress v1.17.9 // indirect - github.com/klauspost/cpuid/v2 v2.2.7 // indirect github.com/magiconair/properties v1.8.7 // indirect github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect @@ -59,13 +54,11 @@ require ( github.com/spf13/pflag v1.0.5 // indirect github.com/stretchr/testify v1.11.1 // indirect github.com/subosito/gotenv v1.6.0 // indirect - github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/valyala/fasthttp v1.51.0 // indirect github.com/valyala/tcplisten v1.0.0 // indirect go.uber.org/atomic v1.9.0 // indirect go.uber.org/multierr v1.9.0 // indirect - golang.org/x/arch v0.8.0 // indirect golang.org/x/crypto v0.36.0 // indirect golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa // indirect golang.org/x/sys v0.31.0 // indirect diff --git a/go.sum b/go.sum index 916bf1c..fa9c4ef 100644 --- a/go.sum +++ b/go.sum @@ -50,10 +50,6 @@ github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= -github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc0= -github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= -github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM= -github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= @@ -61,10 +57,6 @@ github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWR github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= -github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y= -github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= -github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg= -github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= @@ -176,10 +168,6 @@ github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/X github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/klauspost/compress v1.17.9 h1:6KIumPrER1LHsvBVuDa0r5xaG0Es51mhhB9BQB2qeMA= github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw= -github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= -github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM= -github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= -github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= @@ -199,8 +187,8 @@ github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6T github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= -github.com/ollama/ollama v0.11.11 h1:mErMiUGclp47rCDbSUmBiY2L76EpT0uIYRZVBO6qg/k= -github.com/ollama/ollama v0.11.11/go.mod h1:9+1//yWPsDE2u+l1a5mpaKrYw4VdnSsRU3ioq5BvMms= +github.com/ollama/ollama v0.12.0 h1:BRry7G2Skz7Mu+E6rz40tzBXNbLTEhheGT8umc1zvxo= +github.com/ollama/ollama v0.12.0/go.mod h1:9+1//yWPsDE2u+l1a5mpaKrYw4VdnSsRU3ioq5BvMms= github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -244,7 +232,6 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= -github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= @@ -253,8 +240,6 @@ github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU= github.com/tmc/langchaingo v0.1.13 h1:rcpMWBIi2y3B90XxfE4Ao8dhCQPVDMaNPnN5cGB1CaA= github.com/tmc/langchaingo v0.1.13/go.mod h1:vpQ5NOIhpzxDfTZK9B6tf2GM/MoaHewPWM5KXXGh7hg= -github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= -github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= github.com/valyala/fasthttp v1.51.0 h1:8b30A5JlZ6C7AS81RsWjYMQmrZG6feChmgAolCl1SqA= @@ -275,9 +260,6 @@ go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE= go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= -golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= -golang.org/x/arch v0.8.0 h1:3wRIsP3pM4yUptoR96otTUOXI367OS0+c9eeRi9doIc= -golang.org/x/arch v0.8.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= @@ -414,7 +396,6 @@ golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= @@ -603,9 +584,7 @@ honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWh honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg= honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k= honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k= -nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8= -rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA= sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo= diff --git a/internal/biz/router.go b/internal/biz/router.go index dd9c910..6e5fccb 100644 --- a/internal/biz/router.go +++ b/internal/biz/router.go @@ -10,7 +10,7 @@ import ( "ai_scheduler/internal/pkg" "ai_scheduler/internal/pkg/mapstructure" "ai_scheduler/internal/pkg/utils_ollama" - "ai_scheduler/internal/tools" + tools "ai_scheduler/internal/tools" "ai_scheduler/tmpl/dataTemp" "context" "encoding/json" @@ -145,7 +145,7 @@ func (r *AiRouterBiz) RouteWithSocket(c *websocket.Conn, req *entitys.ChatSockRe } }() defer close(ch) - err = r.handleMatch(c, ch, &matchJson, task) + err = r.handleMatch(c, ch, &matchJson, task, sysInfo) if err != nil { return } @@ -169,7 +169,7 @@ func (r *AiRouterBiz) handleOtherTask(c *websocket.Conn, ch chan entitys.Respons return } -func (r *AiRouterBiz) handleMatch(c *websocket.Conn, ch chan entitys.ResponseData, matchJson *entitys.Match, tasks []model.AiTask) (err error) { +func (r *AiRouterBiz) handleMatch(c *websocket.Conn, ch chan entitys.ResponseData, matchJson *entitys.Match, tasks []model.AiTask, sysInfo model.AiSy) (err error) { if !matchJson.IsMatch { ch <- entitys.ResponseData{ @@ -195,6 +195,8 @@ func (r *AiRouterBiz) handleMatch(c *websocket.Conn, ch chan entitys.ResponseDat return r.handleApiTask(ch, c, matchJson, pointTask) case constants.TaskTypeFunc: return r.handleTask(ch, c, matchJson, pointTask) + case constants.TaskTypeKnowle: + return r.handleKnowle(ch, c, matchJson, pointTask, sysInfo) default: return r.handleOtherTask(c, ch, matchJson) } @@ -215,6 +217,90 @@ func (r *AiRouterBiz) handleTask(channel chan entitys.ResponseData, c *websocket return } +// 知识库 +func (r *AiRouterBiz) handleKnowle(channel chan entitys.ResponseData, c *websocket.Conn, matchJson *entitys.Match, task *model.AiTask, sysInfo model.AiSy) (err error) { + + var ( + configData entitys.ConfigDataTool + sessionIdKnowledge string + query string + host string + ) + err = json.Unmarshal([]byte(task.Config), &configData) + if err != nil { + return + } + + // 通过session 找到知识库session + session := c.Headers("X-Session", "") + if len(session) == 0 { + return errors.SessionNotFound + } + sessionInfo, has, err := r.sessionImpl.FindOne(r.sessionImpl.WithSessionId(session)) + if err != nil { + return + } else if !has { + return errors.SessionNotFound + } + + // 找到知识库的host + { + tool, exists := r.toolManager.GetTool(configData.Tool) + if !exists { + return fmt.Errorf("tool not found: %s", configData.Tool) + } + + if knowledgeTool, ok := tool.(*tools.KnowledgeBaseTool); !ok { + return fmt.Errorf("未找到知识库Tool: %s", configData.Tool) + } else { + host = knowledgeTool.GetConfig().BaseURL + } + + } + + // 知识库的session为空，请求知识库获取, 并绑定 + if sessionInfo.KnowlegeSessionID == "" { + // 请求知识库 + if sessionIdKnowledge, err = tools.GetKnowledgeBaseSession(host, sysInfo.KnowlegeBaseID, sysInfo.KnowlegeTenantKey); err != nil { + return + } + + // 绑定知识库session，下次可以使用 + sessionInfo.KnowlegeSessionID = sessionIdKnowledge + if err = r.sessionImpl.Update(&sessionInfo, r.sessionImpl.WithSessionId(sessionInfo.SessionID)); err != nil { + return + } + } + + // 用户输入解析 + var ok bool + input := make(map[string]string) + if err = json.Unmarshal([]byte(matchJson.Parameters), &input); err != nil { + return + } + if query, ok = input["query"]; !ok { + return fmt.Errorf("query不能为空") + } + + knowledgeConfig := tools.KnowledgeBaseRequest{ + Session: sessionInfo.KnowlegeSessionID, + ApiKey: sysInfo.KnowlegeTenantKey, + Query: query, + } + b, err := json.Marshal(knowledgeConfig) + if err != nil { + return + } + + // 执行工具 + err = r.toolManager.ExecuteTool(channel, c, configData.Tool, b) + if err != nil { + return + } + + return +} + func (r *AiRouterBiz) handleApiTask(channels chan entitys.ResponseData, c *websocket.Conn, matchJson *entitys.Match, task *model.AiTask) (err error) { var ( request l_request.Request diff --git a/internal/entitys/types.go b/internal/entitys/types.go index 189f426..1d10a01 100644 --- a/internal/entitys/types.go +++ b/internal/entitys/types.go @@ -76,8 +76,9 @@ type ConfigDataHttp struct { } type ConfigDataTool struct { - Param map[string]interface{} `json:"param"` - Tool string `json:"tool"` + Param map[string]interface{} `json:"param"` + Request map[string]interface{} `json:"request"` + Tool string `json:"tool"` } // Message 消息 diff --git a/internal/pkg/l_request/README.md b/internal/pkg/l_request/README.md new file mode 100644 index 0000000..f895ffe --- /dev/null +++ b/internal/pkg/l_request/README.md @@ -0,0 +1,25 @@ +## 安装 + +```bash +$ go get -u gitea.cdlsxd.cn/rzy_tools/request/tags +``` + + +## 正常使用 +```go + req := request.Request{ + Method: "POST", + Url: reqUrl, + Json: RequestBody, + Headers: header, + } + resp, _ := req.Send() +``` + +## 同时大量请求或者在协程中使用建议使用 + +```go + r := RequestPools.Get() + defer RequestPools.ClearAndPut(r) + ... +``` \ No newline at end of file diff --git a/internal/pkg/l_request/pool.go b/internal/pkg/l_request/pool.go new file mode 100644 index 0000000..6aed978 --- /dev/null +++ b/internal/pkg/l_request/pool.go @@ -0,0 +1,44 @@ +package l_request + +import ( + "sync" +) + +type RequestPool struct { + pool sync.Pool +} + +var RequestPools = &RequestPool{ + pool: sync.Pool{ + New: func() interface{} { + return new(Request) + }, + }, +} + +func (re *RequestPool) Get() *Request { + return re.pool.Get().(*Request) +} + +func (re *RequestPool) Put(r *Request) { + re.pool.Put(r) +} + +// 重置对象 +func (re *RequestPool) Reset(r *Request) { + r.Method = "" + r.Url = "" + r.Params = nil + r.Headers = nil + r.Cookies = nil + r.Data = nil + r.Json = nil + r.Files = nil + r.Raw = "" + r.JsonByte = nil +} + +func (re *RequestPool) ClearAndPut(r *Request) { + re.Reset(r) + re.Put(r) +} diff --git a/internal/pkg/l_request/pool_test.go b/internal/pkg/l_request/pool_test.go new file mode 100644 index 0000000..596177f --- /dev/null +++ b/internal/pkg/l_request/pool_test.go @@ -0,0 +1,11 @@ +package l_request + +import "testing" + +func TestPool(t *testing.T) { + r := RequestPools.Get() + r.Url = "http://www.baidu.com" + RequestPools.ClearAndPut(r) + a := RequestPools.Get() + t.Log(a.Url) +} diff --git a/internal/pkg/l_request/request.go b/internal/pkg/l_request/request.go new file mode 100644 index 0000000..d47730c --- /dev/null +++ b/internal/pkg/l_request/request.go @@ -0,0 +1,169 @@ +package l_request + +import ( + "crypto/tls" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +// 请求结构体 +type Request struct { + Method string `json:"method"` // 请求方法 + Url string `json:"url"` // 请求url + Params map[string]string `json:"params"` // Query参数 + Headers map[string]string `json:"headers"` // 请求头 + Cookies map[string]string `json:"cookies"` // todo 处理 Cookies + Data map[string]string `json:"data"` // 表单格式请求数据 + Json map[string]interface{} `json:"json"` // JSON格式请求数据 todo 多层 嵌套 + Files map[string]string `json:"files"` // todo 处理 Files + Raw string `json:"raw"` // 原始请求数据 + JsonByte []byte `json:"json_raw"` // JSON格式请求数据 todo 多层 嵌套 + Xml []byte `json:"xml"` // xml +} + +// 响应结构体 +type Response struct { + StatusCode int `json:"status_code"` // 状态码 + Reason string `json:"reason"` // 状态码说明 + Elapsed float64 `json:"elapsed"` // 请求耗时(秒) + Content []byte `json:"content"` // 响应二进制内容 + Text string `json:"text"` // 响应文本 + Headers map[string]string `json:"headers"` // 响应头 + Cookies map[string]string `json:"cookies"` // todo 添加响应Cookies + Request *Request `json:"request"` // 原始请求 +} + +// 处理请求方法 +func (r *Request) getMethod() string { + return strings.ToUpper(r.Method) // 必须转为全部大写 +} + +// 组装URL +func (r *Request) getUrl() string { + if r.Params != nil { + urlValues := url.Values{} + Url, _ := url.Parse(r.Url) // todo 处理err + for key, value := range r.Params { + urlValues.Set(key, value) + } + Url.RawQuery = urlValues.Encode() + return Url.String() + } + return r.Url +} + +// 组装请求数据 +func (r *Request) getData() io.Reader { + var reqBody string + if r.Headers == nil { + r.Headers = make(map[string]string, 1) + } + if r.Raw != "" { + reqBody = r.Raw + } else if r.Data != nil { + urlValues := url.Values{} + for key, value := range r.Data { + urlValues.Add(key, value) + } + reqBody = urlValues.Encode() + if _, ex := r.Headers["Content-Type"]; !ex { + r.Headers["Content-Type"] = "application/x-www-form-urlencoded" + } + + } else if r.Json != nil { + bytesData, _ := json.Marshal(r.Json) + reqBody = string(bytesData) + if _, ex := r.Headers["Content-Type"]; !ex { + r.Headers["Content-Type"] = "application/json" + } + } else if r.JsonByte != nil { + reqBody = string(r.JsonByte) + if _, ex := r.Headers["Content-Type"]; !ex { + r.Headers["Content-Type"] = "application/json" + } + } else if r.Xml != nil { + reqBody = string(r.Xml) + if _, ex := r.Headers["Content-Type"]; !ex { + r.Headers["Content-Type"] = "application/xml" + } + } + return strings.NewReader(reqBody) +} + +// 添加请求头-需要在getData后使用 +func (r *Request) addHeaders(req *http.Request) { + if r.Headers != nil { + for key, value := range r.Headers { + req.Header.Add(key, value) + } + } +} + +// 准备请求 +func (r *Request) prepare() *http.Request { + Method := r.getMethod() + Url := r.getUrl() + Data := r.getData() + req, _ := http.NewRequest(Method, Url, Data) + r.addHeaders(req) + return req +} + +// 组装响应对象 +func (r *Request) packResponse(res *http.Response, elapsed float64) Response { + var resp Response + resBody, _ := io.ReadAll(res.Body) + resp.Content = resBody + resp.Text = string(resBody) + resp.StatusCode = res.StatusCode + resp.Reason = strings.Split(res.Status, " ")[1] + resp.Elapsed = elapsed + resp.Headers = map[string]string{} + for key, value := range res.Header { + resp.Headers[key] = strings.Join(value, ";") + } + return resp +} + +// 发送请求 +func (r *Request) Send() (Response, error) { + req := r.prepare() + client := &http.Client{} + start := time.Now() + res, err := client.Do(req) + if err != nil { + return Response{}, err + } + defer res.Body.Close() + elapsed := time.Since(start).Seconds() + return r.packResponse(res, elapsed), nil +} + +// 跳过证书发送请求 +func (r *Request) SendWithoutSsl() (Response, error) { + tr := &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + } + req := r.prepare() + client := &http.Client{Transport: tr} + start := time.Now() + res, err := client.Do(req) + if err != nil { + return Response{}, err + } + defer res.Body.Close() + elapsed := time.Since(start).Seconds() + return r.packResponse(res, elapsed), nil +} + +// 发送请求 +func (r *Request) SendNoParseResponse() (*http.Response, error) { + req := r.prepare() + client := &http.Client{} + res, err := client.Do(req) + return res, err +} diff --git a/internal/pkg/util/http.go b/internal/pkg/util/http.go new file mode 100644 index 0000000..5bddb26 --- /dev/null +++ b/internal/pkg/util/http.go @@ -0,0 +1,168 @@ +package util + +import ( + "ai_scheduler/internal/pkg/l_request" + "bufio" + "context" + "encoding/json" + "fmt" + + "net/http" + "strings" +) + +type KnowledgeBase struct { + session string + url string + apiKey string +} + +func NewKnowledgeBase(url, apiKey, session string) *KnowledgeBase { + return &KnowledgeBase{ + session: session, + url: url, + apiKey: apiKey, + } +} + +// 请求知识库聊天 +func (this *KnowledgeBase) Chat(ctx context.Context, query string) (text string, err error) { + + req := l_request.Request{ + Method: "post", + Url: this.url + "/api/v1/knowledge-chat/" + this.session, + Params: nil, + Headers: map[string]string{ + "Content-Type": "application/json", + "X-API-Key": this.apiKey, + }, + Cookies: nil, + Data: nil, + Json: map[string]interface{}{ + "query": query, + }, + Files: nil, + Raw: "", + JsonByte: nil, + Xml: nil, + } + + rsp, err := req.SendNoParseResponse() + if err != nil { + return + } + defer rsp.Body.Close() + + err = connectAndReadSSE(rsp) + if err != nil { + return + } + + return +} + +// Message 表示解析后的 SSE 消息 +type Message struct { + Event string // 事件类型（默认 "message"） + Data string // 消息内容（可能多行） + ID string // 消息 ID（可选） +} + +// 连接 SSE 并读取数据 +func connectAndReadSSE(resp *http.Response) error { + + // 验证响应状态和格式 + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("非 200 状态码: %d", resp.StatusCode) + } + contentType := resp.Header.Get("Content-Type") + if !strings.Contains(contentType, "text/event-stream") { + return fmt.Errorf("不支持的 Content-Type: %s", contentType) + } + + // 逐行读取响应流 + scanner := bufio.NewScanner(resp.Body) + var currentMsg Message // 当前正在组装的消息 + + for scanner.Scan() { + line := scanner.Text() + if line == "" { + // 空行表示一条消息结束，处理当前消息 + if currentMsg.Data != "" || currentMsg.Event != "" || currentMsg.ID != "" { + printMessage(currentMsg) + currentMsg = Message{} // 重置消息 + } + continue + } + + // 解析字段（格式："field: value"） + parts := strings.SplitN(line, ":", 2) + if len(parts) < 2 { + continue // 无效行（无冒号），跳过 + } + field := strings.TrimSpace(parts[0]) + value := strings.TrimSpace(parts[1]) + + switch field { + case "event": + currentMsg.Event = value + case "data": + // data 可能多行，用换行符拼接（最后一条消息可能无结尾空行） + currentMsg.Data += value + "" + //case "id": + // currentMsg.ID = value + // 可选：处理 "retry" 字段（服务器建议的重连时间，单位秒） + } + } + + // 检查扫描错误（如连接断开） + if err := scanner.Err(); err != nil { + return fmt.Errorf("读取流失败: %w", err) + } + + // 处理最后一条未结束的消息（无结尾空行） + if currentMsg.Data != "" || currentMsg.Event != "" || currentMsg.ID != "" { + printMessage(currentMsg) + } + + return nil +} + +type MegContent struct { + Id string `json:"id"` // 消息 ID + ResponseType string `json:"response_type"` // 响应类型，answer 或 references + Content string `json:"content"` // 消息内容 + Done bool `json:"done"` // 是否完成 + KnowledgeReferences interface{} `json:"knowledge_references"` +} + +// printMessage 打印解析后的 SSE 消息 +func printMessage(msg Message) { + + //fmt.Printf("--- 收到 SSE 消息 ---") + //fmt.Printf("事件类型: %s，", msg.Event) + //fmt.Printf("消息 ID: %s，", msg.ID) + //fmt.Printf("内容:%s，", strings.TrimSpace(msg.Data)) // 去除末尾多余换行 + + var content MegContent + _ = json.Unmarshal([]byte(msg.Data), &content) + fmt.Println(msg.Data) + + //if content.ResponseType == "answer" { + // //fmt.Printf("%s", content.Content) + // fmt.Println(content) + //} else { + // fmt.Printf("--- 收到 SSE 消息 ---") + // fmt.Printf("事件类型: %s，", msg.Event) + // fmt.Printf("消息 ID: %s，", msg.ID) + // fmt.Printf("内容:%s，", strings.TrimSpace(msg.Data)) // 去除末尾多余换行 + //} + +} + +// getRetryAfter 从响应头获取重连时间（示例，需根据实际响应头调整） +func getRetryAfter(url string) int { + // 实际需重新请求并获取响应头（此处简化为固定值） + // 正确做法：在 connectAndReadSSE 中记录响应头的 Retry-After 字段 + return 5 // 示例：等待 5 秒 +} diff --git a/internal/tools/konwledge_base.go b/internal/tools/konwledge_base.go new file mode 100644 index 0000000..bf82b28 --- /dev/null +++ b/internal/tools/konwledge_base.go @@ -0,0 +1,248 @@ +package tools + +import ( + "ai_scheduler/internal/config" + "ai_scheduler/internal/entitys" + "ai_scheduler/internal/pkg/l_request" + "bufio" + "encoding/json" + "fmt" + "github.com/gofiber/websocket/v2" + "net/http" + "strings" +) + +// 知识库工具 +type KnowledgeBaseTool struct { + config config.ToolConfig +} + +// NewKnowledgeBaseTool 创建知识库工具 +func NewKnowledgeBaseTool(config config.ToolConfig) *KnowledgeBaseTool { + return &KnowledgeBaseTool{config: config} +} + +func (k *KnowledgeBaseTool) GetConfig() config.ToolConfig { + return k.config +} + +// Name 返回工具名称 +func (k *KnowledgeBaseTool) Name() string { + return "knowledgeBase" +} + +// Description 返回工具描述 +func (k *KnowledgeBaseTool) Description() string { + return "请求知识库" +} + +// Definition 返回工具定义 +func (k *KnowledgeBaseTool) Definition() entitys.ToolDefinition { + return entitys.ToolDefinition{ + Type: "function", + Function: entitys.FunctionDef{ + Name: k.Name(), + Description: k.Description(), + Parameters: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "知识库查询条件", + }, + }, + "required": []string{"query"}, + }, + }, + } +} + +// Execute 执行知识库查询 +func (k *KnowledgeBaseTool) Execute(channel chan entitys.ResponseData, c *websocket.Conn, args json.RawMessage) error { + + var params KnowledgeBaseRequest + if err := json.Unmarshal(args, ¶ms); err != nil { + return fmt.Errorf("unmarshal args failed: %w", err) + } + + return k.chat(channel, c, params) + +} + +type KnowledgeBaseRequest struct { + Session string // 知识库会话id + ApiKey string // 知识库apiKey + Query string // 用户输入 +} + +// Message 表示解析后的 SSE 消息 +type Message struct { + Event string // 事件类型（默认 "message"） + Data string // 消息内容（可能多行） + ID string // 消息 ID（可选） +} + +type MegContent struct { + Id string `json:"id"` + ResponseType string `json:"response_type"` + Content string `json:"content"` + Done bool `json:"done"` + KnowledgeReferences interface{} `json:"knowledge_references"` +} + +// 请求知识库聊天 +func (this *KnowledgeBaseTool) chat(channel chan entitys.ResponseData, c *websocket.Conn, param KnowledgeBaseRequest) (err error) { + + req := l_request.Request{ + Method: "post", + Url: this.config.BaseURL + "/api/v1/knowledge-chat/" + param.Session, + Params: nil, + Headers: map[string]string{ + "Content-Type": "application/json", + "X-API-Key": param.ApiKey, + }, + Cookies: nil, + Data: nil, + Json: map[string]interface{}{ + "query": param.Query, + }, + Files: nil, + Raw: "", + JsonByte: nil, + Xml: nil, + } + + rsp, err := req.SendNoParseResponse() + if err != nil { + return + } + defer rsp.Body.Close() + + err = connectAndReadSSE(rsp, channel) + if err != nil { + return + } + + return +} + +// 连接 SSE 并读取数据 +func connectAndReadSSE(resp *http.Response, channel chan entitys.ResponseData) error { + + // 验证响应状态和格式 + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("非 200 状态码: %d", resp.StatusCode) + } + contentType := resp.Header.Get("Content-Type") + if !strings.Contains(contentType, "text/event-stream") { + return fmt.Errorf("不支持的 Content-Type: %s", contentType) + } + + // 逐行读取响应流 + scanner := bufio.NewScanner(resp.Body) + var currentMsg Message // 当前正在组装的消息 + + for scanner.Scan() { + line := scanner.Text() + if line == "" { + // 空行表示一条消息结束，处理当前消息 + if currentMsg.Data != "" || currentMsg.Event != "" || currentMsg.ID != "" { + channel <- entitys.ResponseData{ + Done: false, + Content: currentMsg.Data, + Type: entitys.ResponseJson, + } + currentMsg = Message{} // 重置消息 + } + continue + } + + // 解析字段（格式："field: value"） + parts := strings.SplitN(line, ":", 2) + if len(parts) < 2 { + continue // 无效行（无冒号），跳过 + } + field := strings.TrimSpace(parts[0]) + value := strings.TrimSpace(parts[1]) + + switch field { + case "event": + currentMsg.Event = value + case "data": + // data 可能多行，用换行符拼接（最后一条消息可能无结尾空行） + currentMsg.Data += value + "" + default: + // 忽略未知字段 + } + } + + // 检查扫描错误（如连接断开） + if err := scanner.Err(); err != nil { + return fmt.Errorf("读取流失败: %w", err) + } + + // 处理最后一条未结束的消息（无结尾空行） + if currentMsg.Data != "" || currentMsg.Event != "" || currentMsg.ID != "" { + channel <- entitys.ResponseData{ + Done: false, + Content: currentMsg.Data, + Type: entitys.ResponseJson, + } + } + + return nil +} + +// 获取知识库 session +func GetKnowledgeBaseSession(host, baseId, apiKey string) (string, error) { + req := l_request.Request{ + Method: "post", + Url: host + "/api/v1/sessions", + Params: nil, + Headers: map[string]string{ + "Content-Type": "application/json", + "X-API-Key": apiKey, + }, + Cookies: nil, + Data: nil, + Json: map[string]interface{}{ + "knowledge_base_id": baseId, + }, + Files: nil, + Raw: "", + JsonByte: nil, + Xml: nil, + } + + rsp, err := req.Send() + if err != nil { + return "", err + } + + var result sessionRsp + err = json.Unmarshal(rsp.Content, &result) + + return result.Data.Id, err +} + +type sessionRsp struct { + Data struct { + Id string `json:"id"` + Title string `json:"title"` + Description string `json:"description"` + TenantId int `json:"tenant_id"` + KnowledgeBaseId string `json:"knowledge_base_id"` + MaxRounds int `json:"max_rounds"` + EnableRewrite bool `json:"enable_rewrite"` + FallbackStrategy string `json:"fallback_strategy"` + FallbackResponse string `json:"fallback_response"` + EmbeddingTopK int `json:"embedding_top_k"` + KeywordThreshold float64 `json:"keyword_threshold"` + VectorThreshold float64 `json:"vector_threshold"` + RerankModelId string `json:"rerank_model_id"` + RerankTopK int `json:"rerank_top_k"` + RerankThreshold float64 `json:"rerank_threshold"` + SummaryModelId string `json:"summary_model_id"` + } `json:"data"` + Success bool `json:"success"` +} diff --git a/internal/tools/konwledge_base_test.go b/internal/tools/konwledge_base_test.go new file mode 100644 index 0000000..4ea9bed --- /dev/null +++ b/internal/tools/konwledge_base_test.go @@ -0,0 +1,34 @@ +package tools + +import ( + "ai_scheduler/internal/config" + "ai_scheduler/internal/entitys" + "testing" +) + +func TestKnowledgeBaseTool_Execute(t *testing.T) { + + kb := NewKnowledgeBaseTool(config.ToolConfig{}) + channel := make(chan entitys.ResponseData) + err := kb.Execute(channel, nil, nil) + if err != nil { + t.Errorf("Execute() error = %v", err) + } + +} + +// session +func TestKnowledgeBaseTool_Submit(t *testing.T) { + + apiKey := "sk-EfnUANKMj3DUOiEPJZ5xS8SGMsbO6be_qYAg9uZ8T3zyoFM-" + baseId := "kb-00000001" + host := "http://117.175.169.61:10000" + + sessionId, err := GetKnowledgeBaseSession(host, baseId, apiKey) + if err != nil { + t.Errorf("GetKnowledgeBaseSession() error = %v", err) + } + + t.Log("sessionId:", sessionId) + +} diff --git a/internal/tools/manager.go b/internal/tools/manager.go index 4b854ca..1aa0281 100644 --- a/internal/tools/manager.go +++ b/internal/tools/manager.go @@ -65,6 +65,12 @@ func NewManager(config *config.Config, llm *utils_ollama.Client) *Manager { zltxOrderStatisticsTool := NewZltxOrderStatisticsTool(config.Tools.ZltxOrderStatistics) m.tools[zltxOrderStatisticsTool.Name()] = zltxOrderStatisticsTool } + + // 注册知识库工具 + if config.Tools.Knowledge.Enabled { + knowledgeTool := NewKnowledgeBaseTool(config.Tools.Knowledge) + m.tools[knowledgeTool.Name()] = knowledgeTool + } return m }