diff --git a/go/ai/document.go b/go/ai/document.go index 09e80b1f59..2dfb9fe3d2 100644 --- a/go/ai/document.go +++ b/go/ai/document.go @@ -38,6 +38,7 @@ type Part struct { Text string `json:"text,omitempty"` // valid for kind∈{text,blob} ToolRequest *ToolRequest `json:"toolRequest,omitempty"` // valid for kind==partToolRequest ToolResponse *ToolResponse `json:"toolResponse,omitempty"` // valid for kind==partToolResponse + Resource *ResourcePart `json:"resource,omitempty"` // valid for kind==partResource Custom map[string]any `json:"custom,omitempty"` // valid for plugin-specific custom parts Metadata map[string]any `json:"metadata,omitempty"` // valid for all kinds } @@ -52,6 +53,7 @@ const ( PartToolResponse PartCustom PartReasoning + PartResource ) // NewTextPart returns a Part containing text. @@ -118,6 +120,11 @@ func NewReasoningPart(text string, signature []byte) *Part { } } +// NewResourcePart returns a Part containing a resource reference. +func NewResourcePart(uri string) *Part { + return &Part{Kind: PartResource, Resource: &ResourcePart{Uri: uri}} +} + // IsText reports whether the [Part] contains plain text. func (p *Part) IsText() bool { return p.Kind == PartText @@ -153,6 +160,11 @@ func (p *Part) IsReasoning() bool { return p.Kind == PartReasoning } +// IsResource reports whether the [Part] contains a resource reference. +func (p *Part) IsResource() bool { + return p.Kind == PartResource +} + // MarshalJSON is called by the JSON marshaler to write out a Part. func (p *Part) MarshalJSON() ([]byte, error) { // This is not handled by the schema generator because @@ -192,6 +204,12 @@ func (p *Part) MarshalJSON() ([]byte, error) { Metadata: p.Metadata, } return json.Marshal(v) + case PartResource: + v := resourcePart{ + Resource: p.Resource, + Metadata: p.Metadata, + } + return json.Marshal(v) case PartCustom: v := customPart{ Custom: p.Custom, @@ -215,6 +233,7 @@ type partSchema struct { Data string `json:"data,omitempty" yaml:"data,omitempty"` ToolRequest *ToolRequest `json:"toolRequest,omitempty" yaml:"toolRequest,omitempty"` ToolResponse *ToolResponse `json:"toolResponse,omitempty" yaml:"toolResponse,omitempty"` + Resource *ResourcePart `json:"resource,omitempty" yaml:"resource,omitempty"` Custom map[string]any `json:"custom,omitempty" yaml:"custom,omitempty"` Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty"` Reasoning string `json:"reasoning,omitempty" yaml:"reasoning,omitempty"` @@ -233,6 +252,9 @@ func (p *Part) unmarshalPartFromSchema(s partSchema) { case s.ToolResponse != nil: p.Kind = PartToolResponse p.ToolResponse = s.ToolResponse + case s.Resource != nil: + p.Kind = PartResource + p.Resource = s.Resource case s.Custom != nil: p.Kind = PartCustom p.Custom = s.Custom diff --git a/go/ai/gen.go b/go/ai/gen.go index c7c38c2836..941b48855a 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -129,6 +129,7 @@ type GenerateActionOptions struct { Output *GenerateActionOutputConfig `json:"output,omitempty"` Resume *GenerateActionResume `json:"resume,omitempty"` ReturnToolRequests bool `json:"returnToolRequests,omitempty"` + StepName string `json:"stepName,omitempty"` ToolChoice ToolChoice `json:"toolChoice,omitempty"` Tools []string `json:"tools,omitempty"` } @@ -315,6 +316,15 @@ type RerankerResponse struct { Documents []*RankedDocumentData `json:"documents,omitempty"` } +type resourcePart struct { + Metadata map[string]any `json:"metadata,omitempty"` + Resource *ResourcePart `json:"resource,omitempty"` +} + +type ResourcePart struct { + Uri string `json:"uri,omitempty"` +} + type RetrieverRequest struct { Options any `json:"options,omitempty"` Query *Document `json:"query,omitempty"` diff --git a/go/ai/generate.go b/go/ai/generate.go index 607b60ccd4..35438469ec 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -389,6 +389,15 @@ func Generate(ctx context.Context, r *registry.Registry, opts ...GenerateOption) } } + if len(genOpts.Resources) > 0 { + r = r.NewChild() + + // Attach resources + for _, res := range genOpts.Resources { + res.Register(r) + } + } + messages := []*Message{} if genOpts.SystemFn != nil { system, err := genOpts.SystemFn(ctx, nil) @@ -467,6 +476,13 @@ func Generate(ctx context.Context, r *registry.Registry, opts ...GenerateOption) } } + // Process resources in messages + processedMessages, err := processResources(ctx, r, messages) + if err != nil { + return nil, core.NewError(core.INTERNAL, "ai.Generate: error processing resources: %v", err) + } + actionOpts.Messages = processedMessages + return GenerateWithRequest(ctx, r, actionOpts, genOpts.Middleware, genOpts.Stream) } @@ -1048,3 +1064,87 @@ func handleResumeOption(ctx context.Context, r *registry.Registry, genOpts *Gene toolMessage: toolMessage, }, nil } + +// processResources processes messages to replace resource parts with actual content. +func processResources(ctx context.Context, r *registry.Registry, messages []*Message) ([]*Message, error) { + processedMessages := make([]*Message, len(messages)) + for i, msg := range messages { + processedContent := []*Part{} + + for _, part := range msg.Content { + if part.IsResource() { + // Find and execute the matching resource + resourceParts, err := executeResourcePart(ctx, r, part.Resource.Uri) + if err != nil { + return nil, fmt.Errorf("failed to process resource %q: %w", part.Resource, err) + } + // Replace resource part with content parts + processedContent = append(processedContent, resourceParts...) + } else { + // Keep non-resource parts as-is + processedContent = append(processedContent, part) + } + } + + processedMessages[i] = &Message{ + Role: msg.Role, + Content: processedContent, + Metadata: msg.Metadata, + } + } + + return processedMessages, nil +} + +// findMatchingResource finds a resource action in the registry that matches the given URI. +func findMatchingResource(r *registry.Registry, uri string) (core.Action, map[string]string, error) { + // Use our updated FindMatchingResource function + resource, resourceInput, err := FindMatchingResource(r, uri) + if err != nil { + return nil, nil, err + } + + // Execute the resource to get the action - we need to access the underlying action + // Since the resource interface doesn't expose the action directly, we'll look it up + action := r.LookupAction(fmt.Sprintf("/resource/%s", resource.Name())) + if action == nil { + return nil, nil, core.NewError(core.INTERNAL, "failed to lookup resource action") + } + + if coreAction, ok := action.(core.Action); ok { + return coreAction, resourceInput.Variables, nil + } + + return nil, nil, core.NewError(core.INTERNAL, "action does not implement core.Action interface") +} + +// executeResourcePart finds and executes a resource, returning the content parts. +func executeResourcePart(ctx context.Context, r *registry.Registry, resourceURI string) ([]*Part, error) { + action, variables, err := findMatchingResource(r, resourceURI) + if err != nil { + return nil, err + } + + // Create resource input with extracted variables + input := &ResourceInput{ + URI: resourceURI, + Variables: variables, + } + + // Execute the resource action directly + inputJSON, _ := json.Marshal(input) + outputJSON, err := action.RunJSON(ctx, inputJSON, nil) + if err != nil { + return nil, fmt.Errorf("failed to execute resource %q: %w", resourceURI, err) + } + + // Parse resource output - use a compatible structure + var output struct { + Content []*Part `json:"content"` + } + if err := json.Unmarshal(outputJSON, &output); err != nil { + return nil, fmt.Errorf("failed to parse resource output: %w", err) + } + + return output.Content, nil +} diff --git a/go/ai/generate_test.go b/go/ai/generate_test.go index d5ec8530ea..4a6def9042 100644 --- a/go/ai/generate_test.go +++ b/go/ai/generate_test.go @@ -1081,3 +1081,71 @@ func TestToolInterruptsAndResume(t *testing.T) { } }) } + +func TestResourceProcessing(t *testing.T) { + r := registry.New() + + // Create test resources using DefineResource + DefineResource(r, "test-file", &ResourceOptions{ + URI: "file:///test.txt", + Description: "Test file resource", + }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { + return &ResourceOutput{Content: []*Part{NewTextPart("FILE CONTENT")}}, nil + }) + + DefineResource(r, "test-api", &ResourceOptions{ + URI: "api://data/123", + Description: "Test API resource", + }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { + return &ResourceOutput{Content: []*Part{NewTextPart("API DATA")}}, nil + }) + + // Test message with resources + messages := []*Message{ + NewUserMessage( + NewTextPart("Read this:"), + NewResourcePart("file:///test.txt"), + NewTextPart("And this:"), + NewResourcePart("api://data/123"), + NewTextPart("Done."), + ), + } + + // Process resources + processed, err := processResources(context.Background(), r, messages) + if err != nil { + t.Fatalf("resource processing failed: %v", err) + } + + // Verify content + content := processed[0].Content + expected := []string{"Read this:", "FILE CONTENT", "And this:", "API DATA", "Done."} + + if len(content) != len(expected) { + t.Fatalf("expected %d parts, got %d", len(expected), len(content)) + } + + for i, want := range expected { + if content[i].Text != want { + t.Fatalf("part %d: got %q, want %q", i, content[i].Text, want) + } + } +} + +func TestResourceProcessingError(t *testing.T) { + r := registry.New() + + // No resources registered + messages := []*Message{ + NewUserMessage(NewResourcePart("missing://resource")), + } + + _, err := processResources(context.Background(), r, messages) + if err == nil { + t.Fatal("expected error when no resources available") + } + + if !strings.Contains(err.Error(), "no resource found for URI") { + t.Fatalf("wrong error: %v", err) + } +} diff --git a/go/ai/option.go b/go/ai/option.go index 33710732d9..41b6dbd54a 100644 --- a/go/ai/option.go +++ b/go/ai/option.go @@ -106,6 +106,7 @@ type commonGenOptions struct { Model ModelArg // Model to use. MessagesFn MessagesFn // Function to generate messages. Tools []ToolRef // References to tools to use. + Resources []Resource // Resources to be temporarily available during generation. ToolChoice ToolChoice // Whether tool calls are required, disabled, or optional. MaxTurns int // Maximum number of tool call iterations. ReturnToolRequests *bool // Whether to return tool requests instead of making the tool calls and continuing the generation. @@ -147,6 +148,13 @@ func (o *commonGenOptions) applyCommonGen(opts *commonGenOptions) error { opts.Tools = o.Tools } + if o.Resources != nil { + if opts.Resources != nil { + return errors.New("cannot set resources more than once (WithResources)") + } + opts.Resources = o.Resources + } + if o.ToolChoice != "" { if opts.ToolChoice != "" { return errors.New("cannot set tool choice more than once (WithToolChoice)") @@ -816,6 +824,34 @@ func WithToolRestarts(parts ...*Part) GenerateOption { return &generateOptions{RestartParts: parts} } +// WithResources specifies resources to be temporarily available during generation. +// Resources are unregistered resources that get attached to a temporary registry +// during the generation request and cleaned up afterward. +func WithResources(resources ...Resource) CommonGenOption { + return &withResources{resources: resources} +} + +type withResources struct { + resources []Resource +} + +func (w *withResources) applyCommonGen(o *commonGenOptions) error { + o.Resources = w.resources + return nil +} + +func (w *withResources) applyPrompt(o *promptOptions) error { + return w.applyCommonGen(&o.commonGenOptions) +} + +func (w *withResources) applyGenerate(o *generateOptions) error { + return w.applyCommonGen(&o.commonGenOptions) +} + +func (w *withResources) applyPromptExecute(o *promptExecutionOptions) error { + return w.applyCommonGen(&o.commonGenOptions) +} + // promptExecutionOptions are options for generating a model response by executing a prompt. type promptExecutionOptions struct { commonGenOptions diff --git a/go/ai/resource.go b/go/ai/resource.go new file mode 100644 index 0000000000..0fa32a5b7a --- /dev/null +++ b/go/ai/resource.go @@ -0,0 +1,248 @@ +// Copyright 2025 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 +// +// http://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. +// +// SPDX-License-Identifier: Apache-2.0 + +package ai + +import ( + "context" + "encoding/json" + "fmt" + "maps" + + "github.com/firebase/genkit/go/core" + coreresource "github.com/firebase/genkit/go/core/resource" + "github.com/firebase/genkit/go/internal/registry" +) + +// ResourceInput represents the input to a resource function. +type ResourceInput struct { + URI string `json:"uri"` // The resource URI + Variables map[string]string `json:"variables"` // Extracted variables from URI template matching +} + +// ResourceOutput represents the output from a resource function. +type ResourceOutput struct { + Content []*Part `json:"content"` // The content parts returned by the resource +} + +// ResourceOptions configures a resource definition. +type ResourceOptions struct { + URI string // Static URI (mutually exclusive with Template) + Template string // URI template (mutually exclusive with URI) + Description string // Optional description + Metadata map[string]any // Optional metadata +} + +// ResourceFunc is a function that loads content for a resource. +type ResourceFunc = func(context.Context, *ResourceInput) (*ResourceOutput, error) + +// resource is the internal implementation of the Resource interface. +// It holds the underlying core action and allows looking up resources +// by name without knowing their specific input/output types. +type resource struct { + core.Action +} + +// Resource represents an instance of a resource. +type Resource interface { + // Name returns the name of the resource. + Name() string + // Matches reports whether this resource matches the given URI. + Matches(uri string) bool + // ExtractVariables extracts variables from a URI using this resource's template. + ExtractVariables(uri string) (map[string]string, error) + // Execute runs the resource with the given input. + Execute(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) + // Register sets the tracing state on the action and registers it with the registry. + Register(r *registry.Registry) +} + +// DefineResource creates a resource and registers it with the given Registry. +func DefineResource(r *registry.Registry, name string, opts *ResourceOptions, fn ResourceFunc) Resource { + metadata := implementResource(name, *opts) + resourceAction := core.DefineAction(r, name, core.ActionTypeResource, metadata, fn) + return &resource{Action: resourceAction} +} + +// NewResource creates a resource but does not register it in the registry. +// It can be registered later via the Register method. +func NewResource(name string, opts ResourceOptions, fn ResourceFunc) Resource { + metadata := implementResource(name, opts) + metadata["dynamic"] = true + resourceAction := core.NewAction(name, core.ActionTypeResource, metadata, fn) + return &resource{Action: resourceAction} +} + +// implementResource creates the metadata common to both DefineResource and NewResource. +func implementResource(name string, opts ResourceOptions) map[string]any { + // Validate options - panic like other Define* functions + if name == "" { + panic("resource name is required") + } + if opts.URI != "" && opts.Template != "" { + panic("cannot specify both URI and Template") + } + if opts.URI == "" && opts.Template == "" { + panic("must specify either URI or Template") + } + + // Create metadata with resource-specific information + metadata := map[string]any{ + "type": core.ActionTypeResource, + "name": name, + "description": opts.Description, + "resource": map[string]any{ + "uri": opts.URI, + "template": opts.Template, + }, + } + + // Add user metadata + if opts.Metadata != nil { + maps.Copy(metadata, opts.Metadata) + } + + return metadata +} + +// Name returns the resource name. +func (r *resource) Name() string { + return r.Action.Name() +} + +// Matches reports whether this resource matches the given URI. +func (r *resource) Matches(uri string) bool { + desc := r.Action.Desc() + resourceMeta, ok := desc.Metadata["resource"].(map[string]any) + if !ok { + return false + } + + // Check static URI + if staticURI, ok := resourceMeta["uri"].(string); ok && staticURI != "" { + return staticURI == uri + } + + // Check template + if template, ok := resourceMeta["template"].(string); ok && template != "" { + matcher, err := coreresource.NewTemplateMatcher(template) + if err != nil { + return false + } + return matcher.Matches(uri) + } + + return false +} + +// ExtractVariables extracts variables from a URI using this resource's template. +func (r *resource) ExtractVariables(uri string) (map[string]string, error) { + desc := r.Action.Desc() + resourceMeta, ok := desc.Metadata["resource"].(map[string]any) + if !ok { + return nil, fmt.Errorf("no resource metadata found") + } + + // Static URI has no variables + if staticURI, ok := resourceMeta["uri"].(string); ok && staticURI != "" { + if staticURI == uri { + return map[string]string{}, nil + } + return nil, fmt.Errorf("URI %q does not match static URI %q", uri, staticURI) + } + + // 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 nil, fmt.Errorf("no URI or template found in resource metadata") +} + +// Execute runs the resource with the given input. +func (r *resource) Execute(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { + // Marshal input to JSON for action call + inputJSON, err := json.Marshal(input) + if err != nil { + return nil, fmt.Errorf("failed to marshal resource input: %w", err) + } + + // Use the underlying action to execute the resource function + outputJSON, err := r.Action.RunJSON(ctx, inputJSON, nil) + if err != nil { + return nil, err + } + + // Unmarshal output back to ResourceOutput + var output ResourceOutput + if err := json.Unmarshal(outputJSON, &output); err != nil { + return nil, fmt.Errorf("failed to unmarshal resource output: %w", err) + } + + return &output, nil +} + +// Register sets the tracing state on the action and registers it with the registry. +func (r *resource) Register(reg *registry.Registry) { + r.Action.SetTracingState(reg.TracingState()) + reg.RegisterAction(fmt.Sprintf("/%s/%s", core.ActionTypeResource, r.Action.Name()), r.Action) +} + +// FindMatchingResource finds a resource that matches the given URI. +func FindMatchingResource(r *registry.Registry, uri string) (Resource, *ResourceInput, error) { + actions := r.ListActions() + + for _, act := range actions { + action, ok := act.(core.Action) + if !ok { + continue + } + + desc := action.Desc() + if desc.Type != core.ActionTypeResource { + continue + } + + // Parse resource from Action metadata and use for template resolution + resource := &resource{Action: action} + if resource.Matches(uri) { + variables, err := resource.ExtractVariables(uri) + if err != nil { + return nil, nil, err + } + return resource, &ResourceInput{URI: uri, Variables: variables}, nil + } + } + + return nil, nil, fmt.Errorf("no resource found for URI %q", uri) +} + +// LookupResource looks up the resource in the registry by provided name and returns it. +func LookupResource(r *registry.Registry, name string) Resource { + if name == "" { + return nil + } + + action := r.LookupAction(fmt.Sprintf("/%s/%s", core.ActionTypeResource, name)) + if action == nil { + return nil + } + return &resource{Action: action.(core.Action)} +} diff --git a/go/core/action.go b/go/core/action.go index 004088f6b6..875f2cf9c5 100644 --- a/go/core/action.go +++ b/go/core/action.go @@ -63,6 +63,7 @@ const ( ActionTypeFlow ActionType = "flow" ActionTypeModel ActionType = "model" ActionTypeExecutablePrompt ActionType = "executable-prompt" + ActionTypeResource ActionType = "resource" ActionTypeTool ActionType = "tool" ActionTypeUtil ActionType = "util" ActionTypeCustom ActionType = "custom" diff --git a/go/core/resource/matcher.go b/go/core/resource/matcher.go new file mode 100644 index 0000000000..fc8492aba8 --- /dev/null +++ b/go/core/resource/matcher.go @@ -0,0 +1,84 @@ +// Copyright 2025 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 +// +// http://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. +// +// SPDX-License-Identifier: Apache-2.0 + +package resource + +import ( + "fmt" + + "github.com/yosida95/uritemplate/v3" +) + +// URIMatcher handles URI matching for resources. +// This is an internal interface used by resource implementations. +type URIMatcher interface { + Matches(uri string) bool + ExtractVariables(uri string) (map[string]string, error) +} + +// NewStaticMatcher creates a matcher for exact URI matches. +func NewStaticMatcher(uri string) URIMatcher { + return &staticMatcher{uri: uri} +} + +// NewTemplateMatcher creates a matcher for URI template patterns. +func NewTemplateMatcher(templateStr string) (URIMatcher, error) { + template, err := uritemplate.New(templateStr) + if err != nil { + return nil, fmt.Errorf("invalid URI template %q: %w", templateStr, err) + } + return &templateMatcher{template: template}, nil +} + +// staticMatcher matches exact URIs. +type staticMatcher struct { + uri string +} + +func (m *staticMatcher) Matches(uri string) bool { + return m.uri == uri +} + +func (m *staticMatcher) ExtractVariables(uri string) (map[string]string, error) { + if !m.Matches(uri) { + return nil, fmt.Errorf("URI %q does not match static URI %q", uri, m.uri) + } + return map[string]string{}, nil +} + +// templateMatcher matches URI templates. +type templateMatcher struct { + template *uritemplate.Template +} + +func (m *templateMatcher) Matches(uri string) bool { + values := m.template.Match(uri) + return len(values) > 0 +} + +func (m *templateMatcher) ExtractVariables(uri string) (map[string]string, error) { + values := m.template.Match(uri) + 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 +} diff --git a/go/core/schemas.config b/go/core/schemas.config index d5b4150c82..6b32a97edb 100644 --- a/go/core/schemas.config +++ b/go/core/schemas.config @@ -109,6 +109,7 @@ TextPart.toolResponse omit TextPart.custom omit TextPart.metadata type map[string]any TextPart.reasoning omit +TextPart.resource omit MediaPart pkg ai MediaPart name mediaPart @@ -119,6 +120,7 @@ MediaPart.custom omit MediaPart.data omit MediaPart.metadata type map[string]any MediaPart.reasoning omit +MediaPart.resource omit ToolRequestPart pkg ai ToolRequestPart name toolRequestPart @@ -129,6 +131,7 @@ ToolRequestPart.toolResponse omit ToolRequestPart.custom omit ToolRequestPart.metadata type map[string]any ToolRequestPart.reasoning omit +ToolRequestPart.resource omit ToolResponsePart pkg ai ToolResponsePart name toolResponsePart @@ -139,6 +142,7 @@ ToolResponsePart.toolRequest omit ToolResponsePart.custom omit ToolResponsePart.metadata type map[string]any ToolResponsePart.reasoning omit +ToolResponsePart.resource omit DataPart pkg ai DataPart name dataPart @@ -149,6 +153,7 @@ DataPart.toolResponse omit DataPart.custom omit DataPart.metadata type map[string]any DataPart.reasoning omit +DataPart.resource omit ReasoningPart pkg ai ReasoningPart name reasoningPart @@ -159,10 +164,31 @@ ReasoningPart.toolRequest omit ReasoningPart.toolResponse omit ReasoningPart.custom omit ReasoningPart.metadata type map[string]any +ReasoningPart.resource omit CustomPart pkg ai CustomPart name customPart +ResourcePart pkg ai +ResourcePart name resourcePart +ResourcePart.text omit +ResourcePart.media omit +ResourcePart.toolRequest omit +ResourcePart.toolResponse omit +ResourcePart.custom omit +ResourcePart.data omit +ResourcePart.metadata type map[string]any +ResourcePart.reasoning omit +ResourcePart.resource type *ResourcePart + +# Use the auto-generated ResourcePartResource type for consistency with schema +ResourcePartResource pkg ai +ResourcePartResource name ResourcePart + + + + + ModelInfo pkg ai ModelInfoSupports pkg ai ModelInfoSupports.output type []string @@ -230,6 +256,7 @@ ModelResponse.message type *Message ModelResponse.request type *ModelRequest ModelResponse.usage type *GenerationUsage ModelResponse.raw omit +ModelResponse.operation omit # ModelResponseChunk ModelResponseChunk pkg ai @@ -287,6 +314,9 @@ that is passed to a streaming callback. Score omit +Operation omit +OperationError omit + Embedding.embedding type []float32 GenkitError omit diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 7ff3f3f9a8..ea12f285f3 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -974,3 +974,77 @@ func DefineFormat(g *Genkit, name string, formatter ai.Formatter) { func IsDefinedFormat(g *Genkit, name string) bool { return g.reg.LookupValue("/format/"+name) != nil } + +// DefineResource defines a resource and registers it with the Genkit instance. +// Resources provide content that can be referenced in prompts via URI. +// +// Example: +// +// DefineResource(g, "company-docs", &ai.ResourceOptions{ +// URI: "file:///docs/handbook.pdf", +// Description: "Company handbook", +// }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { +// content, err := os.ReadFile("/docs/handbook.pdf") +// if err != nil { +// return nil, err +// } +// return &ai.ResourceOutput{ +// Content: []*ai.Part{ai.NewTextPart(string(content))}, +// }, nil +// }) +func DefineResource(g *Genkit, name string, opts *ai.ResourceOptions, fn ai.ResourceFunc) ai.Resource { + return ai.DefineResource(g.reg, name, opts, fn) +} + +// FindMatchingResource finds a resource that matches the given URI. +func FindMatchingResource(g *Genkit, uri string) (ai.Resource, *ai.ResourceInput, error) { + return ai.FindMatchingResource(g.reg, uri) +} + +// NewResource creates an unregistered resource action that can be temporarily +// attached during generation via WithResources option. +// +// Example: +// +// dynamicRes := NewResource("user-data", &ai.ResourceOptions{ +// Template: "user://profile/{id}", +// }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { +// userID := input.Variables["id"] +// // Load user data dynamically... +// return &ai.ResourceOutput{Content: []*ai.Part{ai.NewTextPart(userData)}}, nil +// }) +// +// // Use in generation: +// ai.Generate(ctx, g, +// ai.WithPrompt([]*ai.Part{ +// ai.NewTextPart("Analyze this user:"), +// ai.NewResourcePart("user://profile/123"), +// }), +// ai.WithResources(dynamicRes), +// ) +func NewResource(name string, opts *ai.ResourceOptions, fn ai.ResourceFunc) ai.Resource { + // Delegate to ai implementation + return ai.NewResource(name, *opts, fn) +} + +// ListResources returns a slice of all resource actions +func ListResources(g *Genkit) []ai.Resource { + acts := g.reg.ListActions() + resources := []ai.Resource{} + for _, act := range acts { + action, ok := act.(core.Action) + if !ok { + continue + } + 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) + } + } + } + } + return resources +} diff --git a/go/genkit/genkit_test.go b/go/genkit/genkit_test.go index 0a6b4777c9..d847db1349 100644 --- a/go/genkit/genkit_test.go +++ b/go/genkit/genkit_test.go @@ -18,8 +18,10 @@ package genkit import ( "context" + "strings" "testing" + "github.com/firebase/genkit/go/ai" "github.com/firebase/genkit/go/core" ) @@ -58,3 +60,242 @@ func count(ctx context.Context, n int, cb func(context.Context, int) error) (int } return n, nil } + +func TestStaticResource(t *testing.T) { + g := Init(context.Background()) + + // Define static resource + DefineResource(g, "test-doc", &ai.ResourceOptions{ + URI: "file:///test.txt", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("test content")}, + }, nil + }) + + // Find and execute + resource, input, err := FindMatchingResource(g, "file:///test.txt") + if err != nil { + t.Fatalf("resource not found: %v", err) + } + + output, err := resource.Execute(context.Background(), input) + if err != nil { + t.Fatalf("resource execution failed: %v", err) + } + + if len(output.Content) != 1 || output.Content[0].Text != "test content" { + t.Fatalf("unexpected output: %v", output.Content) + } +} + +func TestDynamicResourceWithTemplate(t *testing.T) { + g := Init(context.Background()) + + dynResource := NewResource("user-profile", &ai.ResourceOptions{ + Template: "user://profile/{userID}", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + userID := input.Variables["userID"] + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("User: " + userID)}, + }, nil + }) + + // Register the resource to set up tracing state properly + dynResource.Register(g.reg) + + // Test URI matching + if !dynResource.Matches("user://profile/123") { + t.Fatal("should match user://profile/123") + } + + if dynResource.Matches("user://different/123") { + t.Fatal("should not match different URI scheme") + } + + // Test variable extraction and execution + variables, err := dynResource.ExtractVariables("user://profile/alice") + if err != nil { + t.Fatalf("failed to extract variables: %v", err) + } + + if variables["userID"] != "alice" { + t.Fatalf("expected userID=alice, got %s", variables["userID"]) + } + + // Execute with extracted variables + input := &ai.ResourceInput{ + URI: "user://profile/alice", + Variables: variables, + } + + output, err := dynResource.Execute(context.Background(), input) + if err != nil { + t.Fatalf("execution failed: %v", err) + } + + if len(output.Content) != 1 || output.Content[0].Text != "User: alice" { + t.Fatalf("unexpected output: %v", output.Content) + } +} + +func TestResourceInGeneration(t *testing.T) { + g := Init(context.Background()) + + // Define mock model + ai.DefineModel(g.reg, "test", nil, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { + // Extract resource parts from the prompt + var responseText strings.Builder + for _, msg := range req.Messages { + for _, part := range msg.Content { + if part.Text != "" { + responseText.WriteString(part.Text + " ") + } + } + } + + return &ai.ModelResponse{ + Request: req, + Message: &ai.Message{ + Content: []*ai.Part{ai.NewTextPart("Model response: " + responseText.String())}, + Role: "model", + }, + }, nil + }) + + // Define resource + DefineResource(g, "policy", &ai.ResourceOptions{ + URI: "file:///policy.txt", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("POLICY_CONTENT")}, + }, nil + }) + + // Generate with resource reference + resp, err := Generate(context.Background(), g, + ai.WithModelName("test"), + ai.WithMessages(ai.NewUserMessage( + ai.NewTextPart("Read this:"), + ai.NewResourcePart("file:///policy.txt"), + ai.NewTextPart("Done."), + )), + ) + + if err != nil { + t.Fatalf("generation failed: %v", err) + } + + if !contains(resp.Text(), "Model response") { + t.Fatalf("expected model response, got: %s", resp.Text()) + } +} + +func TestDynamicResourceInGeneration(t *testing.T) { + g := Init(context.Background()) + + // Define mock model + ai.DefineModel(g.reg, "test", nil, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { + var responseText strings.Builder + for _, msg := range req.Messages { + for _, part := range msg.Content { + if part.Text != "" { + responseText.WriteString(part.Text + " ") + } + } + } + + return &ai.ModelResponse{ + Request: req, + Message: &ai.Message{ + Content: []*ai.Part{ai.NewTextPart("Model response: " + responseText.String())}, + Role: "model", + }, + }, nil + }) + + // Create dynamic resource (not registered in registry) + dynResource := NewResource("dynamic-policy", &ai.ResourceOptions{ + URI: "dynamic://policy", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("DYNAMIC_CONTENT")}, + }, nil + }) + + // Generate with dynamic resource reference using WithResources + resp, err := Generate(context.Background(), g, + ai.WithModelName("test"), + ai.WithMessages(ai.NewUserMessage( + ai.NewTextPart("Read this:"), + ai.NewResourcePart("dynamic://policy"), + ai.NewTextPart("Done."), + )), + ai.WithResources(dynResource), + ) + + if err != nil { + t.Fatalf("generation failed: %v", err) + } + + if !contains(resp.Text(), "Model response") { + t.Fatalf("expected model response, got: %s", resp.Text()) + } +} + +func TestMultipleDynamicResourcesInGeneration(t *testing.T) { + g := Init(context.Background()) + + // Define mock model + ai.DefineModel(g.reg, "test", nil, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { + return &ai.ModelResponse{ + Request: req, + Message: &ai.Message{ + Content: []*ai.Part{ai.NewTextPart("Model processed all resources")}, + Role: "model", + }, + }, nil + }) + + // Create multiple dynamic resources + userResource := NewResource("user-data", &ai.ResourceOptions{ + Template: "user://profile/{id}", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("User: " + input.Variables["id"])}, + }, nil + }) + + projectResource := NewResource("project-data", &ai.ResourceOptions{ + URI: "project://settings", + }, func(ctx context.Context, input *ai.ResourceInput) (*ai.ResourceOutput, error) { + return &ai.ResourceOutput{ + Content: []*ai.Part{ai.NewTextPart("Project settings")}, + }, nil + }) + + // Generate with multiple dynamic resources + resp, err := Generate(context.Background(), g, + ai.WithModelName("test"), + ai.WithMessages(ai.NewUserMessage( + ai.NewTextPart("User:"), + ai.NewResourcePart("user://profile/alice"), + ai.NewTextPart("Project:"), + ai.NewResourcePart("project://settings"), + ai.NewTextPart("Done."), + )), + ai.WithResources(userResource, projectResource), + ) + + if err != nil { + t.Fatalf("generation failed: %v", err) + } + + if !contains(resp.Text(), "Model processed") { + t.Fatalf("expected model response, got: %s", resp.Text()) + } +} + +func contains(s, substr string) bool { + return strings.Contains(s, substr) +} diff --git a/go/go.mod b/go/go.mod index ed667e6b24..9926c69896 100644 --- a/go/go.mod +++ b/go/go.mod @@ -44,7 +44,7 @@ require ( require ( github.com/spf13/cast v1.7.1 // indirect - github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 ) require (