Improvements to Rag and design prompt

devops-mcp-server/main.go

@@ -30,6 +30,7 @@ import (
 var (
 	httpAddr  = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout. e.g. localhost:8080")
 	pprofAddr = flag.String("pprof", "", "if set, host the pprof debugging server at this address")
+	logFile   = "/tmp/devops-mcp-server.log"
 )
 
 func main() {
@@ -72,4 +73,16 @@ func main() {
 			log.Printf("Server failed: %v", err)
 		}
 	}
+
+	setupLogging()
+}
+
+func setupLogging() {
+	f, err := os.OpenFile(logFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
+	if err != nil {
+		// Fallback if file fails
+		log.SetOutput(os.Stderr)
+		return
+	}
+	log.SetOutput(f)
 }

devops-mcp-server/prompts/cicd.go

@@ -22,7 +22,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-//go:embed cicd_design_prompt.txt
+//go:embed cicd_design_prompt.md
 var promptCICDText string
 
 // Helps design and implement GCP CI/CD pipelines.
@@ -32,7 +32,7 @@ func DesignPrompt(ctx context.Context, server *mcp.Server) {
 			Description: "Helps design and implement GCP CI/CD pipelines.",
 			Messages: []*mcp.PromptMessage{
 				{
-					Role:    "user",
+					Role: "user",
 					Content: &mcp.TextContent{
 						Text: fmt.Sprintf(promptCICDText, req.Params.Arguments["query"]),
 					},
@@ -43,16 +43,15 @@ func DesignPrompt(ctx context.Context, server *mcp.Server) {
 
 	// Create a server with a single prompt.
 	prompt := &mcp.Prompt{
-		Name: "devops:design",
+		Name:  "devops:design",
 		Title: "Design and implement a Google Cloud based CI/CD pipeline.",
 		Arguments: []*mcp.PromptArgument{
 			{
 				Name:        "query",
-				Description: "CICD pipeline description, as explained by te user",
+				Description: "CICD pipeline description",
 				Required:    true,
 			},
 		},
 	}
-	server.AddPrompt(prompt,promptHandler)
+	server.AddPrompt(prompt, promptHandler)
 }
-

devops-mcp-server/prompts/cicd_design_prompt.txt → devops-mcp-server/prompts/cicd_design_prompt.md

@@ -4,10 +4,9 @@ You are a comprehensive Google Cloud DevOps Assistant. Your primary function is
 
 First, analyze the user's request to determine the primary intent.
 
-* If the intent is a high-level goal like **"build a pipeline," "design an architecture,"** or **"migrate my Jenkins pipeline,"** you must follow the two-stage **Workflow A: Design & Implement**.
-* If the intent is a direct, concrete command like **"create an artifact registry repo," "deploy to prod,"** or **"run the main-branch trigger,"** you must follow **Workflow B: Direct Action**.
+* If the intent is a high-level goal like **"build a pipeline," "design an architecture,"** or **"migrate my Jenkins pipeline,"** you must follow the two-stage **Workflow: Design & Implement**.
 
-## Workflow A: Design & Implement
+## Workflow: Design & Implement
 
 This workflow is for high-level, architectural tasks. It consists of a design phase followed by an implementation phase.
 
@@ -29,20 +28,10 @@ Once the user has approved the YAML plan, your sole purpose is to execute it by
 
 1.  **Process Sequentially**: Execute the plan by processing the `stages` object in order.
 2.  **Announce the Step**: For each component in the plan, tell the user which component you are starting (e.g., "Starting step: 'Build and Test'").
-3.  **Consult Knowledge Base**: Use the `query_knowledge` tool to find out how to implement the component based on its `type` and `name`.
-4.  **Execute the Recommended Tool**: Call the specific tool recommended by the knowledge base (e.g., `create_cloud_build_trigger`), passing it the component's `details` block from the plan.
-5.  **Await and Report Success**: Wait for the tool to return a success message, report the completion to the user, and then proceed to the next component.
+3.  **Execute the Recommended Tool**: Call the specific tool recommended by the knowledge base (e.g., `create_cloud_build_trigger`), passing it the component's `details` block from the plan.
+4.  **Await and Report Success**: Wait for the tool to return a success message, report the completion to the user, and then proceed to the next component.
 
 
-## Workflow B: Direct Action
-
-This workflow is for executing single, direct commands.
-
-1.  **Identify the Intent**: Determine the single action the user wants to perform (e.g., `create_artifact_registry_repo`).
-2.  **Gather Parameters**: Analyze the request to find all necessary parameters (e.g., `repo_name: "my-app-images"`).
-3.  **Clarify if Needed**: If any mandatory parameters are missing, you MUST ask the user for them before proceeding. Do not guess or make assumptions.
-4.  **Execute**: Call the single, correct tool to perform the action.
-
 
 ## Universal Protocols & Constraints
 

devops-mcp-server/prompts/deploy.go

@@ -22,7 +22,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-//go:embed deploy_prompt.txt
+//go:embed deploy_prompt.md
 var promptDeployText string
 
 // Helps deploy applications to GCP.
@@ -32,7 +32,7 @@ func DeployPrompt(ctx context.Context, server *mcp.Server) {
 			Description: "Helps deploy applications to GCP.",
 			Messages: []*mcp.PromptMessage{
 				{
-					Role:    "user",
+					Role: "user",
 					Content: &mcp.TextContent{
 						Text: fmt.Sprintf(promptDeployText, req.Params.Arguments["query"]),
 					},
@@ -43,7 +43,7 @@ func DeployPrompt(ctx context.Context, server *mcp.Server) {
 
 	// Create a server with a single prompt.
 	prompt := &mcp.Prompt{
-		Name: "devops:deploy",
+		Name:  "devops:deploy",
 		Title: "Deploy an application to GCP.",
 		Arguments: []*mcp.PromptArgument{
 			{
@@ -53,6 +53,5 @@ func DeployPrompt(ctx context.Context, server *mcp.Server) {
 			},
 		},
 	}
-	server.AddPrompt(prompt,promptHandler)
+	server.AddPrompt(prompt, promptHandler)
 }
-

devops-mcp-server/rag/rag.go

@@ -15,52 +15,48 @@
 package rag
 
 import (
+	"bytes"
 	"context"
 	"devops-mcp-server/auth"
+	_ "embed"
+	"encoding/json"
 	"fmt"
 	"log"
-	"os"
 	"sync"
 
 	chromem "github.com/philippgille/chromem-go"
 )
 
 var initOnce sync.Once
 
-// ListResult defines a generic struct to wrap a list of items.
-type ListResult[T any] struct {
-	Items []T `json:"items"`
-}
+//go:embed devops-rag.db
+var embeddedDB []byte
 
 type RagData struct {
 	DB        *chromem.DB
 	Pattern   *chromem.Collection
 	Knowledge *chromem.Collection
 }
 
+// Only expose what the LLM needs to read.
+type Result struct {
+	Content    string            `json:"content"`
+	Metadata   map[string]string `json:"metadata,omitempty"` // Source info
+	Similarity float32           `json:"relevance_score"`    // Helps LLM weigh confidence
+}
+
 var RagDB RagData
 
 // loadRAG performs the one-time initialization.
 func loadRAG(ctx context.Context) error {
-	dbFile := os.Getenv("RAG_DB_PATH")
 	RagDB = RagData{DB: chromem.NewDB()}
-	if len(dbFile) < 1 {
-		// RETURN AN ERROR
-		return fmt.Errorf("Env variable RAG_DB_PATH is not set for RAG data file")
-	}
-
-	//check if file exists
-	if _, err := os.Stat(dbFile); os.IsNotExist(err) {
-		// RETURN AN ERROR
-		return fmt.Errorf("RAG_DB_PATH file does not exist, skipping import: %v", dbFile)
-	}
-
-	err := RagDB.DB.ImportFromFile(dbFile, "")
+	reader := bytes.NewReader(embeddedDB)
+	err := RagDB.DB.ImportFromReader(reader, "")
 	if err != nil {
-		log.Printf("Unable to import from the RAG DB file:%s - %v", dbFile, err)
-		// This seems non-fatal based on your log, so we continue.
+		log.Printf("Unable to import from the RAG DB file: %v", err)
+		return err
 	}
-	log.Printf("IMPORTED from the RAG DB file:%s - %v", dbFile, len(RagDB.DB.ListCollections()))
+	log.Printf("IMPORTED from the RAG DB collections: %v", len(RagDB.DB.ListCollections()))
 
 	creds, err := auth.GetAuthToken(ctx)
 	if err != nil {
@@ -73,21 +69,17 @@ func loadRAG(ctx context.Context) error {
 		creds.Token,
 		creds.ProjectId,
 		chromem.EmbeddingModelVertexEnglishV4)
-
+	RagDB.Knowledge, err = RagDB.DB.GetOrCreateCollection("knowledge", nil, vertexEmbeddingFunc)
+	if err != nil {
+		return fmt.Errorf("Unable to get collection knowledge: %w", err)
+	}
+	log.Printf("LOADED collection knowledge: %v", RagDB.Pattern.Count())
 	RagDB.Pattern, err = RagDB.DB.GetOrCreateCollection("pattern", nil, vertexEmbeddingFunc)
 	if err != nil {
-		// RETURN AN ERROR
 		return fmt.Errorf("Unable to get collection pattern: %w", err)
 	}
 	log.Printf("LOADED collection pattern: %v", RagDB.Pattern.Count())
 
-	RagDB.Knowledge, err = RagDB.DB.GetOrCreateCollection("knowledge", nil, vertexEmbeddingFunc)
-	if err != nil {
-		// RETURN AN ERROR
-		return fmt.Errorf("Unable to get collection knowledge: %w", err)
-	}
-	log.Printf("LOADED collection knowledge: %v", RagDB.Knowledge.Count())
-
 	log.Print("RAG Init Completed!")
 	return nil // Success
 }
@@ -104,34 +96,46 @@ func GetRAG(ctx context.Context) (RagData, error) {
 	return RagDB, initErr
 }
 
-func QueryKnowledge(ctx context.Context, query string) (*ListResult[chromem.Result], error) {
-	ragDB, err := GetRAG(ctx)
+func (r *RagData) QueryPattern(ctx context.Context, query string) (string, error) {
+	results, err := r.Pattern.Query(ctx, query, 2, nil, nil)
 	if err != nil {
-		// Initialization failed, return the error.
-		return nil, fmt.Errorf("RAG system not initialized: %w", err)
+		log.Fatalf("Unable to Query collection pattern: %v", err)
 	}
-	if ragDB.Knowledge == nil {
-		return nil, fmt.Errorf("RagDB does not contain Knowledge collection, among %d collections", len(RagDB.DB.ListCollections()))
+	cleanResults := make([]Result, len(results))
+	for i, r := range results {
+		cleanResults[i] = Result{
+			Content:    r.Content,
+			Metadata:   r.Metadata,
+			Similarity: r.Similarity,
+		}
 	}
-	result, err := RagDB.Knowledge.Query(ctx, query, 3, nil, nil)
+
+	// Marshal to JSON
+	jsonData, err := json.Marshal(cleanResults)
 	if err != nil {
-		log.Fatalf("Unable to Query collection pattern: %v", err)
+		return "", fmt.Errorf("failed to marshal results: %w", err)
 	}
-	return &ListResult[chromem.Result]{Items: result}, nil
+	return string(jsonData), nil
 }
 
-func QueryPattern(ctx context.Context, query string) (*ListResult[chromem.Result], error) {
-	ragDB, err := GetRAG(ctx)
+func (r *RagData) Queryknowledge(ctx context.Context, query string) (string, error) {
+	results, err := r.Knowledge.Query(ctx, query, 2, nil, nil)
 	if err != nil {
-		// Initialization failed, return the error.
-		return nil, fmt.Errorf("RAG system not initialized: %w", err)
+		log.Fatalf("Unable to Query collection knowledge: %v", err)
 	}
-	if ragDB.Pattern == nil {
-		return nil, fmt.Errorf("RagDB does not contain Pattern collection, among %d collections", len(RagDB.DB.ListCollections()))
+	cleanResults := make([]Result, len(results))
+	for i, r := range results {
+		cleanResults[i] = Result{
+			Content:    r.Content,
+			Metadata:   r.Metadata,
+			Similarity: r.Similarity,
+		}
 	}
-	result, err := RagDB.Pattern.Query(ctx, query, 2, nil, nil)
+
+	// Marshal to JSON
+	jsonData, err := json.Marshal(cleanResults)
 	if err != nil {
-		log.Fatalf("Unable to Query collection pattern: %v", err)
+		return "", fmt.Errorf("failed to marshal results: %w", err)
 	}
-	return &ListResult[chromem.Result]{Items: result}, nil
+	return string(jsonData), nil
 }

devops-mcp-server/version.txt

@@ -1 +1 @@
-0.1.0
+0.1.1