Merge pull request #211 from APIParkLab/feature/1.5-local-model

This pull request includes several changes aimed at improving error handling, adding new functionalities, and refactoring existing code. The changes primarily focus on the AI provider and controller modules.

Error Handling Improvements:
Added checks to ensure the client is initialized before performing operations in multiple functions (PullModel, StopPull, CancelPipeline, RemoveModel, ModelsInstalled) in ai-provider/local/executor.go. [1] [2] [3] [4]
New Functionalities:
Introduced OllamaConfig and OllamaConfigUpdate methods to the ILocalModelController interface and implemented them in controller/ai-local/iml.go. These methods allow for getting and updating the Ollama configuration. [1] [2]
Added functionality to automatically subscribe all applications to new services in the Create method of controller/service/iml.go.
Refactoring:
Refactored the initialization of the Ollama client by replacing the static address with a ResetOllamaAddress function in ai-provider/local/local.go.
Removed unused code and imports, such as the newAIUpstream function and upstream_dto import in controller/service/iml.go. [1] [2]
Codebase Simplification:
Simplified the OnInit method in controller/system/iml.go by consolidating the creation of default entities and adding subscription logic. [1] [2] [3]
Additional Changes:
Added new imports and modules to support the new functionalities and refactoring efforts. [1] [2] [3]

Author: Dot-Liu

ai-provider/local/executor.go

@@ -210,6 +210,9 @@ func (e *AsyncExecutor) DistributeToModelPipelines(model string, msg PullMessage
 type PullCallback func(msg PullMessage) error
 
 func PullModel(model string, id string, fn PullCallback) (*Pipeline, error) {
+	if client == nil {
+		return nil, fmt.Errorf("client not initialized")
+	}
 	mp, has := taskExecutor.GetModelPipeline(model)
 	if !has {
 		mp = newModelPipeline(taskExecutor.ctx, 100)
@@ -279,6 +282,9 @@ func PullModel(model string, id string, fn PullCallback) (*Pipeline, error) {
 }
 
 func StopPull(model string) {
+	if client == nil {
+		return
+	}
 	taskExecutor.CloseModelPipeline(model)
 }
 
@@ -287,6 +293,9 @@ func CancelPipeline(model string, id string) {
 }
 
 func RemoveModel(model string) error {
+	if client == nil {
+		return fmt.Errorf("client not initialized")
+	}
 	taskExecutor.CloseModelPipeline(model)
 	err := client.Delete(context.Background(), &api.DeleteRequest{Model: model})
 	if err != nil {
@@ -298,6 +307,9 @@ func RemoveModel(model string) error {
 }
 
 func ModelsInstalled() ([]Model, error) {
+	if client == nil {
+		return nil, fmt.Errorf("client not initialized")
+	}
 	result, err := client.List(context.Background())
 	if err != nil {
 		return nil, err

ai-provider/local/local.go

@@ -4,27 +4,18 @@ import (
 	"net/http"
 	"net/url"
 
-	"github.com/eolinker/eosc/env"
 	"github.com/ollama/ollama/api"
 )
 
 var (
-	ollamaAddress    = "http://127.0.0.1:11434"
-	EnvOllamaAddress = "OLLAMA_ADDRESS"
-	client           *api.Client
+	client *api.Client
 )
 
-func init() {
-	address, has := env.GetEnv(EnvOllamaAddress)
-	if !has {
-		address = ollamaAddress
-	}
+func ResetOllamaAddress(address string) error {
 	u, err := url.Parse(address)
 	if err != nil {
-		u, err = url.Parse(ollamaAddress)
-		if err != nil {
-			panic(err)
-		}
+		return err
 	}
 	client = api.NewClient(u, http.DefaultClient)
+	return nil
 }

ai-provider/local/ollama.go

@@ -1,7 +1,6 @@
 package ai_provider_local
 
 var (
-	OllamaBase   = "http://apipark-ollama:11434"
 	OllamaConfig = "{\n  \"mirostat\": 0,\n  \"mirostat_eta\": 0.1,\n  \"mirostat_tau\": 5.0,\n  \"num_ctx\": 4096,\n  \"repeat_last_n\":64,\n  \"repeat_penalty\": 1.1,\n  \"temperature\": 0.7,\n  \"seed\": 42,\n  \"num_predict\": 42,\n  \"top_k\": 40,\n  \"top_p\": 0.9,\n  \"min_p\": 0.5\n}\n"
 	OllamaSvg    = `<?xml version="1.0" standalone="no"?>
 <!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 20010904//EN"

controller/ai-local/controller.go

@@ -18,6 +18,8 @@ type ILocalModelController interface {
 	Update(ctx *gin.Context, model string, input *ai_local_dto.Update) error
 	State(ctx *gin.Context, model string) (*ai_local_dto.DeployState, *ai_local_dto.ModelInfo, error)
 	SimpleList(ctx *gin.Context) ([]*ai_local_dto.SimpleItem, error)
+	OllamaConfig(ctx *gin.Context) (*ai_local_dto.OllamaConfig, error)
+	OllamaConfigUpdate(ctx *gin.Context, input *ai_local_dto.OllamaConfig) error
 }
 
 func init() {

controller/ai-local/iml.go

@@ -7,7 +7,15 @@ import (
 	"io"
 	"math"
 	"net/http"
+	"net/url"
 	"strings"
+	"time"
+
+	ai_balance "github.com/APIParkLab/APIPark/module/ai-balance"
+
+	"github.com/APIParkLab/APIPark/module/system"
+	system_dto "github.com/APIParkLab/APIPark/module/system/dto"
+	ollama_api "github.com/ollama/ollama/api"
 
 	"github.com/APIParkLab/APIPark/module/subscribe"
 	subscribe_dto "github.com/APIParkLab/APIPark/module/subscribe/dto"
@@ -47,13 +55,54 @@ type imlLocalModelController struct {
 	serviceModule   service.IServiceModule     `autowired:""`
 	catalogueModule catalogue.ICatalogueModule `autowired:""`
 	aiAPIModule     ai_api.IAPIModule          `autowired:""`
+	aiBalanceModule ai_balance.IBalanceModule  `autowired:""`
 	appModule       service.IAppModule         `autowired:""`
 	routerModule    router.IRouterModule       `autowired:""`
 	subscribeModule subscribe.ISubscribeModule `autowired:""`
 	docModule       service.IServiceDocModule  `autowired:""`
+	settingModule   system.ISettingModule      `autowired:""`
 	transaction     store.ITransaction         `autowired:""`
 }
 
+func (i *imlLocalModelController) OllamaConfig(ctx *gin.Context) (*ai_local_dto.OllamaConfig, error) {
+	cfg := i.settingModule.Get(ctx)
+	return &ai_local_dto.OllamaConfig{
+		Address: cfg.OllamaAddress,
+	}, nil
+}
+
+var (
+	client = &http.Client{
+		Timeout: 2 * time.Second,
+	}
+)
+
+func (i *imlLocalModelController) OllamaConfigUpdate(ctx *gin.Context, input *ai_local_dto.OllamaConfig) error {
+	u, err := url.Parse(input.Address)
+	if err != nil {
+		return nil
+	}
+	ollamaClient := ollama_api.NewClient(u, client)
+	_, err = ollamaClient.Version(ctx)
+	if err != nil {
+		return err
+	}
+	return i.transaction.Transaction(ctx, func(ctx context.Context) error {
+		err = i.module.SyncLocalModels(ctx, input.Address)
+		if err != nil {
+			return err
+		}
+		err = i.aiBalanceModule.SyncLocalBalances(ctx, input.Address)
+		if err != nil {
+			return err
+		}
+
+		return i.settingModule.Set(ctx, &system_dto.InputSetting{
+			OllamaAddress: &input.Address,
+		})
+	})
+}
+
 func (i *imlLocalModelController) SimpleList(ctx *gin.Context) ([]*ai_local_dto.SimpleItem, error) {
 	return i.module.SimpleList(ctx)
 }
@@ -209,6 +258,7 @@ func (i *imlLocalModelController) initAILocalService(ctx context.Context, model
 			ApprovalType: "auto",
 			Kind:         "ai",
 			Provider:     &providerId,
+			Model:        &model,
 		})
 		if err != nil {
 			return err

controller/service/iml.go

@@ -32,8 +32,6 @@ import (
 
 	api_doc "github.com/APIParkLab/APIPark/module/api-doc"
 
-	upstream_dto "github.com/APIParkLab/APIPark/module/upstream/dto"
-
 	"github.com/eolinker/eosc/log"
 
 	application_authorization "github.com/APIParkLab/APIPark/module/application-authorization"
@@ -94,7 +92,10 @@ func (i *imlServiceController) QuickCreateAIService(ctx *gin.Context, input *ser
 		if err != nil {
 			return err
 		}
-
+		p, err := i.providerModule.Provider(ctx, input.Provider)
+		if err != nil {
+			return err
+		}
 		id := uuid.NewString()
 		prefix := fmt.Sprintf("/%s", id[:8])
 		catalogueInfo, err := i.catalogueModule.DefaultCatalogue(ctx)
@@ -111,6 +112,7 @@ func (i *imlServiceController) QuickCreateAIService(ctx *gin.Context, input *ser
 			Catalogue:    catalogueInfo.Id,
 			ApprovalType: "auto",
 			Provider:     &input.Provider,
+			Model:        &p.DefaultLLM,
 			Kind:         "ai",
 		})
 		return err
@@ -294,25 +296,18 @@ func (i *imlServiceController) editAIService(ctx *gin.Context, id string, input
 	if input.Provider == nil {
 		return nil, fmt.Errorf("provider is required")
 	}
-	p, has := model_runtime.GetProvider(*input.Provider)
-	if !has {
-		return nil, fmt.Errorf("provider not found")
+	if *input.Provider != "ollama" {
+		_, has := model_runtime.GetProvider(*input.Provider)
+		if !has {
+			return nil, fmt.Errorf("provider not found")
+		}
 	}
-	info, err := i.module.Get(ctx, id)
-	if err != nil {
 
-	}
-	err = i.transaction.Transaction(ctx, func(txCtx context.Context) error {
-		info, err = i.module.Edit(ctx, id, input)
-		if err != nil {
-			return err
-		}
-		_, err = i.upstreamModule.Save(ctx, id, newAIUpstream(id, *input.Provider, p.URI()))
-		return err
-	})
+	info, err := i.module.Edit(ctx, id, input)
 	if err != nil {
 		return nil, err
 	}
+	//_, err = i.upstreamModule.Save(ctx, id, newAIUpstream(id, *input.Provider, p.URI()))
 
 	return info, nil
 }
@@ -482,13 +477,13 @@ func (i *imlServiceController) Create(ctx *gin.Context, teamID string, input *se
 	}
 	var err error
 	var info *service_dto.Service
-	err = i.transaction.Transaction(ctx, func(txCtx context.Context) error {
-		info, err = i.module.Create(txCtx, teamID, input)
+	err = i.transaction.Transaction(ctx, func(ctx context.Context) error {
+		info, err = i.module.Create(ctx, teamID, input)
 		if err != nil {
 			return err
 		}
 		path := fmt.Sprintf("/%s/", strings.Trim(input.Prefix, "/"))
-		_, err = i.routerModule.Create(txCtx, info.Id, &router_dto.Create{
+		_, err = i.routerModule.Create(ctx, info.Id, &router_dto.Create{
 			Id:          uuid.New().String(),
 			Name:        "",
 			Path:        path + "*",
@@ -504,6 +499,15 @@ func (i *imlServiceController) Create(ctx *gin.Context, teamID string, input *se
 			},
 			Disable: false,
 		})
+		apps, err := i.appModule.Search(ctx, teamID, "")
+		if err != nil {
+			return err
+		}
+		for _, app := range apps {
+			i.subscribeModule.AddSubscriber(ctx, info.Id, &subscribe_dto.AddSubscriber{
+				Application: app.Id,
+			})
+		}
 		return err
 	})
 	return info, err
@@ -590,22 +594,22 @@ func (i *imlAppController) DeleteApp(ctx *gin.Context, appId string) error {
 	return i.module.DeleteApp(ctx, appId)
 }
 
-func newAIUpstream(id string, provider string, uri model_runtime.IProviderURI) *upstream_dto.Upstream {
-	return &upstream_dto.Upstream{
-		Type:            "http",
-		Balance:         "round-robin",
-		Timeout:         300000,
-		Retry:           0,
-		Remark:          fmt.Sprintf("auto create by ai service %s,provider is %s", id, provider),
-		LimitPeerSecond: 0,
-		ProxyHeaders:    nil,
-		Scheme:          uri.Scheme(),
-		PassHost:        "node",
-		Nodes: []*upstream_dto.NodeConfig{
-			{
-				Address: uri.Host(),
-				Weight:  100,
-			},
-		},
-	}
-}
+//func newAIUpstream(id string, provider string, uri model_runtime.IProviderURI) *upstream_dto.Upstream {
+//	return &upstream_dto.Upstream{
+//		Type:            "http",
+//		Balance:         "round-robin",
+//		Timeout:         300000,
+//		Retry:           0,
+//		Remark:          fmt.Sprintf("auto create by ai service %s,provider is %s", id, provider),
+//		LimitPeerSecond: 0,
+//		ProxyHeaders:    nil,
+//		Scheme:          uri.Scheme(),
+//		PassHost:        "node",
+//		Nodes: []*upstream_dto.NodeConfig{
+//			{
+//				Address: uri.Host(),
+//				Weight:  100,
+//			},
+//		},
+//	}
+//}