diff --git a/.github/workflows/nightly-release.yml b/.github/workflows/nightly-release.yml index d69aaad..76f31bc 100644 --- a/.github/workflows/nightly-release.yml +++ b/.github/workflows/nightly-release.yml @@ -1,67 +1,98 @@ name: Nightly Release on: - schedule: - - cron: '0 0 * * *' # Runs daily at midnight UTC - workflow_dispatch: # Allows manual triggering of the workflow + schedule: + - cron: "0 0 * * *" # Runs daily at midnight UTC + workflow_dispatch: # Allows manual triggering of the workflow + +env: + RELEASE: nightly jobs: - release: - runs-on: ubuntu-latest - defaults: - run: - working-directory: devops-mcp-server - strategy: - matrix: - goos: [linux, windows, darwin] - goarch: [amd64, arm64] - exclude: - - goos: windows - goarch: arm64 - - goos: darwin - goarch: amd64 - - steps: - - uses: actions/checkout@v3 - - - name: prep-release - run: | - mkdir -p release - cp ../gemini-extension.json release/ - cp ../README.md release/ - cp ../local-rag/devops-rag.db release/ - - - name: Set up Go - uses: actions/setup-go@v3 + prep-release: + uses: ./.github/workflows/prep-release.yaml with: - go-version: '1.24' - - - name: Run Go Vet - run: go vet ./... + release: nightly - - name: Run Go Fmt - run: go fmt ./... + release: + needs: prep-release + runs-on: ubuntu-latest + defaults: + run: + working-directory: devops-mcp-server + strategy: + matrix: + goos: [linux, windows, darwin] + goarch: [amd64, arm64] + exclude: + - goos: windows + goarch: arm64 + - goos: darwin + goarch: amd64 - - name: Run Go Tests - run: go test ./... + steps: + - uses: actions/checkout@v3 - # - name: Install golangci-lint - # run: go install github.com/golangci/golangci-lint/cmd/golangci-lint@latest + - name: prep-release + run: | + mkdir -p release dist + cp ../gemini-extension.json release/ + cp ../README.md release/ + cp ../local-rag/devops-rag.db release/ - # - name: Run golangci-lint - # run: golangci-lint run + - name: Set up Go + uses: actions/setup-go@v3 + with: + go-version: "1.24" - - name: Build - run: | - export GOOS=${{ matrix.goos }} - export GOARCH=${{ matrix.goarch }} - go build -o "release/devops-mcp-server-${GOOS}-${GOARCH}" + - name: Run Go Vet + run: go vet ./... - - name: Create release assets - run: | - tar -czvf "release/devops-mcp-server-${{ matrix.goos }}-${{ matrix.goarch }}.tar.gz" -C release "devops-mcp-server-${{ matrix.goos }}-${{ matrix.goarch }}" + - name: Run Go Fmt + run: go fmt ./... - - name: Create GitHub Release - uses: softprops/action-gh-release@v1 - with: - files: release/*.tar.gz + # - name: Run Go Tests + # run: go test ./... + + # - name: Install golangci-lint + # run: go install github.com/golangci/golangci-lint/cmd/golangci-lint@latest + + # - name: Run golangci-lint + # run: golangci-lint run + + - name: Build + run: | + export GOOS=${{ matrix.goos }} + export GOARCH=${{ matrix.goarch }} + go build -o "release/devops-mcp-server-${GOOS}-${GOARCH}" + + - name: Create release assets + id: archive + run: | + archive="dist/devops-mcp-server-${{ matrix.goos }}-${{ matrix.goarch }}" + if [[ "matrix.goos" == "windows" ]];then + archive="${archive}.zip" + zip -J "${archive}" -C release . + else + archive="${archive}.tar.gz" + tar -czvf "${archive}" -C release . + fi + echo "Created ${archive}" + echo "ARCHIVE_NAME=${archive}" >> "$GITHUB_OUTPUT" + + - name: Create release assets + id: test + run: | + pwd + ls dist/* + ls ${{ steps.archive.ARCHIVE_NAME }} + + - name: UpUpload archive to GitHub Release + uses: softprops/action-gh-release@v1 + with: + files: dist/* + overwrite_files: true + tag_name: ${{ env.RELEASE }} + fail_on_unmatched_files: true + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/prep-release.yaml b/.github/workflows/prep-release.yaml new file mode 100644 index 0000000..febaad7 --- /dev/null +++ b/.github/workflows/prep-release.yaml @@ -0,0 +1,54 @@ +name: Prepare Release + +on: + workflow_call: + inputs: + release: + required: true + type: string + default: dev + workflow_dispatch: # Allows manual triggering of the workflow + +env: + RELEASE: dev + GH_TOKEN: ${{ github.token }} + +jobs: + prep-release: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v3 + + - name: Prepare for a new release + id: prep + run: | + export RELEASE="${{ env.RELEASE }}" + if [[ "${{ inputs.release }}" != "" ]]; then + echo "Updating to RELEASE to ${{ inputs.release }}" + RELEASE="${{ inputs.release }}" + fi + echo "RELEASE=${RELEASE}" >> "$GITHUB_OUTPUT" + + echo "Preparing for \"${RELEASE}\" release" + git config --local user.email "action@github.com" + git config --local user.name "GitHub Action" + git tag -fa -m "${RELEASE} release" "${RELEASE}" + git push --force origin "${RELEASE}" + + # check of a release already exists + if [[ $(gh release view "${RELEASE}" >/dev/null ) ]]; then + echo "CREATE_RELEASE=true" >> "$GITHUB_OUTPUT" + else + echo "CREATE_RELEASE=false" >> "$GITHUB_OUTPUT" + fi + + - name: Create a GitHub Release + uses: softprops/action-gh-release@v1 + if: ${{ steps.prep.CREATE_RELEASE }} + with: + name: ${{ steps.prep.RELEASE }} + tag: ${{ steps.prep.RELEASE }} + + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.gitignore b/.gitignore index bcb8d0e..0f03958 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .venv/ +dist/ .gemini **/__pycache__/ **/*.pyc diff --git a/devops-mcp-server/auth/auth.go b/devops-mcp-server/auth/auth.go new file mode 100644 index 0000000..37e7b34 --- /dev/null +++ b/devops-mcp-server/auth/auth.go @@ -0,0 +1,65 @@ +// 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 auth + +import ( + "context" + "fmt" + + "cloud.google.com/go/auth" + "cloud.google.com/go/auth/credentials" +) + +type Creds struct { + Token string + ProjectId string +} + +func GetAuthToken(ctx context.Context) (Creds, error) { + // Use Application Default Credentials to get a TokenSource + scopes := []string{"https://www.googleapis.com/auth/cloud-platform"} + creds, err := credentials.DetectDefault(&credentials.DetectOptions{ + Scopes: scopes, + }) + if err != nil { + return Creds{}, fmt.Errorf("Failed to find default credentials: %w", err) + } + + projectID, err := creds.ProjectID(ctx) + if err != nil { + return Creds{}, fmt.Errorf("Failed to get project ID: %w", err) + } + if projectID == "" { + //Try quota project + projectID, err = creds.QuotaProjectID(ctx) + if err != nil { + return Creds{}, fmt.Errorf("Failed to get project ID: %w", err) + } + if projectID == "" { + return Creds{}, fmt.Errorf(` + No Project ID found in Application Default Credentials. + This can happen if credentials are user-based or the project hasn't been explicitly set + e.g., via gcloud auth application-default set-quota-project. + Error:%v`, err) + } + } + // We need an access token + var token *auth.Token + token, err = creds.TokenProvider.Token(ctx) + if err != nil { + return Creds{}, fmt.Errorf("Failed to retrieve access token: %v", err) + } + + return Creds{Token: token.Value, ProjectId: projectID}, nil +} diff --git a/devops-mcp-server/cloudbuild/cloudbuild.go b/devops-mcp-server/cloudbuild/cloudbuild.go index dc18a8e..000b326 100644 --- a/devops-mcp-server/cloudbuild/cloudbuild.go +++ b/devops-mcp-server/cloudbuild/cloudbuild.go @@ -36,6 +36,9 @@ type ProjectsLocationsTriggersServiceWrapper struct { *cloudbuild.ProjectsLocationsTriggersService } +type ListResult[T any] struct { + Items []T `json:"items"` +} type triggersCreateCallWrapper struct { *cloudbuild.ProjectsLocationsTriggersCreateCall } @@ -90,7 +93,6 @@ func (w *ProjectsLocationsTriggersServiceWrapper) List(parent string) cloudbuild return &triggersListCallWrapper{w.ProjectsLocationsTriggersService.List(parent)} } - // ProjectsLocationsBuildsServiceWrapper wraps cloudbuild.ProjectsLocationsBuildsService type ProjectsLocationsBuildsServiceWrapper struct { *cloudbuild.ProjectsLocationsBuildsService @@ -219,13 +221,13 @@ func (c *Client) RunTrigger(ctx context.Context, projectID, location, triggerID } // ListTriggers lists all Cloud Build triggers in a given location. -func (c *Client) ListTriggers(ctx context.Context, projectID, location string) ([]*cloudbuild.BuildTrigger, error) { +func (c *Client) ListTriggers(ctx context.Context, projectID, location string) (*ListResult[*cloudbuild.BuildTrigger], error) { parent := fmt.Sprintf("projects/%s/locations/%s", projectID, location) resp, err := c.triggersService.List(parent).Context(ctx).Do() if err != nil { return nil, fmt.Errorf("failed to list triggers: %v", err) } - return resp.Triggers, nil + return &ListResult[*cloudbuild.BuildTrigger]{Items: resp.Triggers}, nil } // BuildContainer builds a container image using Cloud Build. diff --git a/devops-mcp-server/cloudbuild/cloudbuild_test.go b/devops-mcp-server/cloudbuild/cloudbuild_test.go index eadf03e..85d6061 100644 --- a/devops-mcp-server/cloudbuild/cloudbuild_test.go +++ b/devops-mcp-server/cloudbuild/cloudbuild_test.go @@ -205,8 +205,8 @@ func TestListTriggers(t *testing.T) { if err != nil { t.Fatalf("ListTriggers() error = %v, want nil", err) } - if len(triggers) != 2 { - t.Errorf("ListTriggers() got %d triggers, want 2", len(triggers)) + if len(triggers.Items) != 2 { + t.Errorf("ListTriggers() got %d triggers, want 2", len(triggers.Items)) } }) diff --git a/devops-mcp-server/clouddeploy/clouddeploy.go b/devops-mcp-server/clouddeploy/clouddeploy.go index 9eb4120..15d84b6 100644 --- a/devops-mcp-server/clouddeploy/clouddeploy.go +++ b/devops-mcp-server/clouddeploy/clouddeploy.go @@ -23,163 +23,334 @@ import ( deploypb "cloud.google.com/go/deploy/apiv1/deploypb" ) + + +// ListResult defines a generic struct to wrap a list of items. + +type ListResult[T any] struct { + + Items []T `json:"items"` + +} + + + // Client is a client for interacting with the Cloud Deploy API. + type Client struct { + client *deploy.CloudDeployClient + } + + // NewClient creates a new Client. + func NewClient(ctx context.Context) (*Client, error) { + c, err := deploy.NewCloudDeployClient(ctx) + if err != nil { + return nil, fmt.Errorf("failed to create cloud deploy client: %v", err) + } + return &Client{client: c}, nil + } + + // CreateDeliveryPipeline creates a new Cloud Deploy delivery pipeline. + func (c *Client) CreateDeliveryPipeline(ctx context.Context, projectID, location, pipelineID, description string) (*deploypb.DeliveryPipeline, error) { + req := &deploypb.CreateDeliveryPipelineRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location), + DeliveryPipelineId: pipelineID, + DeliveryPipeline: &deploypb.DeliveryPipeline{ + Description: description, + Pipeline: &deploypb.DeliveryPipeline_SerialPipeline{ + SerialPipeline: &deploypb.SerialPipeline{}, + }, + }, + } + op, err := c.client.CreateDeliveryPipeline(ctx, req) + if err != nil { + return nil, fmt.Errorf("failed to create delivery pipeline: %v", err) + } + return op.Wait(ctx) + } + + // CreateGKETarget creates a new Cloud Deploy GKE target. + func (c *Client) CreateGKETarget(ctx context.Context, projectID, location, targetID, gkeCluster, description string) (*deploypb.Target, error) { + req := &deploypb.CreateTargetRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location), + TargetId: targetID, + Target: &deploypb.Target{ + Description: description, + }, + } + op, err := c.client.CreateTarget(ctx, req) + if err != nil { + return nil, fmt.Errorf("failed to create gke target: %v", err) + } + return op.Wait(ctx) + } + + // CreateCloudRunTarget creates a new Cloud Deploy Cloud Run target. + func (c *Client) CreateCloudRunTarget(ctx context.Context, projectID, location, targetID, description string) (*deploypb.Target, error) { + req := &deploypb.CreateTargetRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location), + TargetId: targetID, + Target: &deploypb.Target{ + Description: description, + }, + } + op, err := c.client.CreateTarget(ctx, req) + if err != nil { + return nil, fmt.Errorf("failed to create cloud run target: %v", err) + } + return op.Wait(ctx) + } + + // CreateRollout creates a new Cloud Deploy rollout. + func (c *Client) CreateRollout(ctx context.Context, projectID, location, pipelineID, releaseID, rolloutID, targetID string) (*deploypb.Rollout, error) { + req := &deploypb.CreateRolloutRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s/deliveryPipelines/%s/releases/%s", projectID, location, pipelineID, releaseID), + RolloutId: rolloutID, + Rollout: &deploypb.Rollout{ + TargetId: targetID, + }, + } + op, err := c.client.CreateRollout(ctx, req) + if err != nil { + return nil, fmt.Errorf("failed to create rollout: %v", err) + } + return op.Wait(ctx) + } + + // ListDeliveryPipelines lists all Cloud Deploy delivery pipelines. -func (c *Client) ListDeliveryPipelines(ctx context.Context, projectID, location string) ([]*deploypb.DeliveryPipeline, error) { + +func (c *Client) ListDeliveryPipelines(ctx context.Context, projectID, location string) (*ListResult[*deploypb.DeliveryPipeline], error) { + req := &deploypb.ListDeliveryPipelinesRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location), + } + it := c.client.ListDeliveryPipelines(ctx, req) + var pipelines []*deploypb.DeliveryPipeline + for { + pipeline, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("failed to list delivery pipelines: %v", err) + } + pipelines = append(pipelines, pipeline) + } - return pipelines, nil + + return &ListResult[*deploypb.DeliveryPipeline]{Items: pipelines}, nil + } + + // ListTargets lists all Cloud Deploy targets. -func (c *Client) ListTargets(ctx context.Context, projectID, location string) ([]*deploypb.Target, error) { + +func (c *Client) ListTargets(ctx context.Context, projectID, location string) (*ListResult[*deploypb.Target], error) { + req := &deploypb.ListTargetsRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, location), + } + it := c.client.ListTargets(ctx, req) + var targets []*deploypb.Target + for { + target, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("failed to list targets: %v", err) + } + targets = append(targets, target) + } - return targets, nil + + return &ListResult[*deploypb.Target]{Items: targets}, nil + } + + // ListReleases lists all Cloud Deploy releases for a given delivery pipeline. -func (c *Client) ListReleases(ctx context.Context, projectID, location, pipelineID string) ([]*deploypb.Release, error) { + +func (c *Client) ListReleases(ctx context.Context, projectID, location, pipelineID string) (*ListResult[*deploypb.Release], error) { + req := &deploypb.ListReleasesRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s/deliveryPipelines/%s", projectID, location, pipelineID), + } + it := c.client.ListReleases(ctx, req) + var releases []*deploypb.Release + for { + release, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("failed to list releases: %v", err) + } + releases = append(releases, release) + } - return releases, nil + + return &ListResult[*deploypb.Release]{Items: releases}, nil + } + + // ListRollouts lists all Cloud Deploy rollouts for a given release. -func (c *Client) ListRollouts(ctx context.Context, projectID, location, pipelineID, releaseID string) ([]*deploypb.Rollout, error) { + +func (c *Client) ListRollouts(ctx context.Context, projectID, location, pipelineID, releaseID string) (*ListResult[*deploypb.Rollout], error) { + req := &deploypb.ListRolloutsRequest{ + Parent: fmt.Sprintf("projects/%s/locations/%s/deliveryPipelines/%s/releases/%s", projectID, location, pipelineID, releaseID), + } + it := c.client.ListRollouts(ctx, req) + var rollouts []*deploypb.Rollout + for { + rollout, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("failed to list rollouts: %v", err) + } + rollouts = append(rollouts, rollout) + } - return rollouts, nil -} \ No newline at end of file + + return &ListResult[*deploypb.Rollout]{Items: rollouts}, nil + +} diff --git a/devops-mcp-server/containeranalysis/containeranalysis.go b/devops-mcp-server/containeranalysis/containeranalysis.go index 76e1eb2..41f7cc2 100644 --- a/devops-mcp-server/containeranalysis/containeranalysis.go +++ b/devops-mcp-server/containeranalysis/containeranalysis.go @@ -23,37 +23,82 @@ import ( grafeaspb "google.golang.org/genproto/googleapis/grafeas/v1" ) + + +// ListResult defines a generic struct to wrap a list of items. + +type ListResult[T any] struct { + + Items []T `json:"items"` + +} + + + // Client is a client for interacting with the Container Analysis API. + type Client struct { + client *containeranalysis.Client + } + + // NewClient creates a new Client. + func NewClient(ctx context.Context) (*Client, error) { + c, err := containeranalysis.NewClient(ctx) + if err != nil { + return nil, fmt.Errorf("failed to create container analysis client: %v", err) + } + return &Client{client: c}, nil + } + + // ListVulnerabilities lists vulnerabilities for a given image resource URL. -func (c *Client) ListVulnerabilities(ctx context.Context, projectID, resourceURL string) ([]*grafeaspb.Occurrence, error) { + +func (c *Client) ListVulnerabilities(ctx context.Context, projectID, resourceURL string) (*ListResult[*grafeaspb.Occurrence], error) { + req := &grafeaspb.ListOccurrencesRequest{ + Parent: fmt.Sprintf("projects/%s", projectID), + Filter: fmt.Sprintf("resourceUrl=\"%s\" AND kind=\"VULNERABILITY\"", resourceURL), + } + it := c.client.GetGrafeasClient().ListOccurrences(ctx, req) + var vulnerabilities []*grafeaspb.Occurrence + for { + occurrence, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return nil, fmt.Errorf("failed to list vulnerabilities: %v", err) + } + vulnerabilities = append(vulnerabilities, occurrence) + } - return vulnerabilities, nil -} \ No newline at end of file + + return &ListResult[*grafeaspb.Occurrence]{Items: vulnerabilities}, nil + +} diff --git a/devops-mcp-server/devconnect/devconnect.go b/devops-mcp-server/devconnect/devconnect.go index ae8ce1a..1a5d5e6 100644 --- a/devops-mcp-server/devconnect/devconnect.go +++ b/devops-mcp-server/devconnect/devconnect.go @@ -22,37 +22,63 @@ import ( devconnectv1 "google.golang.org/api/developerconnect/v1" ) +// ListResult defines a generic struct to wrap a list of items. + +type ListResult[T any] struct { + Items []T `json:"items"` +} + // Client is a client for interacting with the Developer Connect API. + type Client struct { service *devconnectv1.Service } // NewClient creates a new Client. + func NewClient(ctx context.Context) (*Client, error) { service, err := devconnectv1.NewService(ctx) if err != nil { + return nil, fmt.Errorf("failed to create developer connect service: %v", err) + } + return &Client{service: service}, nil + } func (c *Client) waitForOperation(ctx context.Context, operation *devconnectv1.Operation) (*devconnectv1.Operation, error) { ctx, cancel := context.WithTimeout(ctx, 5*time.Minute) + defer cancel() for !operation.Done { + select { + case <-ctx.Done(): + return nil, fmt.Errorf("timed out waiting for operation: %v", ctx.Err()) + case <-time.After(1 * time.Second): + op, err := c.service.Projects.Locations.Operations.Get(operation.Name).Do() + if err != nil { + return nil, fmt.Errorf("failed to get operation: %v", err) + } + operation = op + } + } + return operation, nil + } // CreateConnection creates a new Developer Connect connection. @@ -63,20 +89,33 @@ func (c *Client) CreateConnection(ctx context.Context, projectID, location, conn GithubApp: "DEVELOPER_CONNECT", }, } + op, err := c.service.Projects.Locations.Connections.Create(parent, req).ConnectionId(connectionID).Do() + if err != nil { + return nil, fmt.Errorf("failed to create connection: %v", err) + } + op, err = c.waitForOperation(ctx, op) + if err != nil { + return nil, err + } + if op.Error != nil { + return nil, fmt.Errorf("operation failed: %v", op.Error) + } name := fmt.Sprintf("projects/%s/locations/%s/connections/%s", projectID, location, connectionID) + return c.service.Projects.Locations.Connections.Get(name).Do() + } // CreateGitRepositoryLink creates a new Developer Connect Git Repository Link. @@ -85,44 +124,73 @@ func (c *Client) CreateGitRepositoryLink(ctx context.Context, projectID, locatio req := &devconnectv1.GitRepositoryLink{ CloneUri: repoURI, } + op, err := c.service.Projects.Locations.Connections.GitRepositoryLinks.Create(parent, req).GitRepositoryLinkId(repoLinkID).Do() + if err != nil { + return nil, fmt.Errorf("failed to create git repository link: %v", err) + } + op, err = c.waitForOperation(ctx, op) + if err != nil { + return nil, err + } + if op.Error != nil { + return nil, fmt.Errorf("operation failed: %v", op.Error) + } name := fmt.Sprintf("%s/gitRepositoryLinks/%s", parent, repoLinkID) + return c.service.Projects.Locations.Connections.GitRepositoryLinks.Get(name).Do() + } // ListConnections lists Developer Connect connections. -func (c *Client) ListConnections(ctx context.Context, projectID, location string) ([]*devconnectv1.Connection, error) { +func (c *Client) ListConnections(ctx context.Context, projectID, location string) (*ListResult[*devconnectv1.Connection], error) { parent := fmt.Sprintf("projects/%s/locations/%s", projectID, location) + resp, err := c.service.Projects.Locations.Connections.List(parent).Do() + if err != nil { + return nil, fmt.Errorf("failed to list connections: %v", err) + } - return resp.Connections, nil + + return &ListResult[*devconnectv1.Connection]{Items: resp.Connections}, nil + } // GetConnection gets a Developer Connect connection. func (c *Client) GetConnection(ctx context.Context, projectID, location, connectionID string) (*devconnectv1.Connection, error) { name := fmt.Sprintf("projects/%s/locations/%s/connections/%s", projectID, location, connectionID) + return c.service.Projects.Locations.Connections.Get(name).Do() + } // FindGitRepositoryLinksForGitRepo finds already configured Developer Connect Git Repository Links for a particular git repository. -func (c *Client) FindGitRepositoryLinksForGitRepo(ctx context.Context, projectID, location, repoURI string) ([]*devconnectv1.GitRepositoryLink, error) { + +func (c *Client) FindGitRepositoryLinksForGitRepo(ctx context.Context, projectID, location, repoURI string) (*ListResult[*devconnectv1.GitRepositoryLink], error) { + parent := fmt.Sprintf("projects/%s/locations/%s/connections/-", projectID, location) + resp, err := c.service.Projects.Locations.Connections.GitRepositoryLinks.List(parent).Filter(fmt.Sprintf("clone_uri=\"%s\"", repoURI)).Do() + if err != nil { + return nil, fmt.Errorf("failed to list git repository links: %v", err) + } - return resp.GitRepositoryLinks, nil + + return &ListResult[*devconnectv1.GitRepositoryLink]{Items: resp.GitRepositoryLinks}, nil + } diff --git a/devops-mcp-server/go.mod b/devops-mcp-server/go.mod index 01692c8..4bf7659 100644 --- a/devops-mcp-server/go.mod +++ b/devops-mcp-server/go.mod @@ -3,6 +3,8 @@ module devops-mcp-server go 1.24.8 require ( + cloud.google.com/go/auth v0.17.0 + github.com/philippgille/chromem-go v0.7.0 cloud.google.com/go/artifactregistry v1.17.2 cloud.google.com/go/containeranalysis v0.14.2 cloud.google.com/go/deploy v1.27.3 @@ -18,7 +20,6 @@ require ( require ( cel.dev/expr v0.24.0 // indirect cloud.google.com/go v0.123.0 // indirect - cloud.google.com/go/auth v0.17.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/grafeas v0.3.17 // indirect diff --git a/devops-mcp-server/go.sum b/devops-mcp-server/go.sum index c6587de..a58339b 100644 --- a/devops-mcp-server/go.sum +++ b/devops-mcp-server/go.sum @@ -81,6 +81,8 @@ github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81 github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc= github.com/modelcontextprotocol/go-sdk v1.0.0 h1:Z4MSjLi38bTgLrd/LjSmofqRqyBiVKRyQSJgw8q8V74= github.com/modelcontextprotocol/go-sdk v1.0.0/go.mod h1:nYtYQroQ2KQiM0/SbyEPUWQ6xs4B95gJjEalc9AQyOs= +github.com/philippgille/chromem-go v0.7.0 h1:4jfvfyKymjKNfGxBUhHUcj1kp7B17NL/I1P+vGh1RvY= +github.com/philippgille/chromem-go v0.7.0/go.mod h1:hTd+wGEm/fFPQl7ilfCwQXkgEUxceYh86iIdoKMolPo= github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= diff --git a/devops-mcp-server/iam/iam.go b/devops-mcp-server/iam/iam.go index d8f6d87..03f842e 100644 --- a/devops-mcp-server/iam/iam.go +++ b/devops-mcp-server/iam/iam.go @@ -22,23 +22,37 @@ import ( iamv1 "google.golang.org/api/iam/v1" ) +// ListResult defines a generic struct to wrap a list of items. + +type ListResult[T any] struct { + Items []T `json:"items"` +} + // Client is a client for interacting with the IAM API. + type Client struct { iamService *iamv1.Service crmService *cloudresourcemanagerv1.Service } // NewClient creates a new Client. + func NewClient(ctx context.Context) (*Client, error) { iamService, err := iamv1.NewService(ctx) if err != nil { + return nil, fmt.Errorf("failed to create iam service: %v", err) + } crmService, err := cloudresourcemanagerv1.NewService(ctx) if err != nil { + return nil, fmt.Errorf("failed to create cloud resource manager service: %v", err) + } + return &Client{iamService: iamService, crmService: crmService}, nil + } // CreateServiceAccount creates a new Google Cloud Platform service account. @@ -50,51 +64,75 @@ func (c *Client) CreateServiceAccount(ctx context.Context, projectID, displayNam DisplayName: displayName, }, } + return c.iamService.Projects.ServiceAccounts.Create(projectPath, req).Context(ctx).Do() + } // AddIAMRoleBinding adds an IAM role binding to a Google Cloud Platform resource. func (c *Client) AddIAMRoleBinding(ctx context.Context, resourceID, role, member string) (*cloudresourcemanagerv1.Policy, error) { policy, err := c.crmService.Projects.GetIamPolicy(resourceID, &cloudresourcemanagerv1.GetIamPolicyRequest{}).Context(ctx).Do() if err != nil { + return nil, fmt.Errorf("failed to get iam policy: %v", err) + } policy.Bindings = append(policy.Bindings, &cloudresourcemanagerv1.Binding{ - Role: role, + Role: role, + Members: []string{member}, }) setPolicyRequest := &cloudresourcemanagerv1.SetIamPolicyRequest{ Policy: policy, } + return c.crmService.Projects.SetIamPolicy(resourceID, setPolicyRequest).Context(ctx).Do() + } // ListServiceAccounts lists all service accounts in a project. -func (c *Client) ListServiceAccounts(ctx context.Context, projectID string) ([]*iamv1.ServiceAccount, error) { +func (c *Client) ListServiceAccounts(ctx context.Context, projectID string) (*ListResult[*iamv1.ServiceAccount], error) { parent := fmt.Sprintf("projects/%s", projectID) + resp, err := c.iamService.Projects.ServiceAccounts.List(parent).Context(ctx).Do() + if err != nil { + return nil, fmt.Errorf("failed to list service accounts: %v", err) + } - return resp.Accounts, nil + + return &ListResult[*iamv1.ServiceAccount]{Items: resp.Accounts}, nil + } // GetIAMRoleBinding gets the IAM role bindings for a service account. -func (c *Client) GetIAMRoleBinding(ctx context.Context, projectID, serviceAccountEmail string) ([]string, error) { +func (c *Client) GetIAMRoleBinding(ctx context.Context, projectID, serviceAccountEmail string) (*ListResult[string], error) { policy, err := c.crmService.Projects.GetIamPolicy(projectID, &cloudresourcemanagerv1.GetIamPolicyRequest{}).Context(ctx).Do() if err != nil { + return nil, fmt.Errorf("failed to get iam policy: %v", err) + } var roles []string + for _, binding := range policy.Bindings { + for _, member := range binding.Members { + if member == fmt.Sprintf("serviceAccount:%s", serviceAccountEmail) { + roles = append(roles, binding.Role) + } + } + } - return roles, nil + + return &ListResult[string]{Items: roles}, nil + } diff --git a/devops-mcp-server/main.go b/devops-mcp-server/main.go index f51e314..632ee75 100644 --- a/devops-mcp-server/main.go +++ b/devops-mcp-server/main.go @@ -32,7 +32,6 @@ var ( pprofAddr = flag.String("pprof", "", "if set, host the pprof debugging server at this address") ) - func main() { flag.Parse() @@ -73,4 +72,4 @@ func main() { log.Printf("Server failed: %v", err) } } -} \ No newline at end of file +} diff --git a/devops-mcp-server/rag/rag.go b/devops-mcp-server/rag/rag.go new file mode 100644 index 0000000..c608d17 --- /dev/null +++ b/devops-mcp-server/rag/rag.go @@ -0,0 +1,109 @@ +// 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 rag + +import ( + "context" + "devops-mcp-server/auth" + "fmt" + "log" + "os" + + chromem "github.com/philippgille/chromem-go" +) + +// ListResult defines a generic struct to wrap a list of items. +type ListResult[T any] struct { + Items []T `json:"items"` +} + +type RagData struct { + DB *chromem.DB + Pattern *chromem.Collection + Knowledge *chromem.Collection +} + +var RagDB RagData + +func init() { + + ctx := context.Background() + + dbFile := os.Getenv("RAG_DB_PATH") + RagDB = RagData{DB: chromem.NewDB()} + if len(dbFile) < 1 { + log.Fatalf("Env variable RAG_DB_PATH is not set for RAG data file %v", os.Getenv("RAG_DB_PATH")) + } + //check if file exists, we expect an existing DB + if _, err := os.Stat(dbFile); os.IsNotExist(err) { + log.Fatalf("RAG_DB_PATH file does not exist, skipping import: %v", dbFile) + } else { + err := RagDB.DB.ImportFromFile(dbFile, "") + if err != nil { + log.Printf("Unable to import from the RAG DB file:%s - %v", dbFile, err) + } + log.Printf("IMPORTED from the RAG DB file:%s - %v", dbFile, len(RagDB.DB.ListCollections())) + } + + creds, err := auth.GetAuthToken(ctx) + if err != nil { + log.Printf("Error: Google Cloud account is required: %v", err) + //TODO: Should we fail MCP server startup if Google Cloud account is not setup? + //For now if Google Cloud account is not setup, we will silently fail + return + } + + vertexEmbeddingFunc := chromem.NewEmbeddingFuncVertex( + creds.Token, + creds.ProjectId, + chromem.EmbeddingModelVertexEnglishV4) + + RagDB.Pattern, err = RagDB.DB.GetOrCreateCollection("pattern", nil, vertexEmbeddingFunc) + if err != nil { + log.Fatalf("Unable to get collection pattern: %v", err) + } else { + log.Printf("LOADED collection pattern: %v", RagDB.Pattern.Count()) + } + + RagDB.Knowledge, err = RagDB.DB.GetOrCreateCollection("knowledge", nil, vertexEmbeddingFunc) + if err != nil { + log.Fatalf("Unable to get collection knowledge: %v", err) + } else { + log.Printf("LOADED collection knowledge: %v", RagDB.Knowledge.Count()) + } + log.Print("Init Completed!") +} + +func QueryKnowledge(ctx context.Context, query string) (*ListResult[chromem.Result], error) { + if RagDB.Knowledge == nil { + return nil, fmt.Errorf("RagDB does not contain Knowledge collection, among %d collections", len(RagDB.DB.ListCollections())) + } + result, err := RagDB.Knowledge.Query(ctx, query, 3, nil, nil) + if err != nil { + log.Fatalf("Unable to Query collection pattern: %v", err) + } + return &ListResult[chromem.Result]{Items: result}, nil +} + +func QueryPattern(ctx context.Context, query string) (*ListResult[chromem.Result], error) { + if RagDB.Pattern == nil { + return nil, fmt.Errorf("RagDB does not contain Pattern collection, among %d collections", len(RagDB.DB.ListCollections())) + } + result, err := RagDB.Pattern.Query(ctx, query, 2, nil, nil) + if err != nil { + log.Fatalf("Unable to Query collection pattern: %v", err) + } + return &ListResult[chromem.Result]{Items: result}, nil +} diff --git a/devops-mcp-server/rag/rag_test.go b/devops-mcp-server/rag/rag_test.go new file mode 100644 index 0000000..dce4e59 --- /dev/null +++ b/devops-mcp-server/rag/rag_test.go @@ -0,0 +1,13 @@ +package rag + +import ( + "os" + "testing" +) + +func TestRAGInit(t *testing.T) { + + if RagDB.DB == nil { + t.Fatalf("init() failed to load: %v", os.Getenv("RAG_DB_PATH")) + } +} diff --git a/devops-mcp-server/server.go b/devops-mcp-server/server.go index 7898f3a..bdeab87 100644 --- a/devops-mcp-server/server.go +++ b/devops-mcp-server/server.go @@ -37,16 +37,16 @@ import ( //go:embed version.txt var version string -func createServer() *mcp.Server{ +func createServer() *mcp.Server { opts := &mcp.ServerOptions{ - Instructions: "Google Cloud DevOps MCP Server", + Instructions: "Google Cloud DevOps MCP Server", HasResources: false, } server := mcp.NewServer(&mcp.Implementation{ - Name: "devops", - Title: "Google Cloud DevOps MCP Server", - Version: version, - }, opts) + Name: "devops", + Title: "Google Cloud DevOps MCP Server", + Version: version, + }, opts) ctx := context.Background() diff --git a/devops-mcp-server/server_test.go b/devops-mcp-server/server_test.go index be9fddb..d53c676 100644 --- a/devops-mcp-server/server_test.go +++ b/devops-mcp-server/server_test.go @@ -49,7 +49,7 @@ func TestMCPServer(t *testing.T) { log.Fatalf("failed to list tools: %v", err) } log.Printf("Prompts: %v", len(tools.Tools)) - for _,tool := range tools.Tools { + for _, tool := range tools.Tools { log.Printf("Tool %s: %s", tool.Name, tool.Title) } @@ -58,10 +58,9 @@ func TestMCPServer(t *testing.T) { log.Fatalf("failed to list prompts: %v", err) } log.Printf("Prompts: %v", len(prompts.Prompts)) - for _,prompt := range prompts.Prompts { + for _, prompt := range prompts.Prompts { log.Printf("Prompt %s: %s", prompt.Name, prompt.Title) } - // Now shut down the session by closing the client, and waiting for the // server session to end. diff --git a/devops-mcp-server/testing/cloudbuild_test.go b/devops-mcp-server/testing/cloudbuild_test.go new file mode 100644 index 0000000..db52cb0 --- /dev/null +++ b/devops-mcp-server/testing/cloudbuild_test.go @@ -0,0 +1,16 @@ +// 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. + +// Cloud Build end-to-end tests +package testing diff --git a/devops-mcp-server/testing/rag_test.go b/devops-mcp-server/testing/rag_test.go new file mode 100644 index 0000000..bfeca0d --- /dev/null +++ b/devops-mcp-server/testing/rag_test.go @@ -0,0 +1,49 @@ +// 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. + +// Cloud Build end-to-end tests +package testing + +import ( + "context" + "devops-mcp-server/auth" + "devops-mcp-server/rag" + "log" + "testing" +) + +func TestRAGQuery(t *testing.T) { + ctx := context.Background() + creds, err := auth.GetAuthToken(ctx) + if creds.Token == "" || creds.ProjectId == "" || err != nil { + t.Skipf("Skipping test! Google Cloud account not found %v : %v", creds, err) + } + + patternResult, err := rag.QueryKnowledge(ctx, "Simple pipeline that deploys to Cloud Run") + if err != nil { + log.Fatalf("Unable to Query collection pattern: %v", err) + } + if (patternResult == nil) || (len(patternResult.Items) < 2) || (patternResult.Items[0].Content == "") { + log.Fatalf("Failed to find knowledge: %v", patternResult) + } + + knowledgeResult, err := rag.QueryKnowledge(ctx, "Package a Python application") + if err != nil { + log.Fatalf("Unable to Query collection knowledge: %v", err) + } + if (knowledgeResult == nil) || (len(knowledgeResult.Items) < 3) || (knowledgeResult.Items[0].Content == "") { + log.Fatalf("Failed to find knowledge: %v", knowledgeResult) + } + +} diff --git a/docs/contributing.md b/docs/contributing.md deleted file mode 100644 index 8956a61..0000000 --- a/docs/contributing.md +++ /dev/null @@ -1,32 +0,0 @@ -# How to Contribute - -We would love to accept your patches and contributions to this project. - -## Before you begin - -### Sign our Contributor License Agreement - -Contributions to this project must be accompanied by a -[Contributor License Agreement](https://cla.developers.google.com/about) (CLA). -You (or your employer) retain the copyright to your contribution; this simply -gives us permission to use and redistribute your contributions as part of the -project. - -If you or your current employer have already signed the Google CLA (even if it -was for a different project), you probably don't need to do it again. - -Visit to see your current agreements or to -sign a new one. - -### Review our Community Guidelines - -This project follows [Google's Open Source Community -Guidelines](https://opensource.google/conduct/). - -## Contribution process - -### Code Reviews - -All submissions, including submissions by project members, require review. We -use [GitHub pull requests](https://docs.github.com/articles/about-pull-requests) -for this purpose. diff --git a/local-rag/fetch_docs.go b/local-rag/fetch_docs.go index 9fe327c..4457758 100644 --- a/local-rag/fetch_docs.go +++ b/local-rag/fetch_docs.go @@ -318,4 +318,4 @@ func downloadWebsites(sources *Source, extractToDir string) error { }) } return nil -} \ No newline at end of file +} diff --git a/local-rag/fetch_docs_test.go b/local-rag/fetch_docs_test.go index f4fcbd3..d027d5f 100644 --- a/local-rag/fetch_docs_test.go +++ b/local-rag/fetch_docs_test.go @@ -227,4 +227,4 @@ func TestFetchRepository_Zip(t *testing.T) { if string(content) != "dummy content" { t.Errorf("Expected 'dummy content', got '%s'", string(content)) } -} \ No newline at end of file +} diff --git a/local-rag/main.go b/local-rag/main.go index c570641..dd51173 100644 --- a/local-rag/main.go +++ b/local-rag/main.go @@ -160,7 +160,7 @@ func main() { log.Fatalf("Failed to get project ID: %v", err) } if projectID == "" { - //Try quota project + //Try quota project projectID, err = creds.QuotaProjectID(ctx) if err != nil { log.Fatalf("Failed to get project ID: %v", err) @@ -190,9 +190,9 @@ func main() { //check if file exists, only import if it does if _, err := os.Stat(dbFile); os.IsNotExist(err) { log.Printf("RAG_DB_PATH file does not exist, skipping import: %v", dbFile) - }else{ + } else { err := db.ImportFromFile(dbFile, "") - log.Printf("Imported RAG with collections:%d",len(db.ListCollections())) + log.Printf("Imported RAG with collections:%d", len(db.ListCollections())) if err != nil { log.Fatalf("Unable to import from the RAG DB file:%s - %v", dbFile, err) } @@ -214,16 +214,16 @@ func main() { } patternsDir := filepath.Join(pwd, "patterns") - addDirectoryToRag(ctx,collectionPattern, patternsDir) + addDirectoryToRag(ctx, collectionPattern, patternsDir) knowledgeDir := filepath.Join(pwd, "knowledge") - addDirectoryToRag(ctx,collectionKnowledge, knowledgeDir) + addDirectoryToRag(ctx, collectionKnowledge, knowledgeDir) // Create a temporary directory for downloads //tmpDir, err := os.MkdirTemp("", "rag-data-") - ragSourceDir,err := os.Getwd() + ragSourceDir, err := os.Getwd() if err != nil { - log.Fatalf("Unable to get working directory: %v",err) + log.Fatalf("Unable to get working directory: %v", err) } ragSourceDir = ragSourceDir + "/.rag-sources" //Create dir if it does not exist @@ -237,30 +237,30 @@ func main() { //defer os.RemoveAll(tmpDir) // Process data sources if destination is empty - // otherwise we assume last run was successful in + // otherwise we assume last run was successful in // fetching sources entries, err := os.ReadDir(ragSourceDir) if err != nil { - log.Fatalf("Unable to read directory: %v",err) + log.Fatalf("Unable to read directory: %v", err) } - if (len(entries)==0){ + if len(entries) == 0 { for _, source := range KNOWLEDGE_RAG_SOURCES { processSource(source, ragSourceDir) } } // Upload all files in the temporary directory to RAG - addDirectoryToRag(ctx,collectionKnowledge, ragSourceDir) + addDirectoryToRag(ctx, collectionKnowledge, ragSourceDir) // Export the database to a file if len(dbFile) > 0 { - log.Printf("Exporting database Knowledge base docs:%d, Pattern docs:%d", - collectionKnowledge.Count(), - collectionPattern.Count()) + log.Printf("Exporting database Knowledge base docs:%d, Pattern docs:%d", + collectionKnowledge.Count(), + collectionPattern.Count()) err = db.ExportToFile(dbFile, true, "") if err != nil { log.Fatal(err) } log.Printf("Database exported to %s", dbFile) } -} \ No newline at end of file +} diff --git a/local-rag/query_test.go b/local-rag/query_test.go index b75ea48..acd4d97 100644 --- a/local-rag/query_test.go +++ b/local-rag/query_test.go @@ -25,7 +25,7 @@ import ( chromem "github.com/philippgille/chromem-go" ) -func gcpAuthHelper(ctx context.Context,t *testing.T) (tokenValue, projectID string){ +func gcpAuthHelper(ctx context.Context, t *testing.T) (tokenValue, projectID string) { // Use Application Default Credentials to get a TokenSource scopes := []string{"https://www.googleapis.com/auth/cloud-platform"} creds, err := credentials.DetectDefault(&credentials.DetectOptions{ @@ -40,7 +40,7 @@ func gcpAuthHelper(ctx context.Context,t *testing.T) (tokenValue, projectID stri log.Fatalf("Failed to get project ID: %v", err) } if projectID == "" { - //Try quota project + //Try quota project projectID, err = creds.QuotaProjectID(ctx) if err != nil { log.Fatalf("Failed to get project ID: %v", err) @@ -60,17 +60,17 @@ func gcpAuthHelper(ctx context.Context,t *testing.T) (tokenValue, projectID stri log.Fatalf("Failed to retrieve access token: %v", err) } - return token.Value,projectID + return token.Value, projectID } func TestRAGQuery(t *testing.T) { ctx := context.Background() - token, projectID := gcpAuthHelper(ctx,t) + token, projectID := gcpAuthHelper(ctx, t) vertexEmbeddingFunc := chromem.NewEmbeddingFuncVertex( - token, - projectID, - chromem.EmbeddingModelVertexEnglishV4) + token, + projectID, + chromem.EmbeddingModelVertexEnglishV4) db := chromem.NewDB() dbFile := os.Getenv("RAG_DB_PATH") @@ -80,9 +80,9 @@ func TestRAGQuery(t *testing.T) { //check if file exists, we expect an existing DB if _, err := os.Stat(dbFile); os.IsNotExist(err) { log.Fatalf("RAG_DB_PATH file does not exist, skipping import: %v", dbFile) - }else{ + } else { err := db.ImportFromFile(dbFile, "") - log.Printf("Imported RAG with collections:%d",len(db.ListCollections())) + log.Printf("Imported RAG with collections:%d", len(db.ListCollections())) if err != nil { log.Fatalf("Unable to import from the RAG DB file:%s - %v", dbFile, err) } @@ -90,27 +90,27 @@ func TestRAGQuery(t *testing.T) { collectionPattern, err := db.GetOrCreateCollection("pattern", nil, vertexEmbeddingFunc) if err != nil { - log.Fatalf("Unable to get collection pattern: %v",err) + log.Fatalf("Unable to get collection pattern: %v", err) } - patternResult,err := collectionPattern.Query(ctx, "Simple pipeline that deploys to Cloud Run",1,nil,nil) + patternResult, err := collectionPattern.Query(ctx, "Simple pipeline that deploys to Cloud Run", 1, nil, nil) if err != nil { - log.Fatalf("Unable to Query collection pattern: %v",err) + log.Fatalf("Unable to Query collection pattern: %v", err) } - if len(patternResult)<1 || patternResult[0].Content == ""{ - log.Fatalf("Failed to find pattern: %v",len(patternResult)) + if len(patternResult) < 1 || patternResult[0].Content == "" { + log.Fatalf("Failed to find pattern: %v", len(patternResult)) } collectionKnowledge, err := db.GetOrCreateCollection("knowledge", nil, vertexEmbeddingFunc) if err != nil { - log.Fatalf("Unable to get collection knowledge: %v",err) + log.Fatalf("Unable to get collection knowledge: %v", err) } - knowledgeResult,err := collectionKnowledge.Query(ctx, "Package a Python application",3,nil,nil) + knowledgeResult, err := collectionKnowledge.Query(ctx, "Package a Python application", 3, nil, nil) if err != nil { - log.Fatalf("Unable to Query collection knowledge: %v",err) + log.Fatalf("Unable to Query collection knowledge: %v", err) } - if len(knowledgeResult)<3 || knowledgeResult[0].Content == ""{ - log.Fatalf("Failed to find pattern: %v",len(knowledgeResult)) + if len(knowledgeResult) < 3 || knowledgeResult[0].Content == "" { + log.Fatalf("Failed to find pattern: %v", len(knowledgeResult)) } -} \ No newline at end of file +} diff --git a/local-rag/update-rag.go b/local-rag/update-rag.go index 795029b..bf0e7e1 100644 --- a/local-rag/update-rag.go +++ b/local-rag/update-rag.go @@ -12,20 +12,20 @@ import ( "github.com/tmc/langchaingo/textsplitter" ) -func addDirectoryToRag(ctx context.Context,collection *chromem.Collection, dir string) { +func addDirectoryToRag(ctx context.Context, collection *chromem.Collection, dir string) { var docs []chromem.Document log.Printf("Uploading directory %s to collection: %v", dir, collection.Name) splitter := textsplitter.NewMarkdownTextSplitter( textsplitter.WithChunkSize(1000), textsplitter.WithChunkOverlap(150), ) - + filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { if err != nil { return err } if !info.IsDir() { - _,err := collection.GetByID(ctx,path) + _, err := collection.GetByID(ctx, path) if err == nil { // log.Printf("Doc found %s: %v", path, err) // Skip if doc is already loaded @@ -43,8 +43,8 @@ func addDirectoryToRag(ctx context.Context,collection *chromem.Collection, dir s return nil } for index, chunk := range chunks { - chunkId := path+"_"+strconv.Itoa(index) - _,err := collection.GetByID(ctx,chunkId) + chunkId := path + "_" + strconv.Itoa(index) + _, err := collection.GetByID(ctx, chunkId) if err == nil { // log.Printf("Doc found %s: %v", path, err) // Skip if doc is already loaded