ai_scheduler/internal/pkg/utils_ollama/ollama.go

48 lines
827 B
Go

package utils_ollama
import (
"ai_scheduler/internal/config"
"net/http"
"os"
"github.com/gofiber/fiber/v2/log"
"github.com/tmc/langchaingo/llms/ollama"
)
type UtilOllama struct {
Llm *ollama.LLM
}
func NewUtilOllama(c *config.Config, logger log.AllLogger) *UtilOllama {
llm, err := ollama.New(
ollama.WithModel(c.Ollama.Model),
ollama.WithHTTPClient(http.DefaultClient),
ollama.WithServerURL(getUrl(c)),
)
if err != nil {
logger.Fatal(err)
panic(err)
}
return &UtilOllama{
Llm: llm,
}
}
//func (o *UtilOllama) a() {
// var agent agents.Agent
// agent = agents.NewOneShotAgent(llm, tools, opts...)
//
// agents.NewExecutor()
//}
func getUrl(c *config.Config) string {
baseURL := c.Ollama.BaseURL
envURL := os.Getenv("OLLAMA_BASE_URL")
if envURL != "" {
baseURL = envURL
}
return baseURL
}