Merge pull request #2210 from dgageot/rag-toolset

Refactor RAG from agent-level config to standard toolset type

Author: dgageot

agent-schema.json

@@ -59,9 +59,9 @@
     },
     "rag": {
       "type": "object",
-      "description": "Map of RAG (Retrieval-Augmented Generation) configurations",
+      "description": "Map of reusable RAG source definitions. Define RAG sources here and reference them by name from agent toolsets to avoid duplication.",
       "additionalProperties": {
-        "$ref": "#/definitions/RAGConfig"
+        "$ref": "#/definitions/RAGToolset"
       }
     },
     "metadata": {
@@ -299,13 +299,6 @@
           ],
           "additionalProperties": false
         },
-        "rag": {
-          "type": "array",
-          "description": "List of RAG sources to use for this agent",
-          "items": {
-            "type": "string"
-          }
-        },
         "add_description_parameter": {
           "type": "boolean",
           "description": "Whether to add a 'description' parameter to tool calls, allowing the LLM to provide context about why it is calling a tool"
@@ -807,6 +800,51 @@
       ],
       "additionalProperties": false
     },
+    "RAGToolset": {
+      "type": "object",
+      "description": "Reusable RAG source definition. Define once at the top level and reference by name from agent toolsets. RAG config fields (tool, docs, strategies, results, respect_vcs) are specified directly alongside toolset fields.",
+      "allOf": [
+        {
+          "$ref": "#/definitions/RAGConfig"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "instruction": {
+              "type": "string",
+              "description": "Custom instruction for this RAG source"
+            },
+            "tools": {
+              "type": "array",
+              "description": "Optional list of tools to expose",
+              "items": {
+                "type": "string"
+              }
+            },
+            "name": {
+              "type": "string",
+              "description": "Optional display name override for the RAG tool"
+            },
+            "defer": {
+              "description": "Deferred loading configuration",
+              "oneOf": [
+                {
+                  "type": "boolean",
+                  "description": "Set to true to defer all tools"
+                },
+                {
+                  "type": "array",
+                  "description": "Array of tool names to defer",
+                  "items": {
+                    "type": "string"
+                  }
+                }
+              ]
+            }
+          }
+        }
+      ]
+    },
     "Toolset": {
       "type": "object",
       "description": "Tool configuration",
@@ -830,7 +868,8 @@
             "user_prompt",
             "openapi",
             "model_picker",
-            "background_agents"
+            "background_agents",
+            "rag"
           ]
         },
         "instruction": {
@@ -910,6 +949,10 @@
           "$ref": "#/definitions/ApiConfig",
           "description": "API tool configuration"
         },
+        "rag_config": {
+          "$ref": "#/definitions/RAGConfig",
+          "description": "RAG configuration for type: rag toolsets"
+        },
         "ignore_vcs": {
           "type": "boolean",
           "description": "Whether to ignore VCS files (.git directories and .gitignore patterns) in filesystem operations. Default: true",
@@ -1119,6 +1162,13 @@
               "const": "background_agents"
             }
           }
+        },
+        {
+          "properties": {
+            "type": {
+              "const": "rag"
+            }
+          }
         }
       ]
     },

docs/features/rag/index.md

@@ -36,7 +36,9 @@ agents:
     model: openai/gpt-4o
     instruction: |
       You have access to a knowledge base. Use it to answer questions.
-    rag: [my_docs]
+    toolsets:
+      - type: rag
+        ref: my_docs
 ```
 
 ## Retrieval Strategies

examples/rag.yaml

@@ -1,3 +1,4 @@
+
 agents:
   root:
     model: gpt-5-minimal
@@ -7,8 +8,10 @@ agents:
       can use when it makes sense to do so, based on the user's question.
       If you receive sources from the knowledge base, always include them as
       a markdown list of links to local files at the very end of your response.
-    rag: 
-      - blork_knowledge_base
+    welcome_message: Ask me anything about Blorks.
+    toolsets:
+      - type: rag
+        ref: blork_knowledge_base
 
 models:
   gpt-5-minimal:
@@ -27,4 +30,4 @@ rag:
       - type: chunked-embeddings
         embedding_model: openai/text-embedding-3-small
         database: ./rag/chunked_embeddings.db
-        vector_dimensions: 1536
+        vector_dimensions: 1536

examples/rag/bm25.yaml

@@ -1,12 +1,14 @@
+
 agents:
   root:
     model: openai/gpt-5-mini
     description: a helpful assistant with keyword search
     instruction: |
       You are a helpful assistant that uses BM25 keyword-based search
       to find relevant information in documents.
-    rag: 
-      - blork_knowledge_base
+    toolsets:
+      - type: rag
+        ref: blork_knowledge_base
 
 rag:
   blork_knowledge_base:

examples/rag/hybrid.yaml

@@ -13,8 +13,9 @@ agents:
     instruction: |
       You are a helpful assistant with access to hybrid retrieval
       combining semantic and keyword search for comprehensive results.
-    rag: 
-      - knowledge_base
+    toolsets:
+      - type: rag
+        ref: knowledge_base
 
 rag:
   knowledge_base:

examples/rag/reranking.yaml

@@ -18,8 +18,9 @@ agents:
     instruction: |
       You are a helpful assistant with access to hybrid retrieval
       combining semantic and keyword search for comprehensive results.
-    rag: 
-      - knowledge_base
+    toolsets:
+      - type: rag
+        ref: knowledge_base
 
 rag:
   knowledge_base:

examples/rag/semantic_embeddings.yaml

@@ -17,8 +17,9 @@ agents:
     instruction: |
       You are a helpful coding assistant with access to semantic code search.
       Use the search tool to find relevant code based on meaning, not just keywords.
-    rag:
-      - codebase
+    toolsets:
+      - type: rag
+        ref: codebase
 
 rag:
   codebase:
@@ -78,7 +79,7 @@ rag:
         chunking:
           size: 1000
           respect_word_boundaries: true
-          code_aware: true  # Use tree-sitter for AST-aware chunking
+          code_aware: true  # Use tree-sitter for AST-based chunking
 
     results:
       # Optional: rerank results using an LLM for better relevance
@@ -94,4 +95,3 @@ rag:
       deduplicate: true
       return_full_content: false # return full document content instead of just the matched chunks
       limit: 5
-

pkg/app/app.go

@@ -30,12 +30,6 @@ import (
 	"github.com/docker/docker-agent/pkg/tui/messages"
 )
 
-// RAGInitializer is implemented by runtimes that support background RAG initialization.
-// Local runtimes use this to start indexing early; remote runtimes typically do not.
-type RAGInitializer interface {
-	StartBackgroundRAGInit(ctx context.Context, sendEvent func(runtime.Event))
-}
-
 type App struct {
 	runtime                runtime.Runtime
 	session                *session.Session
@@ -122,18 +116,6 @@ func New(ctx context.Context, rt runtime.Runtime, sess *session.Session, opts ..
 		}
 	}()
 
-	// If the runtime supports background RAG initialization, start it
-	// and forward events to the TUI. Remote runtimes typically handle RAG server-side
-	// and won't implement this optional interface.
-	if ragRuntime, ok := rt.(RAGInitializer); ok {
-		go ragRuntime.StartBackgroundRAGInit(ctx, func(event runtime.Event) {
-			select {
-			case app.events <- event:
-			case <-ctx.Done():
-			}
-		})
-	}
-
 	// Subscribe to tool list changes so the sidebar updates immediately
 	// when an MCP server adds or removes tools (outside of a RunStream).
 	if tcs, ok := rt.(runtime.ToolsChangeSubscriber); ok {

pkg/config/config.go

@@ -124,6 +124,10 @@ func validateConfig(cfg *latest.Config) error {
 		return err
 	}
 
+	if err := resolveRAGDefinitions(cfg); err != nil {
+		return err
+	}
+
 	allNames := map[string]bool{}
 	for _, agent := range cfg.Agents {
 		allNames[agent.Name] = true

pkg/config/latest/parse.go

@@ -26,5 +26,34 @@ func upgradeIfNeeded(c any, _ []byte) (any, error) {
 
 	var config Config
 	types.CloneThroughJSON(old, &config)
+
+	// Migrate AgentConfig.RAG []string → toolsets with type: rag + ref
+	for i, agent := range old.Agents {
+		if len(agent.RAG) == 0 {
+			continue
+		}
+		for _, ragName := range agent.RAG {
+			config.Agents[i].Toolsets = append(config.Agents[i].Toolsets, Toolset{
+				Type: "rag",
+				Ref:  ragName,
+			})
+		}
+	}
+
+	// Migrate top-level RAG map from RAGConfig to RAGToolset
+	if len(old.RAG) > 0 && config.RAG == nil {
+		config.RAG = make(map[string]RAGToolset)
+	}
+	for name, oldRAG := range old.RAG {
+		var ragCfg RAGConfig
+		types.CloneThroughJSON(oldRAG, &ragCfg)
+		config.RAG[name] = RAGToolset{
+			Toolset: Toolset{
+				Type:      "rag",
+				RAGConfig: &ragCfg,
+			},
+		}
+	}
+
 	return config, nil
 }