-
Notifications
You must be signed in to change notification settings - Fork 63
mcp server for ml feature ops #338
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
a0d00kc
wants to merge
1
commit into
develop
Choose a base branch
from
feat/on-fs-mcp
base: develop
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| package main | ||
|
|
||
| import ( | ||
| "os" | ||
| "strings" | ||
|
|
||
| "github.com/Meesho/BharatMLStack/horizon/internal/configs" | ||
| mcpserver "github.com/Meesho/BharatMLStack/horizon/internal/mcp" | ||
| onlinefeaturestore "github.com/Meesho/BharatMLStack/horizon/internal/online-feature-store" | ||
| ofsConfig "github.com/Meesho/BharatMLStack/horizon/internal/online-feature-store/config" | ||
| ofsHandler "github.com/Meesho/BharatMLStack/horizon/internal/online-feature-store/handler" | ||
| "github.com/Meesho/BharatMLStack/horizon/pkg/etcd" | ||
| "github.com/Meesho/BharatMLStack/horizon/pkg/infra" | ||
| "github.com/Meesho/BharatMLStack/horizon/pkg/logger" | ||
| "github.com/mark3labs/mcp-go/server" | ||
| "github.com/rs/zerolog/log" | ||
| ) | ||
|
|
||
| // AppConfig holds the application configuration for the MCP server. | ||
| type AppConfig struct { | ||
| Configs configs.Configs | ||
| DynamicConfigs configs.DynamicConfigs | ||
| } | ||
|
|
||
| func (cfg *AppConfig) GetStaticConfig() interface{} { | ||
| return &cfg.Configs | ||
| } | ||
|
|
||
| func (cfg *AppConfig) GetDynamicConfig() interface{} { | ||
| return &cfg.DynamicConfigs | ||
| } | ||
|
|
||
| func main() { | ||
| var appConfig AppConfig | ||
|
|
||
| configs.InitConfig(&appConfig) | ||
| logger.Init(appConfig.Configs) | ||
| infra.InitDBConnectors(appConfig.Configs) | ||
| etcd.InitFromAppName(&ofsConfig.FeatureRegistry{}, appConfig.Configs.OnlineFeatureStoreAppName, appConfig.Configs) | ||
| onlinefeaturestore.Init(appConfig.Configs) | ||
|
|
||
| configHandler := ofsHandler.NewConfigHandler(1) | ||
| if configHandler == nil { | ||
| log.Fatal().Msg("failed to initialize online-feature-store config handler") | ||
| } | ||
|
|
||
| mcpSrv := mcpserver.NewServer(configHandler) | ||
|
|
||
| addr := os.Getenv("MCP_ADDR") | ||
| if addr == "" { | ||
| addr = ":8080" | ||
| } | ||
|
|
||
| var httpOpts []server.StreamableHTTPOption | ||
|
|
||
| if strings.EqualFold(os.Getenv("MCP_TLS_ENABLED"), "true") { | ||
| certFile := os.Getenv("MCP_TLS_CERT_FILE") | ||
| keyFile := os.Getenv("MCP_TLS_KEY_FILE") | ||
| if certFile == "" || keyFile == "" { | ||
| log.Fatal().Msg("MCP_TLS_CERT_FILE and MCP_TLS_KEY_FILE must be set when MCP_TLS_ENABLED=true") | ||
| } | ||
| httpOpts = append(httpOpts, server.WithTLSCert(certFile, keyFile)) | ||
| log.Info().Str("addr", addr).Msg("Starting Horizon MCP server with TLS (Streamable HTTP)") | ||
| } else { | ||
| log.Info().Str("addr", addr).Msg("Starting Horizon MCP server (Streamable HTTP)") | ||
| } | ||
|
|
||
| httpServer := server.NewStreamableHTTPServer(mcpSrv, httpOpts...) | ||
| if err := httpServer.Start(addr); err != nil { | ||
| log.Fatal().Err(err).Msg("MCP server error") | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| package mcp | ||
|
|
||
| import ( | ||
| "github.com/Meesho/BharatMLStack/horizon/internal/online-feature-store/handler" | ||
| "github.com/mark3labs/mcp-go/mcp" | ||
| "github.com/mark3labs/mcp-go/server" | ||
| ) | ||
|
|
||
| // NewServer creates and configures a new MCP server with all | ||
| // Horizon online-feature-store discovery tools registered. | ||
| func NewServer(configHandler handler.Config) *server.MCPServer { | ||
| s := server.NewMCPServer( | ||
| "horizon-ofs-mcp", | ||
| "1.0.0", | ||
| server.WithToolCapabilities(false), | ||
| server.WithRecovery(), | ||
| ) | ||
|
|
||
| tools := NewToolHandlers(configHandler) | ||
| registerTools(s, tools) | ||
| return s | ||
| } | ||
|
|
||
| func registerTools(s *server.MCPServer, t *ToolHandlers) { | ||
| listEntities := mcp.NewTool("list_entities", | ||
| mcp.WithDescription("List all registered entity labels in the online feature store"), | ||
| ) | ||
| s.AddTool(listEntities, t.ListEntities) | ||
|
|
||
| getEntityDetails := mcp.NewTool("get_entity_details", | ||
| mcp.WithDescription("Get detailed configuration for all entities including keys and cache settings"), | ||
| ) | ||
| s.AddTool(getEntityDetails, t.GetEntityDetails) | ||
|
|
||
| listFeatureGroups := mcp.NewTool("list_feature_groups", | ||
| mcp.WithDescription("List feature groups for a given entity with their labels and data types"), | ||
| mcp.WithString("entity_label", | ||
| mcp.Required(), | ||
| mcp.Description("The entity label to list feature groups for"), | ||
| ), | ||
| ) | ||
| s.AddTool(listFeatureGroups, t.ListFeatureGroups) | ||
|
|
||
| getFeatureGroupDetails := mcp.NewTool("get_feature_group_details", | ||
| mcp.WithDescription("Get detailed configuration for feature groups including features, versions, TTL, and cache settings"), | ||
| mcp.WithString("entity_label", | ||
| mcp.Required(), | ||
| mcp.Description("The entity label to retrieve feature group details for"), | ||
| ), | ||
| ) | ||
| s.AddTool(getFeatureGroupDetails, t.GetFeatureGroupDetails) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| package mcp | ||
|
|
||
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
|
|
||
| "github.com/Meesho/BharatMLStack/horizon/internal/online-feature-store/handler" | ||
| "github.com/mark3labs/mcp-go/mcp" | ||
| ) | ||
|
|
||
| // ToolHandlers holds the handler.Config dependency and implements | ||
| // all MCP tool handler functions for Horizon discovery tools. | ||
| type ToolHandlers struct { | ||
| config handler.Config | ||
| } | ||
|
|
||
| // NewToolHandlers creates a new ToolHandlers with the given config handler. | ||
| func NewToolHandlers(config handler.Config) *ToolHandlers { | ||
| return &ToolHandlers{config: config} | ||
| } | ||
|
|
||
| // ListEntities returns all registered entity labels. | ||
| func (t *ToolHandlers) ListEntities(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { | ||
| entities, err := t.config.GetAllEntities() | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to list entities: %v", err)), nil | ||
| } | ||
|
|
||
| data, err := json.Marshal(entities) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to marshal entities: %v", err)), nil | ||
| } | ||
| return mcp.NewToolResultText(string(data)), nil | ||
| } | ||
|
|
||
| // GetEntityDetails returns full configuration for all entities | ||
| // including keys, in-memory cache, and distributed cache settings. | ||
| func (t *ToolHandlers) GetEntityDetails(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { | ||
| entities, err := t.config.RetrieveEntities() | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve entity details: %v", err)), nil | ||
| } | ||
|
|
||
| data, err := json.Marshal(entities) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to marshal entity details: %v", err)), nil | ||
| } | ||
| return mcp.NewToolResultText(string(data)), nil | ||
| } | ||
|
|
||
| // ListFeatureGroups returns feature groups for a given entity | ||
| // with their labels and data types. | ||
| func (t *ToolHandlers) ListFeatureGroups(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { | ||
| entityLabel, err := request.RequireString("entity_label") | ||
| if err != nil { | ||
| return mcp.NewToolResultError("entity_label is required"), nil | ||
| } | ||
|
|
||
| featureGroups, err := t.config.RetrieveFeatureGroups(entityLabel) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to list feature groups: %v", err)), nil | ||
| } | ||
|
|
||
| type fgSummary struct { | ||
| Label string `json:"label"` | ||
| DataType string `json:"data-type"` | ||
| } | ||
| var summaries []fgSummary | ||
| if featureGroups != nil { | ||
| for _, fg := range *featureGroups { | ||
| summaries = append(summaries, fgSummary{ | ||
| Label: fg.FeatureGroupLabel, | ||
| DataType: string(fg.DataType), | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| data, err := json.Marshal(summaries) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to marshal feature groups: %v", err)), nil | ||
| } | ||
| return mcp.NewToolResultText(string(data)), nil | ||
| } | ||
|
|
||
| // GetFeatureGroupDetails returns full feature group configuration | ||
| // including features, active version, TTL, and cache settings. | ||
| func (t *ToolHandlers) GetFeatureGroupDetails(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { | ||
| entityLabel, err := request.RequireString("entity_label") | ||
| if err != nil { | ||
| return mcp.NewToolResultError("entity_label is required"), nil | ||
| } | ||
|
|
||
| featureGroups, err := t.config.RetrieveFeatureGroups(entityLabel) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to retrieve feature group details: %v", err)), nil | ||
| } | ||
|
|
||
| data, err := json.Marshal(featureGroups) | ||
| if err != nil { | ||
| return mcp.NewToolResultError(fmt.Sprintf("failed to marshal feature group details: %v", err)), nil | ||
| } | ||
| return mcp.NewToolResultText(string(data)), nil | ||
| } | ||
Binary file not shown.
Binary file not shown.
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
json.Marshal(nil)produces"null", not"[]"— initialize the slice.When
featureGroupsis nil or empty,summariesstaysnil, andjson.Marshal(nil)outputsnull. MCP tool consumers likely expect an empty JSON array.Proposed fix
🤖 Prompt for AI Agents