package prompts import ( "context" "embed" "fmt" "strings" "gitee.com/kwdb/kwdb-mcp-server/pkg/db" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" ) //go:embed docs/*.md var docsFS embed.FS // loadMarkdown loads a Markdown file from the embedded file system func loadMarkdown(filename string) (string, error) { content, err := docsFS.ReadFile(fmt.Sprintf("docs/%s", filename)) if err != nil { return "", fmt.Errorf("failed to read file %s: %v", filename, err) } return string(content), nil } // DBDescription is the database description provided to the LLM var DBDescription string // SyntaxGuide is the KWDB syntax guide var SyntaxGuide string // ReadExamplesTemplate is the template for read query examples var ReadExamplesTemplate string // WriteExamplesTemplate is the template for write query examples var WriteExamplesTemplate string // Use case guide content var ClusterManagementGuide string var DataMigrationGuide string var InstallationGuide string var PerformanceTuningGuide string var TroubleshootingGuide string // Additional guides var BackupRestoreGuide string var DBATemplate string // init function loads all Markdown files when the package is imported func init() { var err error // Load database description DBDescription, err = loadMarkdown("DBDescription.md") if err != nil { // Provide default content if file loading fails DBDescription = "# KWDB Database\n\nKWDB is a distributed SQL database compatible with PostgreSQL and CockroachDB." fmt.Printf("Warning: Failed to load database description file: %v\n", err) } // Load syntax guide SyntaxGuide, err = loadMarkdown("SyntaxGuide.md") if err != nil { // Provide default content if file loading fails SyntaxGuide = "# KWDB SQL Syntax Guide\n\nKWDB supports standard SQL syntax compatible with PostgreSQL and CockroachDB." fmt.Printf("Warning: Failed to load syntax guide file: %v\n", err) } // Load read query examples template ReadExamplesTemplate, err = loadMarkdown("ReadExamples.md") if err != nil { // Provide default content if file loading fails ReadExamplesTemplate = "" fmt.Printf("Warning: Failed to load read examples file: %v\n", err) } // Load write query examples template WriteExamplesTemplate, err = loadMarkdown("WriteExamples.md") if err != nil { // Provide default content if file loading fails WriteExamplesTemplate = "" fmt.Printf("Warning: Failed to load write examples file: %v\n", err) } // Load use case guides // Note: These files serve as templates for generating prompts // When implementing auto-generation logic, refer to these templates // and follow the same structure for new prompt templates // Cluster management guide ClusterManagementGuide, err = loadMarkdown("ClusterManagementGuide.md") if err != nil { // Provide default content if file loading fails ClusterManagementGuide = "" fmt.Printf("Warning: Failed to load cluster management guide file: %v\n", err) } // Data migration guide DataMigrationGuide, err = loadMarkdown("DataMigrationGuide.md") if err != nil { // Provide default content if file loading fails DataMigrationGuide = "" fmt.Printf("Warning: Failed to load data migration guide file: %v\n", err) } // Installation guide InstallationGuide, err = loadMarkdown("InstallationGuide.md") if err != nil { // Provide default content if file loading fails InstallationGuide = "" fmt.Printf("Warning: Failed to load installation guide file: %v\n", err) } // Performance tuning guide PerformanceTuningGuide, err = loadMarkdown("PerformanceTuningGuide.md") if err != nil { // Provide default content if file loading fails PerformanceTuningGuide = "" fmt.Printf("Warning: Failed to load performance tuning guide file: %v\n", err) } // Troubleshooting guide TroubleshootingGuide, err = loadMarkdown("TroubleShootingGuide.md") if err != nil { // Provide default content if file loading fails TroubleshootingGuide = "" fmt.Printf("Warning: Failed to load troubleshooting guide file: %v\n", err) } // Backup and restore guide BackupRestoreGuide, err = loadMarkdown("BackupRestoreGuide.md") if err != nil { // Provide default content if file loading fails BackupRestoreGuide = "" fmt.Printf("Warning: Failed to load backup and restore guide file: %v\n", err) } // DBA template DBATemplate, err = loadMarkdown("DBATemplate.md") if err != nil { // Provide default content if file loading fails DBATemplate = "" fmt.Printf("Warning: Failed to load DBA template file: %v\n", err) } } // GetReadExampleQueries returns example read queries for a table // This function demonstrates how to use the ReadExamplesTemplate // to generate table-specific queries by replacing placeholders func GetReadExampleQueries(tableName string) []string { // Replace {table} placeholder with actual table name content := strings.ReplaceAll(ReadExamplesTemplate, "{table}", tableName) // Split content by lines and remove "- " prefix lines := strings.Split(content, "\n") var queries []string for _, line := range lines { line = strings.TrimSpace(line) if line == "" { continue } line = strings.TrimPrefix(line, "- ") queries = append(queries, line) } return queries } // GetWriteExampleQueries returns example write queries for a table // This function demonstrates how to use the WriteExamplesTemplate // to generate table-specific queries by replacing placeholders func GetWriteExampleQueries(tableName string) []string { // Replace {table} placeholder with actual table name content := strings.ReplaceAll(WriteExamplesTemplate, "{table}", tableName) // Split content by lines and remove "- " prefix lines := strings.Split(content, "\n") var queries []string for _, line := range lines { line = strings.TrimSpace(line) if line == "" { continue } line = strings.TrimPrefix(line, "- ") queries = append(queries, line) } return queries } // RegisterPrompts registers all prompts with the MCP server func RegisterPrompts(s *server.MCPServer) { // Register syntax guide prompt registerSyntaxGuidePrompt(s) // Register database description prompt registerDBDescriptionPrompt(s) // Register use case prompts registerUseCasePrompts(s) } // registerSyntaxGuidePrompt registers the parameterized syntax guide prompt func registerSyntaxGuidePrompt(s *server.MCPServer) { // Create syntax guide prompt with parameter support syntaxGuidePrompt := mcp.NewPrompt("syntax_guide", mcp.WithPromptDescription("KWDB (KaiwuDB) syntax guide and examples. Optional parameters: 'database' and 'table' for table-specific guidance"), mcp.WithArgument("database", mcp.ArgumentDescription("Database name to provide specific table information")), mcp.WithArgument("table", mcp.ArgumentDescription("Table name to provide specific table schema and examples")), ) // Add parameterized syntax guide prompt handler s.AddPrompt(syntaxGuidePrompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { arguments := request.Params.Arguments // Extract parameters - Arguments is already map[string]string database := arguments["database"] table := arguments["table"] // Base content baseContent := fmt.Sprintf(`You are a SQL expert specializing in KWDB (KaiwuDB). Help users understand the syntax and capabilities of the database. %s`, SyntaxGuide) // Enhanced content for specific table tableContent := "" if table != "" { // Get table schema information if columns, err := db.GetTableColumnsWithContext(ctx, table); err == nil && len(columns) > 0 { tableContent += fmt.Sprintf("\n\n## Table Schema for '%s'\n", table) for _, col := range columns { columnName, _ := col["column_name"].(string) dataType, _ := col["data_type"].(string) isNullable, _ := col["is_nullable"].(string) columnDefault, _ := col["column_default"].(string) if columnName != "" && dataType != "" { tableContent += fmt.Sprintf("- **%s**: %s", columnName, dataType) if isNullable == "NO" { tableContent += " (NOT NULL)" } if columnDefault != "" { tableContent += fmt.Sprintf(" DEFAULT %s", columnDefault) } tableContent += "\n" } } // Add example queries for this table readExamples := GetReadExampleQueries(table) if len(readExamples) > 0 { tableContent += fmt.Sprintf("\n## Example Read Queries for '%s'\n", table) for _, example := range readExamples { if example != "" { tableContent += fmt.Sprintf("```sql\n%s\n```\n\n", example) } } } writeExamples := GetWriteExampleQueries(table) if len(writeExamples) > 0 { tableContent += fmt.Sprintf("\n## Example Write Queries for '%s'\n", table) for _, example := range writeExamples { if example != "" { tableContent += fmt.Sprintf("```sql\n%s\n```\n\n", example) } } } } } // Database information dbContent := "" if database != "" { if dbInfo, err := db.GetDatabaseInfoByName(database); err == nil { dbContent += fmt.Sprintf("\n\n## Database Information for '%s'\n", database) dbContent += fmt.Sprintf("- **Name**: %s\n", dbInfo.Name) dbContent += fmt.Sprintf("- **Version**: %s\n", dbInfo.Version) dbContent += fmt.Sprintf("- **Engine Type**: %s\n", dbInfo.EngineType) if dbInfo.Comment != "" { dbContent += fmt.Sprintf("- **Comment**: %s\n", dbInfo.Comment) } // Add properties if available for key, value := range dbInfo.Properties { dbContent += fmt.Sprintf("- **%s**: %v\n", key, value) } } } // Construct title title := "KWDB (KaiwuDB) Syntax Guide" if table != "" { title += fmt.Sprintf(" - Table: %s", table) } if database != "" { title += fmt.Sprintf(" - Database: %s", database) } return mcp.NewGetPromptResult( title, []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent(baseContent+tableContent+dbContent), ), }, ), nil }) } // registerDBDescriptionPrompt registers the database description prompt func registerDBDescriptionPrompt(s *server.MCPServer) { // Create database description prompt dbDescriptionPrompt := mcp.NewPrompt("db_description", mcp.WithPromptDescription("KWDB (KaiwuDB) database description and capabilities"), ) // Add database description prompt handler s.AddPrompt(dbDescriptionPrompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB (KaiwuDB) Database Description", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a database expert specializing in KWDB (KaiwuDB). Help users understand the capabilities and features of the database."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(DBDescription), ), }, ), nil }) } // registerUseCasePrompts registers all use case prompts // When implementing auto-generation logic for prompts, follow this pattern: // 1. Create a new Markdown template file in the docs/ directory // 2. Add a variable to store the content // 3. Load the file in the init() function // 4. Create a registration function similar to the ones below // 5. Add the registration function call to this function func registerUseCasePrompts(s *server.MCPServer) { // Register cluster management prompt registerClusterManagementPrompt(s) // Register data migration prompt registerDataMigrationPrompt(s) // Register installation prompt registerInstallationPrompt(s) // Register performance tuning prompt registerPerformanceTuningPrompt(s) // Register troubleshooting prompt registerTroubleshootingPrompt(s) // Register backup and restore prompt registerBackupRestorePrompt(s) // Register DBA template prompt registerDBATemplatePrompt(s) // Note: To add new use case prompts, follow these steps: // 1. Create a new Markdown file in the docs/ directory, e.g., NewUseCase.md // 2. Add a new variable in prompts.go, e.g., var NewUseCaseGuide string // 3. Load the new Markdown file in the init() function // 4. Create a new registration function, e.g., registerNewUseCasePrompt(s *server.MCPServer) // 5. Call the new registration function in the registerUseCasePrompts function // 6. Update the README.md to document the new prompt } // registerClusterManagementPrompt registers the cluster management prompt func registerClusterManagementPrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("cluster_management", mcp.WithPromptDescription("KWDB Cluster Management Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Cluster Management Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB cluster management expert. Help users understand and implement KWDB cluster management operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(ClusterManagementGuide), ), }, ), nil }) } // registerDataMigrationPrompt registers the data migration prompt func registerDataMigrationPrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("data_migration", mcp.WithPromptDescription("KWDB Data Migration Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Data Migration Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB data migration expert. Help users understand and implement data migration operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(DataMigrationGuide), ), }, ), nil }) } // registerInstallationPrompt registers the installation prompt func registerInstallationPrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("installation", mcp.WithPromptDescription("KWDB Installation and Deployment Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Installation and Deployment Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB installation and deployment expert. Help users understand and implement KWDB installation and deployment operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(InstallationGuide), ), }, ), nil }) } // registerPerformanceTuningPrompt registers the performance tuning prompt func registerPerformanceTuningPrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("performance_tuning", mcp.WithPromptDescription("KWDB Performance Tuning Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Performance Tuning Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB performance tuning expert. Help users understand and implement performance optimization operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(PerformanceTuningGuide), ), }, ), nil }) } // registerTroubleshootingPrompt registers the troubleshooting prompt func registerTroubleshootingPrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("troubleshooting", mcp.WithPromptDescription("KWDB Troubleshooting Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Troubleshooting Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB troubleshooting expert. Help users diagnose and resolve KWDB issues."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(TroubleshootingGuide), ), }, ), nil }) } // registerBackupRestorePrompt registers the backup and restore prompt // This is an example of how to add a new prompt using the template pattern func registerBackupRestorePrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("backup_restore", mcp.WithPromptDescription("KWDB Backup and Restore Guide and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Backup and Restore Guide", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB backup and restore expert. Help users understand and implement backup and restore operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(BackupRestoreGuide), ), }, ), nil }) } // registerDBATemplatePrompt registers the DBA template prompt // This is an example of how to add a new prompt using the template pattern func registerDBATemplatePrompt(s *server.MCPServer) { prompt := mcp.NewPrompt("dba_template", mcp.WithPromptDescription("KWDB Database Administration Template and Best Practices"), ) s.AddPrompt(prompt, func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { return mcp.NewGetPromptResult( "KWDB Database Administration Template", []mcp.PromptMessage{ mcp.NewPromptMessage( mcp.RoleUser, mcp.NewTextContent("You are a KWDB database administration expert. Help users understand and implement database administration operations."), ), mcp.NewPromptMessage( mcp.RoleAssistant, mcp.NewTextContent(DBATemplate), ), }, ), nil }) }