From fe692b39cd324729e56771b44dbc7114561d46d0 Mon Sep 17 00:00:00 2001 From: Mark Nefedov Date: Wed, 25 Dec 2024 00:48:58 +0300 Subject: [PATCH] Initial commit --- .gitignore | 27 +++++ .idea/.gitignore | 8 ++ .idea/modules.xml | 8 ++ .idea/ollama-to-openrouter-proxy.iml | 9 ++ .idea/vcs.xml | 6 + go.mod | 37 ++++++ go.sum | 91 +++++++++++++++ main.go | 161 +++++++++++++++++++++++++++ provider.go | 132 ++++++++++++++++++++++ 9 files changed, 479 insertions(+) create mode 100644 .gitignore create mode 100644 .idea/.gitignore create mode 100644 .idea/modules.xml create mode 100644 .idea/ollama-to-openrouter-proxy.iml create mode 100644 .idea/vcs.xml create mode 100644 go.mod create mode 100644 go.sum create mode 100644 main.go create mode 100644 provider.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f597e4b --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +### Go template +# If you prefer the allow list template instead of the deny list, see community template: +# https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore +# +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib + +# Test binary, built with `go test -c` +*.test + +# Output of the go coverage tool, specifically when used with LiteIDE +*.out + +# Dependency directories (remove the comment below to include it) +# vendor/ + +# Go workspace file +go.work +go.work.sum + +# env file +.env + diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..13566b8 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,8 @@ +# Default ignored files +/shelf/ +/workspace.xml +# Editor-based HTTP Client requests +/httpRequests/ +# Datasource local storage ignored files +/dataSources/ +/dataSources.local.xml diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..b171fd0 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/ollama-to-openrouter-proxy.iml b/.idea/ollama-to-openrouter-proxy.iml new file mode 100644 index 0000000..5e764c4 --- /dev/null +++ b/.idea/ollama-to-openrouter-proxy.iml @@ -0,0 +1,9 @@ + + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..9f2a42b --- /dev/null +++ b/go.mod @@ -0,0 +1,37 @@ +module ollama-to-openrouter-proxy + +go 1.23 + +require ( + github.com/gin-gonic/gin v1.10.0 + github.com/sashabaranov/go-openai v1.36.0 +) + +require ( + github.com/bytedance/sonic v1.11.6 // indirect + github.com/bytedance/sonic/loader v0.1.1 // indirect + github.com/cloudwego/base64x v0.1.4 // indirect + github.com/cloudwego/iasm v0.2.0 // indirect + github.com/gabriel-vasile/mimetype v1.4.3 // indirect + github.com/gin-contrib/sse v0.1.0 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/go-playground/validator/v10 v10.20.0 // indirect + github.com/goccy/go-json v0.10.2 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.7 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/pelletier/go-toml/v2 v2.2.2 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/ugorji/go/codec v1.2.12 // indirect + golang.org/x/arch v0.8.0 // indirect + golang.org/x/crypto v0.23.0 // indirect + golang.org/x/net v0.25.0 // indirect + golang.org/x/sys v0.20.0 // indirect + golang.org/x/text v0.15.0 // indirect + google.golang.org/protobuf v1.34.1 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..0b5043f --- /dev/null +++ b/go.sum @@ -0,0 +1,91 @@ +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/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/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= +github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= +github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= +github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= +github.com/gin-gonic/gin v1.10.0 h1:nTuyha1TYqgedzytsKYqna+DfLos46nTv2ygFy86HFU= +github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.20.0 h1:K9ISHbSaI0lyB2eWMPJo+kOS/FBExVwjEviJTixqxL8= +github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= +github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= +github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +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/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +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/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/sashabaranov/go-openai v1.36.0 h1:fcSrn8uGuorzPWCBp8L0aCR95Zjb/Dd+ZSML0YZy9EI= +github.com/sashabaranov/go-openai v1.36.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +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 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +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/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= +github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +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.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI= +golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8= +golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= +golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= +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.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y= +golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.15.0 h1:h1V/4gjBv8v9cjcR6+AR5+/cIYK5N/WAgiv4xlsEtAk= +golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= +rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= diff --git a/main.go b/main.go new file mode 100644 index 0000000..d0efce0 --- /dev/null +++ b/main.go @@ -0,0 +1,161 @@ +package main + +import ( + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "os" + "time" + + "github.com/gin-gonic/gin" + openai "github.com/sashabaranov/go-openai" +) + +func main() { + r := gin.Default() + // Load the API key from environment variables or command-line arguments. + apiKey := os.Getenv("OPENAI_API_KEY") + if apiKey == "" { + if len(os.Args) > 1 { + apiKey = os.Args[1] + } else { + slog.Error("OPENAI_API_KEY environment variable or command-line argument not set.") + return + } + } + + provider := NewOpenrouterProvider(apiKey) + + r.GET("/", func(c *gin.Context) { + c.String(http.StatusOK, "Ollama is running") + }) + + r.GET("/api/tags", func(c *gin.Context) { + models, err := provider.GetModels() + if err != nil { + slog.Error("Error getting models", "Error", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + // Respond with the list of models + c.JSON(http.StatusOK, gin.H{"models": models}) + }) + + r.POST("/api/show", func(c *gin.Context) { + var request map[string]string + if err := c.BindJSON(&request); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid JSON payload"}) + return + } + + modelName := request["name"] + if modelName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "Model name is required"}) + return + } + + details, err := provider.GetModelDetails(modelName) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, details) + }) + + r.POST("/api/chat", func(c *gin.Context) { + var request struct { + Model string `json:"model"` + Messages []openai.ChatCompletionMessage `json:"messages"` + } + + // Parse the JSON request + if err := c.ShouldBindJSON(&request); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid JSON payload"}) + return + } + + fullModelName, err := provider.GetFullModelName(request.Model) + if err != nil { + slog.Error("Error getting full model name", "Error", err) + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // Call ChatStream to get the stream + stream, err := provider.ChatStream(request.Messages, fullModelName) + if err != nil { + slog.Error("Failed to create stream", "Error", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + defer stream.Close() // Ensure stream closure + + // Set headers for streaming response + c.Writer.Header().Set("Content-Type", "application/json") + c.Writer.Header().Set("Transfer-Encoding", "chunked") + c.Status(http.StatusOK) + + // Stream responses back to the client + for { + response, err := stream.Recv() + if errors.Is(err, io.EOF) { + // End of stream + break + } + if err != nil { + // Handle errors + slog.Error("Stream error", "Error", err) + break + } + + // Build JSON response structure + responseJSON := map[string]interface{}{ + "model": fullModelName, + "created_at": time.Now().Format(time.RFC3339), + "message": map[string]string{ + "role": "assistant", + "content": response.Choices[0].Delta.Content, + }, + "done": false, + "total_duration": 0, + "load_duration": 0, + "prompt_eval_count": nil, // Replace with actual prompt tokens if available + "eval_count": nil, // Replace with actual completion tokens if available + "eval_duration": 0, + } + + // Marshal and send the JSON response + if err := json.NewEncoder(c.Writer).Encode(responseJSON); err != nil { + slog.Error("Error encoding response", "Error", err) + break + } + + // Flush data to send it immediately + c.Writer.Flush() + } + + // Final response indicating the stream has ended + endResponse := map[string]interface{}{ + "model": fullModelName, + "created_at": time.Now().Format(time.RFC3339), + "message": map[string]string{ + "role": "assistant", + "content": "", + }, + "done": true, + "total_duration": 0, + "load_duration": 0, + "prompt_eval_count": nil, + "eval_count": nil, + "eval_duration": 0, + } + if err := json.NewEncoder(c.Writer).Encode(endResponse); err != nil { + slog.Error("Error encoding end response", "Error", err) + } + }) + + r.Run(":8080") +} diff --git a/provider.go b/provider.go new file mode 100644 index 0000000..12d6d8c --- /dev/null +++ b/provider.go @@ -0,0 +1,132 @@ +package main + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/sashabaranov/go-openai" +) + +type OpenrouterProvider struct { + client *openai.Client + modelNames []string // Shared storage for model names +} + +func NewOpenrouterProvider(apiKey string) *OpenrouterProvider { + config := openai.DefaultConfig(apiKey) + config.BaseURL = "https://openrouter.ai/api/v1/" // Custom endpoint if needed + return &OpenrouterProvider{ + client: openai.NewClientWithConfig(config), + modelNames: []string{}, + } +} + +func (o *OpenrouterProvider) ChatStream(messages []openai.ChatCompletionMessage, modelName string) (*openai.ChatCompletionStream, error) { + // Create a chat completion request + req := openai.ChatCompletionRequest{ + Model: modelName, + Messages: messages, + Stream: true, + } + + // Call the OpenAI API to get a streaming response + stream, err := o.client.CreateChatCompletionStream(context.Background(), req) + if err != nil { + return nil, err + } + + // Return the stream for further processing + return stream, nil +} + +type ModelDetails struct { + ParentModel string `json:"parent_model"` + Format string `json:"format"` + Family string `json:"family"` + Families []string `json:"families"` + ParameterSize string `json:"parameter_size"` + QuantizationLevel string `json:"quantization_level"` +} + +type Model struct { + Name string `json:"name"` + Model string `json:"model,omitempty"` + ModifiedAt string `json:"modified_at,omitempty"` + Size int64 `json:"size,omitempty"` + Digest string `json:"digest,omitempty"` + Details ModelDetails `json:"details,omitempty"` +} + +func (o *OpenrouterProvider) GetModels() ([]Model, error) { + currentTime := time.Now().Format(time.RFC3339) + + // Fetch models from the OpenAI API + modelsResponse, err := o.client.ListModels(context.Background()) + if err != nil { + return nil, err + } + + // Clear shared model storage + o.modelNames = []string{} + + var models []Model + for _, apiModel := range modelsResponse.Models { + // Split model name + parts := strings.Split(apiModel.ID, "/") + name := parts[len(parts)-1] + + // Store name in shared storage + o.modelNames = append(o.modelNames, apiModel.ID) + + // Create model struct + model := Model{ + Name: name, + Model: name, + ModifiedAt: currentTime, + Size: 0, // Stubbed size + Digest: name, + Details: ModelDetails{ + ParentModel: "", + Format: "gguf", + Family: "claude", + Families: []string{"claude"}, + ParameterSize: "175B", + QuantizationLevel: "Q4_K_M", + }, + } + models = append(models, model) + } + + return models, nil +} + +func (o *OpenrouterProvider) GetModelDetails(modelName string) (map[string]interface{}, error) { + // Stub response; replace with actual model details if available + currentTime := time.Now().Format(time.RFC3339) + return map[string]interface{}{ + "license": "STUB License", + "system": "STUB SYSTEM", + "modifiedAt": currentTime, + "details": map[string]interface{}{ + "format": "gguf", + "parameter_size": "200B", + "quantization_level": "Q4_K_M", + }, + "model_info": map[string]interface{}{ + "architecture": "STUB", + "context_length": 200000, + "parameter_count": 200_000_000_000, + }, + }, nil +} + +func (o *OpenrouterProvider) GetFullModelName(alias string) (string, error) { + for _, fullName := range o.modelNames { + if strings.HasSuffix(fullName, alias) { // Match by alias suffix + return fullName, nil + } + } + return "", fmt.Errorf("model alias '%s' not found", alias) +}