refactor

Author: huangjeff5

go/ai/resource.go

@@ -21,12 +21,69 @@ import (
 	"encoding/json"
 	"fmt"
 	"maps"
+	"net/url"
+	"strings"
 
 	"github.com/firebase/genkit/go/core"
-	coreresource "github.com/firebase/genkit/go/core/resource"
 	"github.com/firebase/genkit/go/internal/registry"
+	"github.com/yosida95/uritemplate/v3"
 )
 
+// normalizeURI normalizes a URI for template matching by removing query parameters,
+// fragments, and trailing slashes from the path.
+func normalizeURI(rawURI string) string {
+	// Parse the URI
+	u, err := url.Parse(rawURI)
+	if err != nil {
+		// If parsing fails, return the original URI
+		return rawURI
+	}
+
+	// Remove query parameters and fragment
+	u.RawQuery = ""
+	u.Fragment = ""
+
+	// Remove trailing slash from path (but not from the root path)
+	if len(u.Path) > 1 && strings.HasSuffix(u.Path, "/") {
+		u.Path = strings.TrimSuffix(u.Path, "/")
+	}
+
+	return u.String()
+}
+
+// matches checks if a URI matches the given URI template.
+func matches(templateStr, uri string) (bool, error) {
+	template, err := uritemplate.New(templateStr)
+	if err != nil {
+		return false, fmt.Errorf("invalid URI template %q: %w", templateStr, err)
+	}
+
+	normalizedURI := normalizeURI(uri)
+	values := template.Match(normalizedURI)
+	return len(values) > 0, nil
+}
+
+// extractVariables extracts variables from a URI using the given URI template.
+func extractVariables(templateStr, uri string) (map[string]string, error) {
+	template, err := uritemplate.New(templateStr)
+	if err != nil {
+		return nil, fmt.Errorf("invalid URI template %q: %w", templateStr, err)
+	}
+
+	normalizedURI := normalizeURI(uri)
+	values := template.Match(normalizedURI)
+	if len(values) == 0 {
+		return nil, fmt.Errorf("URI %q does not match template", uri)
+	}
+
+	// Convert uritemplate.Values to string map
+	result := make(map[string]string)
+	for name, value := range values {
+		result[name] = value.String()
+	}
+	return result, nil
+}
+
 // ResourceInput represents the input to a resource function.
 type ResourceInput struct {
 	URI       string            `json:"uri"`       // The resource URI
@@ -75,10 +132,6 @@ func DefineResource(r *registry.Registry, name string, opts *ResourceOptions, fn
 	metadata := implementResource(name, opts)
 	resourceAction := core.DefineAction(r, name, core.ActionTypeResource, metadata, fn)
 	resource := &resource{Action: resourceAction}
-
-	// Register the resource wrapper so ListResources can find it
-	r.RegisterValue(fmt.Sprintf("resource/%s", name), resource)
-
 	return resource
 }
 
@@ -143,11 +196,11 @@ func (r *resource) Matches(uri string) bool {
 
 	// Check template
 	if template, ok := resourceMeta["template"].(string); ok && template != "" {
-		matcher, err := coreresource.NewTemplateMatcher(template)
+		matches, err := matches(template, uri)
 		if err != nil {
 			return false
 		}
-		return matcher.Matches(uri)
+		return matches
 	}
 
 	return false
@@ -171,11 +224,7 @@ func (r *resource) ExtractVariables(uri string) (map[string]string, error) {
 
 	// Extract from template
 	if template, ok := resourceMeta["template"].(string); ok && template != "" {
-		matcher, err := coreresource.NewTemplateMatcher(template)
-		if err != nil {
-			return nil, fmt.Errorf("invalid template %q: %w", template, err)
-		}
-		return matcher.ExtractVariables(uri)
+		return extractVariables(template, uri)
 	}
 
 	return nil, fmt.Errorf("no URI or template found in resource metadata")

go/core/resource.go

@@ -25,7 +25,6 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
-	"strings"
 	"syscall"
 
 	"github.com/firebase/genkit/go/ai"
@@ -373,11 +372,9 @@ func ListTools(g *Genkit) []ai.Tool {
 			continue
 		}
 		actionDesc := action.Desc()
-		if strings.HasPrefix(actionDesc.Key, "/"+string(core.ActionTypeTool)+"/") {
-			// Extract tool name from key
-			toolName := strings.TrimPrefix(actionDesc.Key, "/"+string(core.ActionTypeTool)+"/")
-			// Lookup the actual tool
-			tool := LookupTool(g, toolName)
+		if actionDesc.Type == core.ActionTypeTool {
+			// Lookup the actual tool using the action name
+			tool := LookupTool(g, actionDesc.Name)
 			if tool != nil {
 				tools = append(tools, tool)
 			}
@@ -1038,11 +1035,10 @@ func ListResources(g *Genkit) []ai.Resource {
 		}
 		actionDesc := action.Desc()
 		if actionDesc.Type == core.ActionTypeResource {
-			// Look up the resource wrapper
-			if resourceValue := g.reg.LookupValue(fmt.Sprintf("resource/%s", actionDesc.Name)); resourceValue != nil {
-				if resource, ok := resourceValue.(ai.Resource); ok {
-					resources = append(resources, resource)
-				}
+			// Lookup the actual resource using the action name
+			resource := ai.LookupResource(g.reg, actionDesc.Name)
+			if resource != nil {
+				resources = append(resources, resource)
 			}
 		}
 	}

go/core/resource/matcher.go

@@ -102,11 +102,6 @@ func (h *MCPHost) Connect(ctx context.Context, g *genkit.Genkit, serverName stri
 
 	h.clients[serverName] = client
 
-	// Auto-register tools and prompts from the connected server
-	if _, err := client.GetActiveTools(ctx, g); err != nil {
-		logger.FromContext(ctx).Warn("Failed to register tools from server", "server", serverName, "error", err)
-	}
-
 	return nil
 }
 

go/genkit/genkit.go

@@ -23,7 +23,6 @@ import (
 	"time"
 
 	"github.com/firebase/genkit/go/ai"
-	"github.com/firebase/genkit/go/core"
 	"github.com/firebase/genkit/go/genkit"
 	"github.com/stretchr/testify/assert"
 )
@@ -263,7 +262,7 @@ func TestMCPContentFetch(t *testing.T) {
 
 	// Find resource that matches our test URI
 	testURI := "file://data/example.txt"
-	var matchingResource core.DetachedResourceAction
+	var matchingResource ai.Resource
 	for _, res := range resources {
 		if res.Matches(testURI) {
 			matchingResource = res