Add rag mcp tools

devops-mcp-server/go.mod

@@ -6,6 +6,7 @@ require (
 	cloud.google.com/go/artifactregistry v1.17.2
 	cloud.google.com/go/auth v0.17.0
 	cloud.google.com/go/cloudbuild v1.23.1
+	cloud.google.com/go/iam v1.5.3
 	cloud.google.com/go/resourcemanager v1.10.7
 	cloud.google.com/go/run v1.12.1
 	cloud.google.com/go/storage v1.57.0
@@ -26,7 +27,6 @@ require (
 	cloud.google.com/go v0.123.0 // indirect
 	cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect
 	cloud.google.com/go/compute/metadata v0.9.0 // indirect
-	cloud.google.com/go/iam v1.5.3 // indirect
 	cloud.google.com/go/longrunning v0.7.0 // indirect
 	cloud.google.com/go/monitoring v1.24.3 // indirect
 	cyphar.com/go-pathrs v0.2.1 // indirect

devops-mcp-server/rag/client/ragclient.go

@@ -0,0 +1,153 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ragclient
+
+import (
+	"bytes"
+	"context"
+	"devops-mcp-server/auth"
+	_ "embed"
+	"encoding/json"
+	"fmt"
+	"log"
+
+	chromem "github.com/philippgille/chromem-go"
+)
+
+//go:embed devops-rag.db
+var embeddedDB []byte
+
+type RagClientImpl 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
+}
+
+type RagClient interface {
+	Queryknowledge(ctx context.Context, query string) (string, error)
+	QueryPatterns(ctx context.Context, query string) (string, error)
+}
+
+// loadRAG performs the one-time initialization.
+func loadRAG(ctx context.Context) (RagClient, error) {
+	ragClient := &RagClientImpl{DB: chromem.NewDB()}
+	reader := bytes.NewReader(embeddedDB)
+	err := ragClient.DB.ImportFromReader(reader, "")
+	if err != nil {
+		log.Printf("Unable to import from the RAG DB file: %v", err)
+		return nil, err
+	}
+	log.Printf("IMPORTED from the RAG DB collections: %v", len(ragClient.DB.ListCollections()))
+
+	creds, err := auth.GetAuthToken(ctx)
+	if err != nil {
+		log.Printf("Error: Google Cloud account is required: %v", err)
+		// RETURN AN ERROR
+		return nil, fmt.Errorf("Google Cloud account is required: %w", err)
+	}
+
+	vertexEmbeddingFunc := chromem.NewEmbeddingFuncVertex(
+		creds.Token,
+		creds.ProjectId,
+		chromem.EmbeddingModelVertexEnglishV4)
+	ragClient.Knowledge, err = ragClient.DB.GetOrCreateCollection("knowledge", nil, vertexEmbeddingFunc)
+	if err != nil {
+		return nil, fmt.Errorf("Unable to get collection knowledge: %w", err)
+	}
+	log.Printf("LOADED collection knowledge: %v", ragClient.Knowledge.Count())
+	ragClient.Pattern, err = ragClient.DB.GetOrCreateCollection("pattern", nil, vertexEmbeddingFunc)
+	if err != nil {
+		return nil, fmt.Errorf("Unable to get collection pattern: %w", err)
+	}
+	log.Printf("LOADED collection pattern: %v", ragClient.Pattern.Count())
+
+	log.Print("RAG Init Completed!")
+	return ragClient, nil // Success
+}
+
+// contextKey is a private type to use as a key for context values.
+type contextKey string
+
+const (
+	ragClientKey contextKey = "ragClient"
+)
+
+// ClientFrom returns the RagClient stored in the context, if any.
+func ClientFrom(ctx context.Context) (RagClient, bool) {
+	client, ok := ctx.Value(ragClientKey).(RagClient)
+	return client, ok
+}
+
+// ContextWithClient returns a new context with the provided RagClient.
+func ContextWithClient(ctx context.Context, client RagClient) context.Context {
+	return context.WithValue(ctx, ragClientKey, client)
+}
+
+// NewClient creates a new Client.
+func NewClient(ctx context.Context) (RagClient, error) {
+	return loadRAG(ctx)
+}
+
+
+func (r *RagClientImpl) QueryPatterns(ctx context.Context, query string) (string, error) {
+	results, err := r.Pattern.Query(ctx, query, 2, nil, nil)
+	if err != nil {
+		log.Fatalf("Unable to Query collection pattern: %v", err)
+	}
+	cleanResults := make([]Result, len(results))
+	for i, r := range results {
+		cleanResults[i] = Result{
+			Content:    r.Content,
+			Metadata:   r.Metadata,
+			Similarity: r.Similarity,
+		}
+	}
+
+	// Marshal to JSON
+	jsonData, err := json.Marshal(cleanResults)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal results: %w", err)
+	}
+	return string(jsonData), nil
+}
+
+func (r *RagClientImpl) Queryknowledge(ctx context.Context, query string) (string, error) {
+	results, err := r.Knowledge.Query(ctx, query, 2, nil, nil)
+	if err != nil {
+		log.Fatalf("Unable to Query collection knowledge: %v", err)
+	}
+	cleanResults := make([]Result, len(results))
+	for i, r := range results {
+		cleanResults[i] = Result{
+			Content:    r.Content,
+			Metadata:   r.Metadata,
+			Similarity: r.Similarity,
+		}
+	}
+
+	// Marshal to JSON
+	jsonData, err := json.Marshal(cleanResults)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal results: %w", err)
+	}
+	return string(jsonData), nil
+}