Skip to content

Commit bf7e9e7

Browse files
committed
feat(prompts): consolidate prompt definitions and handlers
This change refactors the prompt registration and handling logic. Instead of registering each prompt individually in main.go and defining separate handler functions for each in prompt_handlers.go, a new `PromptDefinition` struct is introduced in prompts.go. This struct encapsulates the prompt name, description, and its associated system prompt. The `GeminiServer` now has a single generic `promptHandler` method that takes a `PromptDefinition`. This handler dynamically constructs the instructions using the provided system prompt and the problem statement. The registration of all prompts is now done by iterating over a slice of `PromptDefinition`s, making the code more concise and easier to manage. This consolidation simplifies the addition of new prompts and reduces code duplication.
1 parent 66896db commit bf7e9e7

File tree

5 files changed

+180
-296
lines changed

5 files changed

+180
-296
lines changed

main.go

Lines changed: 6 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -368,30 +368,12 @@ func setupGeminiServer(ctx context.Context, mcpServer *server.MCPServer, config
368368
mcpServer.AddTool(GeminiModelsTool, wrapHandlerWithLogger(geminiSvc.GeminiModelsHandler, "gemini_models", logger))
369369
logger.Info("Registered tool: gemini_models")
370370

371-
// Register prompts
372-
mcpServer.AddPrompt(CodeReviewPrompt, wrapPromptHandlerWithLogger(geminiSvc.CodeReviewHandler, "code_review", logger))
373-
logger.Info("Registered prompt: code_review")
374-
375-
mcpServer.AddPrompt(ExplainCodePrompt, wrapPromptHandlerWithLogger(geminiSvc.ExplainCodeHandler, "explain_code", logger))
376-
logger.Info("Registered prompt: explain_code")
377-
378-
mcpServer.AddPrompt(DebugHelpPrompt, wrapPromptHandlerWithLogger(geminiSvc.DebugHelpHandler, "debug_help", logger))
379-
logger.Info("Registered prompt: debug_help")
380-
381-
mcpServer.AddPrompt(RefactorSuggestionsPrompt, wrapPromptHandlerWithLogger(geminiSvc.RefactorSuggestionsHandler, "refactor_suggestions", logger))
382-
logger.Info("Registered prompt: refactor_suggestions")
383-
384-
mcpServer.AddPrompt(ArchitectureAnalysisPrompt, wrapPromptHandlerWithLogger(geminiSvc.ArchitectureAnalysisHandler, "architecture_analysis", logger))
385-
logger.Info("Registered prompt: architecture_analysis")
386-
387-
mcpServer.AddPrompt(DocGeneratePrompt, wrapPromptHandlerWithLogger(geminiSvc.DocGenerateHandler, "doc_generate", logger))
388-
logger.Info("Registered prompt: doc_generate")
389-
390-
mcpServer.AddPrompt(TestGeneratePrompt, wrapPromptHandlerWithLogger(geminiSvc.TestGenerateHandler, "test_generate", logger))
391-
logger.Info("Registered prompt: test_generate")
392-
393-
mcpServer.AddPrompt(SecurityAnalysisPrompt, wrapPromptHandlerWithLogger(geminiSvc.SecurityAnalysisHandler, "security_analysis", logger))
394-
logger.Info("Registered prompt: security_analysis")
371+
// Register all prompts from the definitions
372+
for _, p := range Prompts {
373+
handler := geminiSvc.promptHandler(p)
374+
mcpServer.AddPrompt(*p.Prompt, wrapPromptHandlerWithLogger(handler, p.Name, logger))
375+
logger.Info("Registered prompt: %s", p.Name)
376+
}
395377

396378
// Log file handling configuration
397379
logger.Info("File handling: max size %s, allowed types: %v",

prompt_handlers.go

Lines changed: 17 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"fmt"
66

77
"github.com/mark3labs/mcp-go/mcp"
8+
"github.com/mark3labs/mcp-go/server"
89
)
910

1011
// createTaskInstructions generates the instructional text for the MCP client
@@ -20,136 +21,22 @@ func createTaskInstructions(problemStatement, systemPrompt string) string {
2021
"<problem_statement>\n%s\n</problem_statement>", systemPrompt, problemStatement)
2122
}
2223

23-
// genericPromptHandler is a generic handler for all prompts
24-
func genericPromptHandler(ctx context.Context, req mcp.GetPromptRequest, systemPrompt string) (*mcp.GetPromptResult, error) {
25-
problemStatement, ok := req.Params.Arguments["problem_statement"]
26-
if !ok || problemStatement == "" {
27-
return nil, fmt.Errorf("problem_statement argument is required")
28-
}
29-
30-
instructions := createTaskInstructions(problemStatement, systemPrompt)
31-
32-
return mcp.NewGetPromptResult(
33-
req.Params.Name,
34-
[]mcp.PromptMessage{
35-
mcp.NewPromptMessage(mcp.RoleAssistant, mcp.NewTextContent(instructions)),
36-
},
37-
), nil
38-
}
39-
40-
// CodeReviewHandler handles the code_review prompt
41-
func (s *GeminiServer) CodeReviewHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
42-
systemPrompt := `You are an expert code reviewer with years of experience in software engineering. Your task is to conduct a thorough analysis of the provided code.
43-
44-
Focus on the following areas:
45-
- **Code Quality & Best Practices:** Adherence to language-specific idioms, code formatting, and established best practices.
46-
- **Potential Bugs:** Logical errors, race conditions, null pointer issues, and other potential bugs.
47-
- **Security Vulnerabilities:** Identify any potential security risks, such as injection vulnerabilities, insecure data handling, or authentication/authorization flaws. Follow OWASP Top 10 guidelines.
48-
- **Performance Concerns:** Look for inefficient algorithms, memory leaks, or other performance bottlenecks.
49-
- **Maintainability & Readability:** Assess the code's clarity, modularity, and ease of maintenance.
50-
51-
Provide specific, actionable feedback. For each issue, include the file path (if available), the relevant line number(s), and a clear explanation of the problem and your suggested improvement.`
52-
return genericPromptHandler(ctx, req, systemPrompt)
53-
}
54-
55-
// ExplainCodeHandler handles the explain_code prompt
56-
func (s *GeminiServer) ExplainCodeHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
57-
systemPrompt := `You are an expert software engineer and a skilled educator. Your goal is to explain the provided code in a clear, comprehensive, and easy-to-understand manner.
24+
// promptHandler is the generic handler for all prompts
25+
func (s *GeminiServer) promptHandler(p *PromptDefinition) server.PromptHandlerFunc {
26+
return func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
27+
problemStatement, ok := req.Params.Arguments["problem_statement"]
28+
if !ok || problemStatement == "" {
29+
return nil, fmt.Errorf("problem_statement argument is required")
30+
}
5831

59-
Structure your explanation as follows:
60-
1. **High-Level Overview:** Start with a summary of what the code does and its primary purpose.
61-
2. **Detailed Breakdown:** Go through the code section by section, explaining the logic, algorithms, and data structures used.
62-
3. **Key Concepts:** Highlight any important design patterns, architectural decisions, or programming concepts demonstrated in the code.
63-
4. **Usage:** If applicable, provide a simple example of how to use the code.
32+
systemPrompt := p.SystemPrompt.GetSystemPrompt()
33+
instructions := createTaskInstructions(problemStatement, systemPrompt)
6434

65-
Tailor the complexity of your explanation to be suitable for an intermediate-level developer.`
66-
return genericPromptHandler(ctx, req, systemPrompt)
67-
}
68-
69-
// DebugHelpHandler handles the debug_help prompt
70-
func (s *GeminiServer) DebugHelpHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
71-
systemPrompt := `You are an expert debugger. Your mission is to analyze the provided code and the user's problem description to identify the root cause of a bug and suggest a solution.
72-
73-
Follow this systematic debugging process:
74-
1. **Analyze the Code:** Carefully review the provided code for potential logical errors, incorrect assumptions, or other issues related to the problem description.
75-
2. **Identify the Root Cause:** Based on your analysis, pinpoint the most likely cause of the bug.
76-
3. **Propose a Fix:** Provide a specific, corrected code snippet to fix the bug.
77-
4. **Explain the Solution:** Clearly explain why the bug occurred and why your proposed solution resolves it.`
78-
return genericPromptHandler(ctx, req, systemPrompt)
79-
}
80-
81-
// RefactorSuggestionsHandler handles the refactor_suggestions prompt
82-
func (s *GeminiServer) RefactorSuggestionsHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
83-
systemPrompt := `You are an expert software architect specializing in code modernization and refactoring. Your task is to analyze the provided code and suggest concrete improvements.
84-
85-
Your suggestions should focus on:
86-
- **Improving Code Structure:** Enhancing modularity, separation of concerns, and overall organization.
87-
- **Applying Design Patterns:** Identifying opportunities to use appropriate design patterns to solve common problems.
88-
- **Increasing Readability & Maintainability:** Making the code easier to understand and modify in the future.
89-
- **Optimizing Performance:** Where applicable, suggest changes to improve efficiency without sacrificing clarity.
90-
91-
For each suggestion, provide a code example demonstrating the change and explain the benefits of the proposed refactoring.`
92-
return genericPromptHandler(ctx, req, systemPrompt)
93-
}
94-
95-
// ArchitectureAnalysisHandler handles the architecture_analysis prompt
96-
func (s *GeminiServer) ArchitectureAnalysisHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
97-
systemPrompt := `You are a seasoned software architect. Your task is to conduct a high-level analysis of the provided codebase to understand its architecture.
98-
99-
Your analysis should cover:
100-
- **Overall Design:** Describe the main architectural pattern (e.g., Monolith, Microservices, MVC, etc.).
101-
- **Component Breakdown:** Identify the key components, their responsibilities, and how they interact.
102-
- **Data Flow:** Explain how data flows through the system.
103-
- **Dependencies:** List the major external dependencies and their roles.
104-
- **Potential Issues:** Highlight any potential architectural weaknesses, bottlenecks, or areas for improvement regarding scalability, maintainability, or security.
105-
106-
Provide a clear and concise summary of the architecture.`
107-
return genericPromptHandler(ctx, req, systemPrompt)
108-
}
109-
110-
// DocGenerateHandler handles the doc_generate prompt
111-
func (s *GeminiServer) DocGenerateHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
112-
systemPrompt := `You are a professional technical writer. Your task is to generate clear, concise, and comprehensive documentation for the provided code.
113-
114-
The documentation should be in Markdown format and include the following sections for each major component or function:
115-
- **Purpose:** A brief description of what the code does.
116-
- **Parameters:** A list of all input parameters, their types, and a description of each.
117-
- **Return Value:** A description of what the function or component returns.
118-
- **Usage Example:** A simple code snippet demonstrating how to use the code.
119-
120-
Ensure the documentation is accurate and easy for other developers to understand.`
121-
return genericPromptHandler(ctx, req, systemPrompt)
122-
}
123-
124-
// TestGenerateHandler handles the test_generate prompt
125-
func (s *GeminiServer) TestGenerateHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
126-
systemPrompt := `You are a test engineering expert. Your task is to generate comprehensive unit tests for the provided code.
127-
128-
The generated tests should:
129-
- Be written using the standard testing library for the given language.
130-
- Cover happy-path scenarios, edge cases, and error conditions.
131-
- Follow best practices for testing, including clear test descriptions, and proper assertions.
132-
- Be easy to read and maintain.
133-
134-
For each function or method, provide a set of corresponding test cases.`
135-
return genericPromptHandler(ctx, req, systemPrompt)
136-
}
137-
138-
// SecurityAnalysisHandler handles the security_analysis prompt
139-
func (s *GeminiServer) SecurityAnalysisHandler(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
140-
systemPrompt := `You are a cybersecurity expert specializing in secure code review. Your task is to analyze the provided code for security vulnerabilities and risks.
141-
142-
Focus on identifying common vulnerabilities, including but not limited to:
143-
- Injection attacks (SQL, Command, etc.)
144-
- Cross-Site Scripting (XSS)
145-
- Insecure Deserialization
146-
- Broken Authentication and Access Control
147-
- Security Misconfiguration
148-
- Sensitive Data Exposure
149-
150-
For each vulnerability you identify, provide:
151-
- A description of the vulnerability and its potential impact.
152-
- The file path and line number where the vulnerability exists.
153-
- A clear recommendation on how to remediate the vulnerability, including a corrected code snippet where possible.`
154-
return genericPromptHandler(ctx, req, systemPrompt)
35+
return mcp.NewGetPromptResult(
36+
req.Params.Name,
37+
[]mcp.PromptMessage{
38+
mcp.NewPromptMessage(mcp.RoleAssistant, mcp.NewTextContent(instructions)),
39+
},
40+
), nil
41+
}
15542
}

0 commit comments

Comments
 (0)