diff --git a/agent/llm/capabilities.go b/agent/llm/capabilities.go index 62f2b5f1..92fe8834 100644 --- a/agent/llm/capabilities.go +++ b/agent/llm/capabilities.go @@ -83,6 +83,12 @@ func capabilitiesFromMap(m map[string]interface{}) *goullm.Capabilities { if v, ok := m["temperature_adjustable"].(bool); ok { caps.TemperatureAdjustable = v } + if v, ok := m["embedding"].(bool); ok { + caps.Embedding = v + } + if v, ok := m["image_generation"].(bool); ok { + caps.ImageGeneration = v + } return caps } @@ -110,26 +116,8 @@ func GetCapabilitiesMap(connectorID string) map[string]interface{} { return ToMap(caps) } -// ToMap converts Capabilities to map[string]interface{} +// ToMap converts Capabilities to map[string]interface{}. +// Delegates to the canonical Capabilities.ToMap() method in gou/llm. func ToMap(caps *goullm.Capabilities) map[string]interface{} { - if caps == nil { - return nil - } - - result := make(map[string]interface{}) - - if caps.Vision != nil { - result["vision"] = caps.Vision - } - - result["audio"] = caps.Audio - result["stt"] = caps.STT - result["tool_calls"] = caps.ToolCalls - result["reasoning"] = caps.Reasoning - result["streaming"] = caps.Streaming - result["json"] = caps.JSON - result["multimodal"] = caps.Multimodal - result["temperature_adjustable"] = caps.TemperatureAdjustable - - return result + return caps.ToMap() } diff --git a/llmprovider/models.go b/llmprovider/models.go index 4052f262..90f9bed2 100644 --- a/llmprovider/models.go +++ b/llmprovider/models.go @@ -338,23 +338,9 @@ func capabilitiesFromConn(conn connector.Connector) *goullm.Capabilities { } // capsToMap converts Capabilities to map[string]interface{} for process handlers. +// Delegates to the canonical Capabilities.ToMap() method in gou/llm. func capsToMap(caps *goullm.Capabilities) map[string]interface{} { - if caps == nil { - return nil - } - result := make(map[string]interface{}) - if caps.Vision != nil { - result["vision"] = caps.Vision - } - result["audio"] = caps.Audio - result["stt"] = caps.STT - result["tool_calls"] = caps.ToolCalls - result["reasoning"] = caps.Reasoning - result["streaming"] = caps.Streaming - result["json"] = caps.JSON - result["multimodal"] = caps.Multimodal - result["temperature_adjustable"] = caps.TemperatureAdjustable - return result + return caps.ToMap() } func defaultCaps() *goullm.Capabilities { diff --git a/llmprovider/sync.go b/llmprovider/sync.go index 087e2f43..74c54718 100644 --- a/llmprovider/sync.go +++ b/llmprovider/sync.go @@ -361,6 +361,12 @@ func capabilitiesFromCapabilities(c *goullm.Capabilities) []string { if c.Multimodal { out = append(out, "multimodal") } + if c.Embedding { + out = append(out, "embedding") + } + if c.ImageGeneration { + out = append(out, "image_generation") + } return out } diff --git a/openapi/llm/llm.go b/openapi/llm/llm.go index b17e2a52..2f6de5d8 100644 --- a/openapi/llm/llm.go +++ b/openapi/llm/llm.go @@ -72,6 +72,11 @@ func listProviders(c *gin.Context) { } capabilities := getCapabilitiesFromConn(conn) + + if isNonChatModel(capabilities) && !hasFilter(filters, "embedding") && !hasFilter(filters, "image_generation") { + continue + } + if len(filters) > 0 && !matchesFilters(capabilities, filters) { continue } @@ -109,6 +114,27 @@ func getCapabilitiesFromConn(conn connector.Connector) map[string]interface{} { return agentllm.ToMap(caps) } +// isNonChatModel returns true if capabilities indicate a non-chat model (embedding or image generation). +func isNonChatModel(caps map[string]interface{}) bool { + if v, ok := caps["embedding"].(bool); ok && v { + return true + } + if v, ok := caps["image_generation"].(bool); ok && v { + return true + } + return false +} + +// hasFilter checks whether a specific filter string is present in the filters list. +func hasFilter(filters []string, name string) bool { + for _, f := range filters { + if f == name { + return true + } + } + return false +} + // matchesFilters checks if capabilities match all requested filters // Filters are matched case-insensitively and support the following capability keys: // - vision: true or string value like "openai", "claude" @@ -119,6 +145,8 @@ func getCapabilitiesFromConn(conn connector.Connector) map[string]interface{} { // - streaming: bool // - json: bool // - multimodal: bool +// - embedding: bool +// - image_generation: bool // - temperature_adjustable: bool func matchesFilters(capabilities map[string]interface{}, filters []string) bool { if capabilities == nil { diff --git a/openapi/setting/cloud.go b/openapi/setting/cloud.go index 4194a81e..4f9163e9 100644 --- a/openapi/setting/cloud.go +++ b/openapi/setting/cloud.go @@ -204,6 +204,7 @@ func handleCloudUpdate(c *gin.Context) { respondError(c, http.StatusInternalServerError, err.Error()) return } + invalidateCloudModelCache() def := cloudDefaultRegion() result := CloudPageData{ @@ -317,6 +318,43 @@ func handleCloudTest(c *gin.Context) { }) } +// handleCloudRefresh invalidates the cloud model cache and re-fetches the model list. +// POST /setting/cloud/refresh +func handleCloudRefresh(c *gin.Context) { + if !guardOwner(c) { + return + } + info := authorized.GetInfo(c) + scope := cloudScope(info) + + saved, _ := setting.Global.Get(scope, cloudNS) + if saved == nil { + respondError(c, http.StatusBadRequest, "cloud service not configured") + return + } + + status, _ := saved["status"].(string) + if status != "connected" { + respondError(c, http.StatusBadRequest, "cloud service not connected") + return + } + + encKey, _ := saved["api_key"].(string) + if encKey == "" { + respondError(c, http.StatusBadRequest, "no API key configured") + return + } + + apiURL := resolveCloudAPIURL(saved) + invalidateCloudModelCache() + models := fetchCloudModels(apiURL, cloudDecrypt(encKey)) + + response.RespondWithSuccess(c, http.StatusOK, map[string]interface{}{ + "success": true, + "count": len(models), + }) +} + // --------------------------------------------------------------------------- // Crypto helpers (AES-256-GCM, same scheme as llmprovider) // --------------------------------------------------------------------------- diff --git a/openapi/setting/cloud_model_cache_test.go b/openapi/setting/cloud_model_cache_test.go new file mode 100644 index 00000000..550eeaa8 --- /dev/null +++ b/openapi/setting/cloud_model_cache_test.go @@ -0,0 +1,92 @@ +package setting + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" +) + +func TestFetchCloudModels_CachesAfterFirstCall(t *testing.T) { + var hits int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&hits, 1) + json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"id": "gpt-4o", "object": "model"}, + }, + }) + })) + defer srv.Close() + + invalidateCloudModelCache() + + models := fetchCloudModels(srv.URL, "test-key") + if len(models) == 0 { + t.Fatal("expected models from first fetch, got none") + } + if atomic.LoadInt64(&hits) != 1 { + t.Fatalf("expected 1 HTTP hit after first fetch, got %d", atomic.LoadInt64(&hits)) + } + + models2 := fetchCloudModels(srv.URL, "test-key") + if len(models2) == 0 { + t.Fatal("expected models from cached fetch, got none") + } + if atomic.LoadInt64(&hits) != 1 { + t.Fatalf("expected still 1 HTTP hit after second fetch (cache), got %d", atomic.LoadInt64(&hits)) + } +} + +func TestFetchCloudModels_InvalidateForcesRefetch(t *testing.T) { + var hits int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&hits, 1) + json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"id": "gpt-4o", "object": "model"}, + }, + }) + })) + defer srv.Close() + + invalidateCloudModelCache() + + fetchCloudModels(srv.URL, "test-key") + if atomic.LoadInt64(&hits) != 1 { + t.Fatalf("expected 1 HTTP hit, got %d", atomic.LoadInt64(&hits)) + } + + invalidateCloudModelCache() + + fetchCloudModels(srv.URL, "test-key") + if atomic.LoadInt64(&hits) != 2 { + t.Fatalf("expected 2 HTTP hits after invalidation, got %d", atomic.LoadInt64(&hits)) + } +} + +func TestFetchCloudModels_URLChangeForcesRefetch(t *testing.T) { + var hits int64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt64(&hits, 1) + json.NewEncoder(w).Encode(map[string]interface{}{ + "data": []map[string]interface{}{ + {"id": "gpt-4o", "object": "model"}, + }, + }) + })) + defer srv.Close() + + invalidateCloudModelCache() + + fetchCloudModels(srv.URL, "test-key") + if atomic.LoadInt64(&hits) != 1 { + t.Fatalf("expected 1 HTTP hit, got %d", atomic.LoadInt64(&hits)) + } + + fetchCloudModels(srv.URL+"/other", "test-key") + if atomic.LoadInt64(&hits) != 2 { + t.Fatalf("expected 2 HTTP hits after URL change, got %d", atomic.LoadInt64(&hits)) + } +} diff --git a/openapi/setting/llm.go b/openapi/setting/llm.go index fa4af29e..e2a3086b 100644 --- a/openapi/setting/llm.go +++ b/openapi/setting/llm.go @@ -123,10 +123,8 @@ func llmValidateKey(providerType, apiURL, apiKey string) error { var ( cloudModelCache []map[string]interface{} - cloudModelCacheAt time.Time cloudModelCacheURL string cloudModelCacheMu sync.Mutex - cloudModelCacheTTL = 5 * time.Minute ) func buildCloudPreset(info *oauthTypes.AuthorizedInfo) { @@ -177,7 +175,7 @@ func resolveCloudAPIURL(saved map[string]interface{}) string { func fetchCloudModels(apiURL, apiKey string) []map[string]interface{} { cloudModelCacheMu.Lock() - if cloudModelCache != nil && cloudModelCacheURL == apiURL && time.Since(cloudModelCacheAt) < cloudModelCacheTTL { + if cloudModelCache != nil && cloudModelCacheURL == apiURL { cached := cloudModelCache cloudModelCacheMu.Unlock() return cached @@ -230,13 +228,19 @@ func fetchCloudModels(apiURL, apiKey string) []map[string]interface{} { cloudModelCacheMu.Lock() cloudModelCache = models - cloudModelCacheAt = time.Now() cloudModelCacheURL = apiURL cloudModelCacheMu.Unlock() return models } +func invalidateCloudModelCache() { + cloudModelCacheMu.Lock() + cloudModelCache = nil + cloudModelCacheURL = "" + cloudModelCacheMu.Unlock() +} + func mapCloudModel(item map[string]interface{}) map[string]interface{} { id, _ := item["id"].(string) if id == "" { diff --git a/openapi/setting/setting.go b/openapi/setting/setting.go index a8001501..5a754418 100644 --- a/openapi/setting/setting.go +++ b/openapi/setting/setting.go @@ -36,6 +36,7 @@ func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) { cloud.GET("", handleCloudGet) cloud.PUT("", handleCloudUpdate) cloud.POST("/test", handleCloudTest) + cloud.POST("/refresh", handleCloudRefresh) llm := group.Group("/llm") llm.GET("", handleLLMGet)