feat: add search tool with RAG (#57)

- add a new MCP Server tool `firebolt_search`. This tool allows to
search Firebolt Documentation using RAG endpoint
- remove full list of documentation articles from `friebolt_docs` tool
- do not require docs_proof in `firebolt_connect` by default. Can be
forced in config

---------

Signed-off-by: Alex Kaplun <o.kaplun@firebolt.io>

Author: oleksandrkaplunvir

README.md

@@ -168,14 +168,25 @@ To enable SSE, set the `--transport` CLI flag (or the `FIREBOLT_MCP_TRANSPORT` e
 
 Optionally, you can specify the address the server should listen on by setting the `--transport-sse-listen-address` CLI flag (or the `FIREBOLT_MCP_TRANSPORT_SSE_LISTEN_ADDRESS` environment variable).
 
+
+#### Requiring LLMs to present docs read proof before connecting
+
+To provide wider context to LLMs before connecting to Firebolt and running queries, by default `firebolt_connect` tool
+requires the LLM to present a read proof of the Firebolt documentation (by querying the `firebolt_docs_overview` tool).
+
+While this provides a good starting point for LLMs ensuring it has full context of Firebolt documentation, at the same time this may lead to slower responses and higher token consumption.
+
+To disable this requirement, set the `--skip-docs-proof` CLI bool flag (or the `FIREBOLT_MCP_SKIP_DOCS_PROOF` environment variable) to `false`.
+
 ## Architecture
 
 Firebolt MCP Server implements the [Model Context Protocol](https://modelcontextprotocol.io/introduction), providing:
 
 1. **Tools** - Task-specific capabilities provided to the LLM:
-    - `firebolt_docs`: Access Firebolt documentation
+    - `firebolt_docs_overview`: Access basic Firebolt documentation overview
     - `firebolt_connect`: Establish connections to Firebolt engines and databases
     - `firebolt_query`: Execute SQL queries against Firebolt
+    - `firebolt_docs_search`: Search Firebolt documentation for any details
 
 2. **Resources** - Data that can be referenced by the LLM:
     - Documentation articles

cmd/firebolt-mcp-server/main.go

@@ -20,7 +20,10 @@ import (
 	"github.com/firebolt-db/mcp-server/pkg/prompts"
 	"github.com/firebolt-db/mcp-server/pkg/resources"
 	"github.com/firebolt-db/mcp-server/pkg/server"
-	"github.com/firebolt-db/mcp-server/pkg/tools"
+	"github.com/firebolt-db/mcp-server/pkg/tools/tool_connect"
+	"github.com/firebolt-db/mcp-server/pkg/tools/tool_docs"
+	"github.com/firebolt-db/mcp-server/pkg/tools/tool_query"
+	"github.com/firebolt-db/mcp-server/pkg/tools/tool_search"
 )
 
 var (
@@ -94,6 +97,14 @@ func main() {
 				Usage:    "Firebolt environment to connect to",
 				Sources:  cli.EnvVars("FIREBOLT_MCP_ENVIRONMENT"),
 			},
+			&cli.BoolFlag{
+				Name:     "skip-docs-proof",
+				Category: "MCP Tools Configuration",
+				Value:    false,
+				Usage: "Skip the requirement for LLM to provide a token as a proof it has reviewed documentation overview. When enabled, LLM " +
+					"will not be forced to gather more starting context and become smarter, but this means more tokens consumed and slower responses.",
+				Sources: cli.EnvVars("FIREBOLT_MCP_SKIP_DOCS_PROOF"),
+			},
 		},
 		Action: run,
 	}
@@ -129,21 +140,40 @@ func run(ctx context.Context, cmd *cli.Command) error {
 	}
 
 	// Initialize MCP server
-	docsProof := generateRandomSecret()
+	docsProofToken := generateRandomSecret()
 	disableResources := cmd.Bool("disable-resources")
-	resourceDocs := resources.NewDocs(fireboltdocs.FS, docsProof)
+	resourceDocs := resources.NewDocs(fireboltdocs.FS, docsProofToken)
 	resourceAccounts := resources.NewAccounts(discoveryClient)
 	resourceDatabases := resources.NewDatabases(dbPool)
 	resourceEngines := resources.NewEngines(dbPool)
+
+	searchCfg := tool_search.Config{
+		BaseURL:      fmt.Sprintf("https://api.%s", cmd.String("environment")),
+		ClientID:     clientID,
+		ClientSecret: clientSecret,
+		TokenURL:     fmt.Sprintf("https://id.%s/oauth/token", cmd.String("environment")),
+	}
+
+	searchTool, err := tool_search.NewSearch(ctx, searchCfg)
+	if err != nil {
+		return fmt.Errorf("failed to create search tool: %w", err)
+	}
+
+	var docsProof *string
+	if !cmd.Bool("skip-docs-proof") {
+		docsProof = &docsProofToken
+	}
+
 	srv := server.NewServer(
 		logger,
 		fullVersion(),
 		cmd.String("transport"),
 		cmd.String("transport-sse-listen-address"),
 		[]server.Tool{
-			tools.NewConnect(resourceAccounts, resourceDatabases, resourceEngines, docsProof, disableResources),
-			tools.NewDocs(resourceDocs, disableResources),
-			tools.NewQuery(dbPool),
+			tool_connect.NewConnect(resourceAccounts, resourceDatabases, resourceEngines, docsProof, disableResources),
+			tool_docs.NewDocs(resourceDocs, disableResources),
+			tool_query.NewQuery(dbPool),
+			searchTool,
 		},
 		[]server.Prompt{
 			prompts.NewFireboltExpert(),

docs/overview.md

@@ -20,9 +20,10 @@ graph TD
         Server["MCP Server Core"]
         
         subgraph "Tools"
-            DocsT["firebolt_docs"]
+            DocsT["firebolt_docs_overview"]
             ConnectT["firebolt_connect"]
             QueryT["firebolt_query"]
+            SearchT["firebolt_docs_search"]
         end
         
         subgraph "Resources"
@@ -48,7 +49,8 @@ graph TD
     
     Server --> DocsT
     Server --> ConnectT
-    Server --> QueryT
+    Server --> QueryT 
+    Server --> SearchT
     Server --> DocsR
     Server --> AccountsR
     Server --> DatabasesR
@@ -81,21 +83,25 @@ Key features:
 
 Tools are executable capabilities exposed to LLMs through the MCP interface:
 
-1. **firebolt_docs** (`pkg/tools/docs.go`)
+1. **firebolt_docs_overview** (`pkg/tools/tool_docs/docs.go`)
    - Provides access to Firebolt documentation
    - Returns embedded markdown content for various documentation articles
    - Helps LLMs understand Firebolt concepts, SQL syntax, and best practices
 
-2. **firebolt_connect** (`pkg/tools/connect.go`)
+2. **firebolt_connect** (`pkg/tools/tool_connect/connect.go`)
    - Lists available Firebolt accounts, databases, and engines
    - Requires a "proof" from documentation to ensure the LLM has read basic Firebolt information
    - Enables discovery of resources before executing queries
 
-3. **firebolt_query** (`pkg/tools/query.go`)
+3. **firebolt_query** (`pkg/tools/tool_query/query.go`)
    - Executes SQL queries against Firebolt databases
    - Manages connections to the specified account, database, and engine
    - Returns query results in JSON format
 
+4. **firebolt_docs_search** (`pkg/tools/tool_docs/search.go`)
+   - Provides RAG search functionality for Firebolt documentation
+   - Allows LLMs to find specific articles in the documentation using semantic search
+
 ### 3. Resources
 
 Resources are information objects that can be accessed by the LLM:
@@ -135,9 +141,10 @@ The MCP server includes specialized clients to interact with Firebolt services:
 sequenceDiagram
     participant LLM as LLM Client
     participant Server as MCP Server
-    participant Docs as firebolt_docs
+    participant Docs as firebolt_docs_overview
     participant Connect as firebolt_connect
     participant Query as firebolt_query
+    participant Search as firebolt_docs_search
     participant FireboltAPI as Firebolt API
     participant FireboltDB as Firebolt Database
 
@@ -146,10 +153,18 @@ sequenceDiagram
     
     Note over LLM, Server: Documentation Flow
     LLM->>Server: Request documentation
-    Server->>Docs: Call firebolt_docs
+    Server->>Docs: Call firebolt_docs_overview
     Docs->>Server: Return documentation resources
     Server->>LLM: Documentation content
     
+    Note over LLM, Server: Search Flow
+    LLM->>Server: Request documentation search
+    Server->>Search: Call firebolt_docs_search
+    Search->>FireboltAPI: Call API to run search
+    FireboltAPI->>Search: Return search results
+    Search->>Server: Return search results
+    Server->>LLM: Search results content
+    
     Note over LLM, Server: Connection Flow
     LLM->>Server: Request resources
     Server->>Connect: Call firebolt_connect

go.mod

@@ -5,6 +5,7 @@ go 1.24.1
 require (
 	github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.0
 	github.com/firebolt-db/firebolt-go-sdk v1.13.0
+	github.com/go-ozzo/ozzo-validation/v4 v4.3.0
 	github.com/gocolly/colly/v2 v2.3.0
 	github.com/modelcontextprotocol/go-sdk v1.2.0
 	github.com/neilotoole/slogt v1.1.0
@@ -21,6 +22,7 @@ require (
 	github.com/antchfx/htmlquery v1.3.5 // indirect
 	github.com/antchfx/xmlquery v1.5.0 // indirect
 	github.com/antchfx/xpath v1.3.5 // indirect
+	github.com/asaskevich/govalidator v0.0.0-20200108200545-475eaeb16496 // indirect
 	github.com/astaxie/beego v1.12.3 // indirect
 	github.com/bits-and-blooms/bitset v1.24.4 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect