tools.go

package main

import (
	"context"
	"fmt"
	"strings"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"google.golang.org/genai"
)

type Tool struct{}

// FetchKnownGeminiModels retrieves the list of available Gemini models.
func (t Tool) FetchKnownGeminiModels(ctx context.Context) (string, error) {
	var res []string
	client, err := genai.NewClient(ctx, nil)
	if err != nil {
		return "", err
	}
	for m, err := range client.Models.All(ctx) {
		if err != nil {
			return "", err
		}
		res = append(res, fmt.Sprintf("%s %s", m.Name, m.Description))
	}
	return strings.Join(res, "\n"), nil
}

// FetchAWSServices returns a list of services via Steampipe.
func (t Tool) FetchAWSServices(ctx context.Context) (string, error) {
	return queryPostgres(ctx, "SELECT DISTINCT foreign_table_name FROM information_schema.foreign_tables WHERE foreign_table_schema='aws'")
}

// GetPrompt lists available prompts from available MCP servers.
func (t Tool) ListPrompts(ctx context.Context) (string, error) {
	params, ok := ctx.Value("params").(*Parameters)
	if !ok {
		return "", fmt.Errorf("ListPrompts: params not found in context")
	}
	var res []string
	for _, sess := range params.MCPSessions {
		for p, err := range sess.Prompts(ctx, nil) {
			if err != nil {
				continue // skip this MCP server
			}
			desc := fmt.Sprintf("%s: %s", p.Name, p.Description)
			if len(p.Arguments) > 0 {
				desc += " ("
			}
			var args []string
			for _, arg := range p.Arguments {
				name := arg.Name
				if !arg.Required {
					name += " optional"
				}
				args = append(args, name)
			}
			if len(p.Arguments) > 0 {
				desc += fmt.Sprintf("%s) ", strings.Join(args, ","))
			}
			res = append(res, desc)
		}
	}
	return strings.Join(res, "\n"), nil
}

// GetPrompt retrieves a specific prompt by name from available MCP servers.
func (t Tool) GetPrompt(ctx context.Context, name string) (string, error) {
	params, ok := ctx.Value("params").(*Parameters)
	if !ok {
		return "", fmt.Errorf("GetPrompt: params not found in context")
	}
	keyVals, ok := ctx.Value("keyVals").(ParamMap)
	if !ok {
		return "", fmt.Errorf("GetPrompt: keyVals not found in context")
	}
	for _, sess := range params.MCPSessions {
		for p, err := range sess.Prompts(ctx, nil) {
			if err != nil {
				continue // skip this MCP server
			}
			if name == p.Name {
				prompt, err := sess.GetPrompt(ctx, &mcp.GetPromptParams{
					Name:      name,
					Arguments: keyVals,
				})
				if err != nil {
					return "", err
				}
				var res []string
				for _, msg := range prompt.Messages {
					switch msg.Content.(type) {
					case *mcp.TextContent:
						res = append(res, msg.Content.(*mcp.TextContent).Text)
					case *mcp.ResourceLink:
						res = append(res, fmt.Sprintf("%s %s",
							msg.Content.(*mcp.ResourceLink).URI,
							msg.Content.(*mcp.ResourceLink).MIMEType))
					case *mcp.EmbeddedResource:
						if msg.Content.(*mcp.EmbeddedResource).Resource != nil {
							res = append(res, fmt.Sprintf("%s %s",
								msg.Content.(*mcp.EmbeddedResource).Resource.URI,
								msg.Content.(*mcp.EmbeddedResource).Resource.MIMEType))
						}
					}
				}
				return strings.Join(res, "\n"), nil
			}
		}
	}
	return fmt.Sprintf("GetPrompt: '%s' not found", name), nil
}

// ListResources returns resources available from MCP servers.
func (t Tool) ListResources(ctx context.Context) (string, error) {
	params, ok := ctx.Value("params").(*Parameters)
	if !ok {
		return "", fmt.Errorf("ListResources: params not found in context")
	}
	var res []string
	for _, sess := range params.MCPSessions {
		for r, err := range sess.Resources(ctx, nil) {
			if err != nil {
				continue // skip this MCP server
			}
			res = append(res, r.URI)
		}
	}
	return strings.Join(res, "\n"), nil
}

// GetResource retrieves a specific resource by name from available MCP servers.
func (t Tool) GetResource(ctx context.Context, name string) (string, error) {
	params, ok := ctx.Value("params").(*Parameters)
	if !ok {
		return "", fmt.Errorf("GetResource: params not found in context")
	}
	for _, sess := range params.MCPSessions {
		for r, err := range sess.Resources(ctx, nil) {
			if err != nil {
				continue // skip this MCP server
			}
			if name == r.URI {
				res, err := sess.ReadResource(ctx, &mcp.ReadResourceParams{
					URI: name,
				})
				if err != nil {
					return "", err
				}
				for _, c := range res.Contents {
					if len(c.Text) > 0 {
						return c.Text, nil
					}
					if len(c.Blob) > 0 {
						return string(c.Blob), nil
					}
				}
			}
		}
	}
	return fmt.Sprintf("GetResource: '%s' not found", name), nil
}