From 3ecbbf3eaceff4b3f6bd9791494f1f09bd892bbb Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 6 Mar 2026 16:39:30 +0100 Subject: [PATCH] chore: minor enhancement to be fully importable Signed-off-by: Ettore Di Giacinto --- core/sse/sse.go | 13 +++++++++++++ webui/app.go | 2 +- webui/collections_backend_inprocess.go | 11 ++++++----- webui/collections_handlers.go | 11 +++++++---- webui/collections_internal.go | 23 +++++++++++++++-------- webui/routes.go | 2 +- 6 files changed, 43 insertions(+), 19 deletions(-) diff --git a/core/sse/sse.go b/core/sse/sse.go index b4b5093..e868c86 100644 --- a/core/sse/sse.go +++ b/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) { diff --git a/webui/app.go b/webui/app.go index f96249a..191031a 100644 --- a/webui/app.go +++ b/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 } ) diff --git a/webui/collections_backend_inprocess.go b/webui/collections_backend_inprocess.go index ddd4c33..edab690 100644 --- a/webui/collections_backend_inprocess.go +++ b/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{}), } diff --git a/webui/collections_handlers.go b/webui/collections_handlers.go index 4b2b2c1..1146805 100644 --- a/webui/collections_handlers.go +++ b/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) } diff --git a/webui/collections_internal.go b/webui/collections_internal.go index 52bcd33..af5936b 100644 --- a/webui/collections_internal.go +++ b/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) +} diff --git a/webui/routes.go b/webui/routes.go index f066951..c000eb2 100644 --- a/webui/routes.go +++ b/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 }