chore: minor enhancement to be fully importable

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/sse/sse.go

@@ -27,6 +27,8 @@ type (
 	Manager interface {
 		Send(message Envelope)
 		Handle(ctx *fiber.Ctx, cl Listener)
+		Register(cl Listener)
+		Unregister(id string)
 		Clients() []string
 	}
 
@@ -110,6 +112,17 @@ func (manager *broadcastManager) Send(message Envelope) {
 	manager.broadcast <- message
 }
 
+// Register adds a client to the broadcast list and sends message history.
+func (manager *broadcastManager) Register(cl Listener) {
+	manager.register(cl)
+	manager.messageHistory.Send(cl)
+}
+
+// Unregister removes a client from the broadcast list and closes its channel.
+func (manager *broadcastManager) Unregister(id string) {
+	manager.unregister(id)
+}
+
 // Handle sets up a new client and handles the connection.
 func (manager *broadcastManager) Handle(c *fiber.Ctx, cl Listener) {
 

webui/app.go

@@ -36,7 +36,7 @@ type (
 		config           *Config
 		*fiber.App
 		sharedState      *internalTypes.AgentSharedState
-		collectionsState *collectionsState // set when RegisterCollectionRoutes runs; used for in-process RAG
+		collectionsState *CollectionsState // set when RegisterCollectionRoutes runs; used for in-process RAG
 	}
 )
 

webui/collections_backend_inprocess.go

@@ -41,7 +41,7 @@ func newVectorEngine(
 
 // collectionsBackendInProcess implements CollectionsBackend using in-process state.
 type collectionsBackendInProcess struct {
-	state        *collectionsState
+	state        *CollectionsState
 	cfg          *Config
 	openAIClient *openai.Client
 }
@@ -206,10 +206,11 @@ func (b *collectionsBackendInProcess) EntryExists(collection, entry string) bool
 }
 
 // NewInProcessCollectionsBackend creates in-process state (load from disk, start sourceManager) and returns
-// a CollectionsBackend and the state. The caller should set app.collectionsState = state for RAG provider.
-func NewInProcessCollectionsBackend(cfg *Config) (CollectionsBackend, *collectionsState) {
-	state := &collectionsState{
-		collections:   collectionList{},
+// a CollectionsBackend and the CollectionsState. The caller can use CollectionsRAGProviderFromState
+// to create a RAG provider from the state.
+func NewInProcessCollectionsBackend(cfg *Config) (CollectionsBackend, *CollectionsState) {
+	state := &CollectionsState{
+		collections:   CollectionList{},
 		sourceManager: rag.NewSourceManager(&sources.Config{}),
 	}
 

webui/collections_handlers.go

@@ -12,12 +12,15 @@ import (
 	"github.com/mudler/xlog"
 )
 
-type collectionList map[string]*rag.PersistentKB
+// CollectionList maps collection names to their persistent knowledge bases.
+type CollectionList map[string]*rag.PersistentKB
 
-// collectionsState holds in-memory state for the collections API.
-type collectionsState struct {
+// CollectionsState holds in-memory state for the collections API.
+// Exported so that external consumers (e.g. LocalAI) can use NewInProcessCollectionsBackend
+// and CollectionsRAGProviderFromState without going through the webui App.
+type CollectionsState struct {
 	mu               sync.RWMutex
-	collections      collectionList
+	collections      CollectionList
 	sourceManager    *rag.SourceManager
 	ensureCollection func(name string) (*rag.PersistentKB, bool) // get-or-create for internal RAG (agent name as collection)
 }

webui/collections_internal.go

@@ -145,11 +145,12 @@ func (a *internalCompactionAdapter) DeleteEntry(entry string) error {
 	return kb.RemoveEntry(entry)
 }
 
-// CollectionsRAGProvider returns a provider that the pool can use when no LocalRAG URL is set.
-// It returns (RAGDB, KBCompactionClient, true) for a collection name, creating the collection on first use if needed.
-func (app *App) CollectionsRAGProvider() func(collectionName string) (agent.RAGDB, state.KBCompactionClient, bool) {
+// CollectionsRAGProviderFromState returns a RAG provider function from a CollectionsState.
+// This is the standalone version that does not require the webui App. External consumers
+// (e.g. LocalAI) can call NewInProcessCollectionsBackend to get the state, then pass it here.
+func CollectionsRAGProviderFromState(cs *CollectionsState) func(collectionName string) (agent.RAGDB, state.KBCompactionClient, bool) {
 	return func(collectionName string) (agent.RAGDB, state.KBCompactionClient, bool) {
-		if app.collectionsState == nil {
+		if cs == nil {
 			return nil, nil, false
 		}
 		name := strings.TrimSpace(strings.ToLower(collectionName))
@@ -157,10 +158,10 @@ func (app *App) CollectionsRAGProvider() func(collectionName string) (agent.RAGD
 			return nil, nil, false
 		}
 		var kb *rag.PersistentKB
-		app.collectionsState.mu.RLock()
-		kb, ok := app.collectionsState.collections[name]
-		ensure := app.collectionsState.ensureCollection
-		app.collectionsState.mu.RUnlock()
+		cs.mu.RLock()
+		kb, ok := cs.collections[name]
+		ensure := cs.ensureCollection
+		cs.mu.RUnlock()
 		if !ok || kb == nil {
 			if ensure == nil {
 				xlog.Debug("internal RAG: no ensureCollection", "collection", name)
@@ -178,3 +179,9 @@ func (app *App) CollectionsRAGProvider() func(collectionName string) (agent.RAGD
 		return ragAdapter, compAdapter, true
 	}
 }
+
+// CollectionsRAGProvider returns a provider that the pool can use when no LocalRAG URL is set.
+// It delegates to CollectionsRAGProviderFromState using the App's internal collectionsState.
+func (app *App) CollectionsRAGProvider() func(collectionName string) (agent.RAGDB, state.KBCompactionClient, bool) {
+	return CollectionsRAGProviderFromState(app.collectionsState)
+}

webui/routes.go

@@ -220,7 +220,7 @@ func (app *App) registerRoutes(pool *state.AgentPool, webapp *fiber.App) {
 		client := localrag.NewClient(app.config.LocalRAGURL, app.config.LLMAPIKey)
 		collectionsBackend = NewCollectionsBackendHTTP(client)
 	} else {
-		var state *collectionsState
+		var state *CollectionsState
 		collectionsBackend, state = NewInProcessCollectionsBackend(app.config)
 		app.collectionsState = state
 	}