feat: implementation of semanticly retrieved dynamic toolsets (#818)

- Implements experimental dynamic tool selection through `Gram-Mode`
header `semantic_search` | `progressive_search`
    - question should we call this mode `semantic` instead?
- Vector storage and semantic search of toolsets is done via pgvector,
all encapsulated into the `rag/search_tools.go` client.
- For embedding this will be done through openrouter base client right
now using the `text-embedding-3-small`
- Given this is experimental and a toolset cannot be marked as dynamic.
Indexing is done on demand as necessary in find_tools. We do not need to
reindex if we recognize there has been no change to the toolset version
or tools. Even in cases where toolset indexing is necessary it is still
quite fast, easily under a second (spans tracking)
- Right now we just hard delete old embeddings. There is no immediate
reason to keep them around so it preserves space. The data model is
setup for soft deletions

Following this we'll complete a documentation page on these experimental
modes and start our own testing across large toolsets between different
modes.

Author: ryan-timothy-albert

.changeset/light-meals-beam.md

@@ -0,0 +1,5 @@
+---
+"server": patch
+---
+
+adds experimental gram-mode:embedding for dynamic MCP tool selection based on semantic search

server/cmd/gram/start.go

@@ -18,6 +18,7 @@ import (
 	"github.com/sourcegraph/conc/pool"
 	"github.com/speakeasy-api/gram/server/internal/logs"
 	"github.com/speakeasy-api/gram/server/internal/productfeatures"
+	"github.com/speakeasy-api/gram/server/internal/rag"
 	"github.com/urfave/cli/v2"
 	"github.com/urfave/cli/v2/altsrc"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
@@ -516,6 +517,7 @@ func newStartCommand() *cli.Command {
 			shutdownFuncs = append(shutdownFuncs, shutdown)
 
 			chatClient := chat.NewChatClient(logger, tracerProvider, meterProvider, db, openRouter, baseChatClient, env, encryptionClient, cache.NewRedisCacheAdapter(redisClient), guardianPolicy, functionsOrchestrator)
+			ragService := rag.NewToolsetVectorStore(tracerProvider, db, baseChatClient)
 			mux := goahttp.NewMuxer()
 
 			mux.Use(middleware.CORSMiddleware(c.String("environment"), c.String("server-url")))
@@ -550,7 +552,7 @@ func newStartCommand() *cli.Command {
 			instances.Attach(mux, instances.NewService(logger, tracerProvider, meterProvider, db, sessionManager, env, encryptionClient, cache.NewRedisCacheAdapter(redisClient), guardianPolicy, functionsOrchestrator, billingTracker, tcm))
 			mcpMetadataService := mcpmetadata.NewService(logger, db, sessionManager, serverURL, siteURL, cache.NewRedisCacheAdapter(redisClient))
 			mcpmetadata.Attach(mux, mcpMetadataService)
-			mcp.Attach(mux, mcp.NewService(logger, tracerProvider, meterProvider, db, sessionManager, env, posthogClient, serverURL, encryptionClient, cache.NewRedisCacheAdapter(redisClient), guardianPolicy, functionsOrchestrator, oauthService, billingTracker, billingRepo, tcm), mcpMetadataService)
+			mcp.Attach(mux, mcp.NewService(logger, tracerProvider, meterProvider, db, sessionManager, env, posthogClient, serverURL, encryptionClient, cache.NewRedisCacheAdapter(redisClient), guardianPolicy, functionsOrchestrator, oauthService, billingTracker, billingRepo, tcm, ragService), mcpMetadataService)
 			chat.Attach(mux, chat.NewService(logger, db, sessionManager, openRouter))
 			if slackClient.Enabled() {
 				slack.Attach(mux, slack.NewService(logger, db, sessionManager, encryptionClient, redisClient, slackClient, temporalClient, slack.Configurations{

server/database/sqlc.yaml

@@ -54,6 +54,10 @@ overrides:
         go_type:
           import: github.com/speakeasy-api/gram/server/internal/urn
           type: Tool
+      - db_type: "vector"
+        go_type:
+          import: github.com/pgvector/pgvector-go
+          type: Vector
 
 sql:
   - schema: schema.sql
@@ -341,3 +345,13 @@ sql:
         out: "../internal/resources/repo"
         sql_package: "pgx/v5"
         omit_unused_structs: true
+
+  - schema: schema.sql
+    queries: ../internal/rag/repo/queries.sql
+    engine: postgresql
+    gen:
+      go:
+        package: "repo"
+        out: "../internal/rag/repo"
+        sql_package: "pgx/v5"
+        omit_unused_structs: true

server/design/shared/toolset.go

@@ -50,6 +50,7 @@ var Toolset = Type("Toolset", func() {
 	Attribute("function_environment_variables", ArrayOf(FunctionEnvironmentVariable), "The function environment variables that are relevant to the toolset")
 	Attribute("tools", ArrayOf(Tool), "The tools in this toolset")
 	Attribute("tool_urns", ArrayOf(String), "The tool URNs in this toolset")
+	Attribute("toolset_version", Int64, "The version of the toolset (will be 0 if none exists)")
 	Attribute("resources", ArrayOf(Resource), "The resources in this toolset")
 	Attribute("resource_urns", ArrayOf(String), "The resource URNs in this toolset")
 
@@ -68,7 +69,7 @@ var Toolset = Type("Toolset", func() {
 		Description("When the toolset was last updated.")
 		Format(FormatDateTime)
 	})
-	Required("id", "project_id", "organization_id", "account_type", "name", "slug", "tools", "prompt_templates", "tool_urns", "resources", "resource_urns", "created_at", "updated_at")
+	Required("id", "project_id", "organization_id", "account_type", "name", "slug", "tools", "toolset_version", "prompt_templates", "tool_urns", "resources", "resource_urns", "created_at", "updated_at")
 })
 
 var ToolsetEntry = Type("ToolsetEntry", func() {