diff --git a/admin/server/ai.go b/admin/server/ai.go index b49f19cc49d3..4349610bfa64 100644 --- a/admin/server/ai.go +++ b/admin/server/ai.go @@ -50,6 +50,7 @@ func (s *Server) Complete(ctx context.Context, req *adminv1.CompleteRequest) (*a Messages: messages, Tools: req.Tools, OutputSchema: outputSchema, + CacheKey: req.CacheKey, }) if err != nil { return nil, err diff --git a/proto/gen/rill/admin/v1/admin.swagger.yaml b/proto/gen/rill/admin/v1/admin.swagger.yaml index c89fe22d8406..308798ecee49 100644 --- a/proto/gen/rill/admin/v1/admin.swagger.yaml +++ b/proto/gen/rill/admin/v1/admin.swagger.yaml @@ -5171,6 +5171,11 @@ definitions: outputJsonSchema: type: string title: Optional output JSON schema + cacheKey: + type: string + description: |- + Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + Providers may use it to improve prompt cache routing. v1CompleteResponse: type: object properties: diff --git a/proto/gen/rill/admin/v1/ai.pb.go b/proto/gen/rill/admin/v1/ai.pb.go index e5ed7ee2e806..afdc80ec9ef5 100644 --- a/proto/gen/rill/admin/v1/ai.pb.go +++ b/proto/gen/rill/admin/v1/ai.pb.go @@ -33,6 +33,9 @@ type CompleteRequest struct { Tools []*v1.Tool `protobuf:"bytes,2,rep,name=tools,proto3" json:"tools,omitempty"` // Optional output JSON schema OutputJsonSchema string `protobuf:"bytes,3,opt,name=output_json_schema,json=outputJsonSchema,proto3" json:"output_json_schema,omitempty"` + // Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + // Providers may use it to improve prompt cache routing. + CacheKey string `protobuf:"bytes,4,opt,name=cache_key,json=cacheKey,proto3" json:"cache_key,omitempty"` } func (x *CompleteRequest) Reset() { @@ -88,6 +91,13 @@ func (x *CompleteRequest) GetOutputJsonSchema() string { return "" } +func (x *CompleteRequest) GetCacheKey() string { + if x != nil { + return x.CacheKey + } + return "" +} + type CompleteResponse struct { state protoimpl.MessageState sizeCache protoimpl.SizeCache @@ -180,7 +190,7 @@ var file_rill_admin_v1_ai_proto_rawDesc = []byte{ 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x1a, 0x1c, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x61, 0x70, 0x69, 0x2f, 0x61, 0x6e, 0x6e, 0x6f, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x13, 0x72, 0x69, 0x6c, 0x6c, 0x2f, 0x61, 0x69, 0x2f, 0x76, - 0x31, 0x2f, 0x61, 0x69, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xa2, 0x01, 0x0a, 0x0f, 0x43, + 0x31, 0x2f, 0x61, 0x69, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xbf, 0x01, 0x0a, 0x0f, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x39, 0x0a, 0x08, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, 0x69, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, @@ -190,40 +200,42 @@ var file_rill_admin_v1_ai_proto_rawDesc = []byte{ 0x61, 0x69, 0x2e, 0x76, 0x31, 0x2e, 0x54, 0x6f, 0x6f, 0x6c, 0x52, 0x05, 0x74, 0x6f, 0x6f, 0x6c, 0x73, 0x12, 0x2c, 0x0a, 0x12, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x6a, 0x73, 0x6f, 0x6e, 0x5f, 0x73, 0x63, 0x68, 0x65, 0x6d, 0x61, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x6f, - 0x75, 0x74, 0x70, 0x75, 0x74, 0x4a, 0x73, 0x6f, 0x6e, 0x53, 0x63, 0x68, 0x65, 0x6d, 0x61, 0x22, - 0xdf, 0x01, 0x0a, 0x10, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, - 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x37, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, - 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, 0x69, 0x2e, - 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x69, 0x6f, 0x6e, 0x4d, 0x65, 0x73, - 0x73, 0x61, 0x67, 0x65, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x21, 0x0a, - 0x0c, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x0d, 0x52, 0x0b, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, - 0x12, 0x23, 0x0a, 0x0d, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, - 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0c, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x54, - 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x2e, 0x0a, 0x13, 0x63, 0x61, 0x63, 0x68, 0x65, 0x64, 0x5f, - 0x69, 0x6e, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x04, 0x20, 0x01, - 0x28, 0x0d, 0x52, 0x11, 0x63, 0x61, 0x63, 0x68, 0x65, 0x64, 0x49, 0x6e, 0x70, 0x75, 0x74, 0x54, - 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, - 0x72, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, - 0x72, 0x32, 0x74, 0x0a, 0x09, 0x41, 0x49, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x67, - 0x0a, 0x08, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x12, 0x1e, 0x2e, 0x72, 0x69, 0x6c, - 0x6c, 0x2e, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x6c, - 0x65, 0x74, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x72, 0x69, 0x6c, - 0x6c, 0x2e, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x6c, - 0x65, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x1a, 0x82, 0xd3, 0xe4, - 0x93, 0x02, 0x14, 0x3a, 0x01, 0x2a, 0x22, 0x0f, 0x2f, 0x76, 0x31, 0x2f, 0x61, 0x69, 0x2f, 0x63, - 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x42, 0xac, 0x01, 0x0a, 0x11, 0x63, 0x6f, 0x6d, 0x2e, - 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x42, 0x07, 0x41, - 0x69, 0x50, 0x72, 0x6f, 0x74, 0x6f, 0x50, 0x01, 0x5a, 0x38, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, - 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x72, 0x69, 0x6c, 0x6c, 0x64, 0x61, 0x74, 0x61, 0x2f, 0x72, 0x69, - 0x6c, 0x6c, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x72, 0x69, 0x6c, - 0x6c, 0x2f, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2f, 0x76, 0x31, 0x3b, 0x61, 0x64, 0x6d, 0x69, 0x6e, - 0x76, 0x31, 0xa2, 0x02, 0x03, 0x52, 0x41, 0x58, 0xaa, 0x02, 0x0d, 0x52, 0x69, 0x6c, 0x6c, 0x2e, - 0x41, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x56, 0x31, 0xca, 0x02, 0x0d, 0x52, 0x69, 0x6c, 0x6c, 0x5c, - 0x41, 0x64, 0x6d, 0x69, 0x6e, 0x5c, 0x56, 0x31, 0xe2, 0x02, 0x19, 0x52, 0x69, 0x6c, 0x6c, 0x5c, - 0x41, 0x64, 0x6d, 0x69, 0x6e, 0x5c, 0x56, 0x31, 0x5c, 0x47, 0x50, 0x42, 0x4d, 0x65, 0x74, 0x61, - 0x64, 0x61, 0x74, 0x61, 0xea, 0x02, 0x0f, 0x52, 0x69, 0x6c, 0x6c, 0x3a, 0x3a, 0x41, 0x64, 0x6d, - 0x69, 0x6e, 0x3a, 0x3a, 0x56, 0x31, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x75, 0x74, 0x70, 0x75, 0x74, 0x4a, 0x73, 0x6f, 0x6e, 0x53, 0x63, 0x68, 0x65, 0x6d, 0x61, 0x12, + 0x1b, 0x0a, 0x09, 0x63, 0x61, 0x63, 0x68, 0x65, 0x5f, 0x6b, 0x65, 0x79, 0x18, 0x04, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x08, 0x63, 0x61, 0x63, 0x68, 0x65, 0x4b, 0x65, 0x79, 0x22, 0xdf, 0x01, 0x0a, + 0x10, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, + 0x65, 0x12, 0x37, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x0b, 0x32, 0x1d, 0x2e, 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, 0x69, 0x2e, 0x76, 0x31, 0x2e, + 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x69, 0x6f, 0x6e, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, + 0x65, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x21, 0x0a, 0x0c, 0x69, 0x6e, + 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0d, + 0x52, 0x0b, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x23, 0x0a, + 0x0d, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x03, + 0x20, 0x01, 0x28, 0x0d, 0x52, 0x0c, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x54, 0x6f, 0x6b, 0x65, + 0x6e, 0x73, 0x12, 0x2e, 0x0a, 0x13, 0x63, 0x61, 0x63, 0x68, 0x65, 0x64, 0x5f, 0x69, 0x6e, 0x70, + 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0d, 0x52, + 0x11, 0x63, 0x61, 0x63, 0x68, 0x65, 0x64, 0x49, 0x6e, 0x70, 0x75, 0x74, 0x54, 0x6f, 0x6b, 0x65, + 0x6e, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x18, 0x05, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x32, 0x74, + 0x0a, 0x09, 0x41, 0x49, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x67, 0x0a, 0x08, 0x43, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, 0x12, 0x1e, 0x2e, 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, + 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, + 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x72, 0x69, 0x6c, 0x6c, 0x2e, 0x61, + 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x74, 0x65, + 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x1a, 0x82, 0xd3, 0xe4, 0x93, 0x02, 0x14, + 0x3a, 0x01, 0x2a, 0x22, 0x0f, 0x2f, 0x76, 0x31, 0x2f, 0x61, 0x69, 0x2f, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x74, 0x65, 0x42, 0xac, 0x01, 0x0a, 0x11, 0x63, 0x6f, 0x6d, 0x2e, 0x72, 0x69, 0x6c, + 0x6c, 0x2e, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x76, 0x31, 0x42, 0x07, 0x41, 0x69, 0x50, 0x72, + 0x6f, 0x74, 0x6f, 0x50, 0x01, 0x5a, 0x38, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, + 0x6d, 0x2f, 0x72, 0x69, 0x6c, 0x6c, 0x64, 0x61, 0x74, 0x61, 0x2f, 0x72, 0x69, 0x6c, 0x6c, 0x2f, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x67, 0x65, 0x6e, 0x2f, 0x72, 0x69, 0x6c, 0x6c, 0x2f, 0x61, + 0x64, 0x6d, 0x69, 0x6e, 0x2f, 0x76, 0x31, 0x3b, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x76, 0x31, 0xa2, + 0x02, 0x03, 0x52, 0x41, 0x58, 0xaa, 0x02, 0x0d, 0x52, 0x69, 0x6c, 0x6c, 0x2e, 0x41, 0x64, 0x6d, + 0x69, 0x6e, 0x2e, 0x56, 0x31, 0xca, 0x02, 0x0d, 0x52, 0x69, 0x6c, 0x6c, 0x5c, 0x41, 0x64, 0x6d, + 0x69, 0x6e, 0x5c, 0x56, 0x31, 0xe2, 0x02, 0x19, 0x52, 0x69, 0x6c, 0x6c, 0x5c, 0x41, 0x64, 0x6d, + 0x69, 0x6e, 0x5c, 0x56, 0x31, 0x5c, 0x47, 0x50, 0x42, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, + 0x61, 0xea, 0x02, 0x0f, 0x52, 0x69, 0x6c, 0x6c, 0x3a, 0x3a, 0x41, 0x64, 0x6d, 0x69, 0x6e, 0x3a, + 0x3a, 0x56, 0x31, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( diff --git a/proto/gen/rill/admin/v1/ai.pb.validate.go b/proto/gen/rill/admin/v1/ai.pb.validate.go index e7518dead177..749875110f0d 100644 --- a/proto/gen/rill/admin/v1/ai.pb.validate.go +++ b/proto/gen/rill/admin/v1/ai.pb.validate.go @@ -127,6 +127,8 @@ func (m *CompleteRequest) validate(all bool) error { // no validation rules for OutputJsonSchema + // no validation rules for CacheKey + if len(errors) > 0 { return CompleteRequestMultiError(errors) } diff --git a/proto/gen/rill/admin/v1/openapi.yaml b/proto/gen/rill/admin/v1/openapi.yaml index 340183efc661..ba2cfa132687 100644 --- a/proto/gen/rill/admin/v1/openapi.yaml +++ b/proto/gen/rill/admin/v1/openapi.yaml @@ -382,6 +382,11 @@ components: type: object v1CompleteRequest: properties: + cacheKey: + description: |- + Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + Providers may use it to improve prompt cache routing. + type: string messages: description: Input message(s) for the AI to complete. items: diff --git a/proto/gen/rill/admin/v1/public.openapi.yaml b/proto/gen/rill/admin/v1/public.openapi.yaml index 7fd099606ac2..9c1d4c86df12 100644 --- a/proto/gen/rill/admin/v1/public.openapi.yaml +++ b/proto/gen/rill/admin/v1/public.openapi.yaml @@ -382,6 +382,11 @@ components: type: object v1CompleteRequest: properties: + cacheKey: + description: |- + Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + Providers may use it to improve prompt cache routing. + type: string messages: description: Input message(s) for the AI to complete. items: diff --git a/proto/rill/admin/v1/ai.proto b/proto/rill/admin/v1/ai.proto index 6284613a0a08..4f94f2d15fc5 100644 --- a/proto/rill/admin/v1/ai.proto +++ b/proto/rill/admin/v1/ai.proto @@ -23,6 +23,9 @@ message CompleteRequest { repeated rill.ai.v1.Tool tools = 2; // Optional output JSON schema string output_json_schema = 3; + // Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + // Providers may use it to improve prompt cache routing. + string cache_key = 4; } message CompleteResponse { diff --git a/runtime/ai/ai.go b/runtime/ai/ai.go index 6dcc274641d8..a5388c734e17 100644 --- a/runtime/ai/ai.go +++ b/runtime/ai/ai.go @@ -29,6 +29,7 @@ import ( "go.uber.org/zap" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/structpb" ) @@ -36,6 +37,11 @@ import ( // Exceeding it will result in an error. const maxMessageSizeBytes = 100 * 1024 // 100 KB +// reservedInputTokens is headroom subtracted from the model's input token limit to make room for +// tool definitions, model output (some providers share one context window between input and output), +// and token estimation error. +const reservedInputTokens = 64_000 + // Tracer for instrumenting requests. var tracer = otel.Tracer("github.com/rilldata/rill/runtime/ai") @@ -1212,6 +1218,14 @@ func (s *Session) Complete(ctx context.Context, name string, out any, opts *Comp } defer release() + // Determine the token budget for input messages based on the model's input token limit. + // The reserved headroom is capped at a quarter of the limit so small limits keep a usable budget. + maxInputTokens := llm.MaxInputTokens() + if maxInputTokens <= 0 { + maxInputTokens = drivers.DefaultAIMaxInputTokens + } + messagesTokenBudget := maxInputTokens - min(reservedInputTokens, maxInputTokens/4) + // Setup input messages. messages := slices.Clone(opts.Messages) @@ -1287,7 +1301,7 @@ func (s *Session) Complete(ctx context.Context, name string, out any, opts *Comp } // Truncate messages to fit within LLM context window. - truncMessages := maybeTruncateMessages(messages) + truncMessages := maybeTruncateMessages(messages, messagesTokenBudget) // Telemetry iterations++ @@ -1304,6 +1318,7 @@ func (s *Session) Complete(ctx context.Context, name string, out any, opts *Comp Messages: truncMessages, Tools: tools, OutputSchema: outputSchema, + CacheKey: s.id, }) llmCancel() @@ -1600,26 +1615,59 @@ func NewTextCompletionMessage(role Role, content string) *aiv1.CompletionMessage } } +// Tuning constants for maybeTruncateMessages. +const ( + truncateKeepFirst = 4 // Always keep the first messages for context + truncateStep = 20 // Skip messages in multiples of this +) + // maybeTruncateMessages keeps recent messages and a few early ones for context. // It's a simple placeholder strategy. In the future, we'll enhance this with AI summarization. -func maybeTruncateMessages(messages []*aiv1.CompletionMessage) []*aiv1.CompletionMessage { - const ( - maxMessages = 20 // Keep up to 20 messages total - keepFirst = 4 // Always keep first 4 messages for context - keepLast = 16 // Keep last 16 messages - ) +// +// Truncation triggers when the estimated token count of the messages exceeds maxTokens. +// +// LLM prompt caching requires the message prefix to be byte-stable across requests, so the number +// of skipped messages only changes in coarse steps of truncateStep: while up to truncateStep new +// messages accumulate, the result is an append-only extension of the previous result, causing a +// cache miss once per step instead of on every request. +func maybeTruncateMessages(messages []*aiv1.CompletionMessage, maxTokens int) []*aiv1.CompletionMessage { + if maxTokens <= 0 || len(messages) <= truncateKeepFirst+1 { + return messages + } + + // Determine the smallest number of messages to skip such that the estimated token count of the + // kept messages fits within maxTokens. Skipping never reaches the last message: if the first + // and last messages alone exceed the budget, truncation can't help and we send them anyway. + var skipped int + var keptTokens int + est := make([]int, len(messages)) + for i, m := range messages { + est[i] = estimateMessageTokens(m) + keptTokens += est[i] + } + for keptTokens > maxTokens && truncateKeepFirst+skipped < len(messages)-1 { + keptTokens -= est[truncateKeepFirst+skipped] + skipped++ + } - if len(messages) <= maxMessages { + if skipped == 0 { return messages } + // Round the skipped count up to a multiple of truncateStep to keep the skipped range stable + // until a step's worth of additional tokens has accumulated. + skipped = ((skipped + truncateStep - 1) / truncateStep) * truncateStep + // do not skip + if skipped > len(messages)-truncateKeepFirst-1 { + skipped = len(messages) - truncateKeepFirst - 1 + } + var result []*aiv1.CompletionMessage // Keep first messages - result = append(result, messages[:keepFirst]...) + result = append(result, messages[:truncateKeepFirst]...) // Add truncation indicator - skipped := len(messages) - keepFirst - keepLast result = append(result, &aiv1.CompletionMessage{ Role: "system", Content: []*aiv1.ContentBlock{ @@ -1631,9 +1679,8 @@ func maybeTruncateMessages(messages []*aiv1.CompletionMessage) []*aiv1.Completio }, }) - // Keep last messages - start := len(messages) - keepLast - result = append(result, messages[start:]...) + // Keep messages after the skipped range + result = append(result, messages[truncateKeepFirst+skipped:]...) // Make sure there are no partial tool calls/results unbalancedIDs := make(map[string]bool) @@ -1660,6 +1707,13 @@ func maybeTruncateMessages(messages []*aiv1.CompletionMessage) []*aiv1.Completio return result } +// estimateMessageTokens conservatively estimates the number of LLM tokens in a message. +// It uses the proto-encoded size at 3 bytes per token: typical English text is ~4 bytes per token, +// but dense JSON and numeric content (common in tool results) can be as low as ~2-3. +func estimateMessageTokens(m *aiv1.CompletionMessage) int { + return proto.Size(m) / 3 +} + // completionMessageID turns a UUID into a truncated ID suitable for use in completion messages (which don't require IDs to be globally unique). func completionMessageID(id string) string { return strings.ReplaceAll(id, "-", "")[0:16] diff --git a/runtime/ai/ai_internal_test.go b/runtime/ai/ai_internal_test.go new file mode 100644 index 000000000000..5ebad71604d2 --- /dev/null +++ b/runtime/ai/ai_internal_test.go @@ -0,0 +1,127 @@ +package ai + +import ( + "fmt" + "strings" + "testing" + + aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1" + "github.com/stretchr/testify/require" +) + +func TestMaybeTruncateMessagesWithinBudget(t *testing.T) { + messages := newTextMessages(500) + result := maybeTruncateMessages(messages, sumEstimatedTokens(messages)) + require.Equal(t, messages, result) +} + +func TestMaybeTruncateMessagesNoBudget(t *testing.T) { + // A budget <= 0 disables truncation entirely. + messages := newTextMessages(500) + result := maybeTruncateMessages(messages, 0) + require.Equal(t, messages, result) +} + +func TestMaybeTruncateMessagesOverBudget(t *testing.T) { + messages := newTextMessages(121) + // A budget that fits all but one message forces a skip, which rounds up to truncateStep. + result := maybeTruncateMessages(messages, sumEstimatedTokens(messages)-1) + + // Expect the first messages, a truncation indicator, then everything after the skipped range. + require.Len(t, result, truncateKeepFirst+1+len(messages)-truncateKeepFirst-truncateStep) + require.Equal(t, messages[:truncateKeepFirst], result[:truncateKeepFirst]) + require.Equal(t, fmt.Sprintf("... [%d messages omitted for brevity] ...", truncateStep), result[truncateKeepFirst].Content[0].GetText()) + require.Equal(t, messages[truncateKeepFirst+truncateStep:], result[truncateKeepFirst+1:]) + require.LessOrEqual(t, sumEstimatedTokens(result), sumEstimatedTokens(messages)) +} + +// TestMaybeTruncateMessagesStablePrefix verifies that as messages accumulate, each truncated +// result is an append-only extension of the previous one until the skipped count steps up. +// This property is what keeps LLM prompt caching effective across completion iterations. +func TestMaybeTruncateMessagesStablePrefix(t *testing.T) { + // Uniform-sized messages so the token budget corresponds to an exact message count. + const keepMessages = 100 + const maxN = 200 + budget := sumEstimatedTokens(newTextMessages(keepMessages)) + + var steps int + prev := maybeTruncateMessages(newTextMessages(keepMessages+1), budget) + require.Less(t, len(prev), keepMessages+1, "expected truncation to trigger") + for n := keepMessages + 2; n <= maxN; n++ { + curr := maybeTruncateMessages(newTextMessages(n), budget) + if len(curr) > len(prev) { + // Within a step: the previous result must be a prefix of the current result. + require.Equal(t, prev, curr[:len(prev)], "prefix changed at %d messages", n) + } else { + // The skipped count stepped up: the result shrinks by truncateStep-1 (truncateStep more skipped, 1 appended). + require.Len(t, curr, len(prev)-(truncateStep-1), "unexpected step at %d messages", n) + steps++ + } + prev = curr + } + require.Equal(t, (maxN-keepMessages-1)/truncateStep, steps) +} + +// TestMaybeTruncateMessagesBudgetTooSmall verifies the degenerate case where even heavy +// truncation can't fit the budget: everything between the first messages and the last message +// is skipped, and the result is still returned. +func TestMaybeTruncateMessagesBudgetTooSmall(t *testing.T) { + messages := make([]*aiv1.CompletionMessage, 10) + for i := range messages { + messages[i] = NewTextCompletionMessage(RoleUser, strings.Repeat("x", 3000)) + } + + result := maybeTruncateMessages(messages, 1) + require.Len(t, result, truncateKeepFirst+1+1) + require.Equal(t, messages[:truncateKeepFirst], result[:truncateKeepFirst]) + require.Equal(t, messages[len(messages)-1], result[len(result)-1]) +} + +// TestMaybeTruncateMessagesUnbalancedToolCalls verifies that tool results whose calls were +// truncated away are also removed. +func TestMaybeTruncateMessagesUnbalancedToolCalls(t *testing.T) { + messages := newTextMessages(121) + // Make the last skipped message a tool call and the first kept message its result + // (which must then be removed as unbalanced). + messages[truncateKeepFirst+truncateStep-1] = &aiv1.CompletionMessage{ + Role: "assistant", + Content: []*aiv1.ContentBlock{ + {BlockType: &aiv1.ContentBlock_ToolCall{ToolCall: &aiv1.ToolCall{Id: "call-1"}}}, + }, + } + messages[truncateKeepFirst+truncateStep] = &aiv1.CompletionMessage{ + Role: "tool", + Content: []*aiv1.ContentBlock{ + {BlockType: &aiv1.ContentBlock_ToolResult{ToolResult: &aiv1.ToolResult{Id: "call-1"}}}, + }, + } + + // A budget that exactly fits everything after the skipped range forces skipping truncateStep messages. + budget := sumEstimatedTokens(messages) - sumEstimatedTokens(messages[truncateKeepFirst:truncateKeepFirst+truncateStep]) + result := maybeTruncateMessages(messages, budget) + require.Len(t, result, truncateKeepFirst+1+len(messages)-truncateKeepFirst-truncateStep-1) + for _, msg := range result { + for _, block := range msg.Content { + if res := block.GetToolResult(); res != nil { + require.NotEqual(t, "call-1", res.Id) + } + } + } +} + +// newTextMessages generates n messages with identical estimated token counts. +func newTextMessages(n int) []*aiv1.CompletionMessage { + messages := make([]*aiv1.CompletionMessage, n) + for i := range messages { + messages[i] = NewTextCompletionMessage(RoleUser, fmt.Sprintf("message %03d", i)) + } + return messages +} + +func sumEstimatedTokens(messages []*aiv1.CompletionMessage) int { + var total int + for _, m := range messages { + total += estimateMessageTokens(m) + } + return total +} diff --git a/runtime/ai/ai_test.go b/runtime/ai/ai_test.go index 9623d61bcb92..2cf1b70058fc 100644 --- a/runtime/ai/ai_test.go +++ b/runtime/ai/ai_test.go @@ -122,6 +122,14 @@ type recordingAIMessage struct { var _ drivers.AIService = &recordingAIService{} +// MaxInputTokens implements drivers.AIService. +func (r *recordingAIService) MaxInputTokens() int { + if r.ai == nil { + return 0 + } + return r.ai.MaxInputTokens() +} + // Complete(ctx context.Context, msgs []*aiv1.CompletionMessage, tools []*aiv1.Tool, outputSchema *jsonschema.Schema) (*aiv1.CompletionMessage, error) func (r *recordingAIService) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { // Create a recorded call diff --git a/runtime/ai/analyst_agent.go b/runtime/ai/analyst_agent.go index 666ab4b569c2..72f92a38e11a 100644 --- a/runtime/ai/analyst_agent.go +++ b/runtime/ai/analyst_agent.go @@ -105,33 +105,37 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An // Determine if it's the first invocation of the agent in this session. first := len(s.Messages(FilterByType(MessageTypeCall), FilterByTool(AnalystAgentName))) == 1 + // Resolve the metrics views tied to the explore/canvas being viewed (if any). + // They are referenced in the context prompt and in the pre-invoked tool calls below. + var metricsViewNames []string + if args.Explore != "" { + _, metricsView, err := t.getValidExploreAndMetricsView(ctx, args.Explore) + if err != nil { + return nil, err + } + metricsViewNames = append(metricsViewNames, metricsView.Meta.Name.Name) + } else if args.Canvas != "" { + _, metricsViews, err := t.getValidCanvasAndMetricsViews(ctx, args.Canvas) + if err != nil { + return nil, err + } + for _, res := range metricsViews { + metricsViewNames = append(metricsViewNames, res.Meta.Name.Name) + } + slices.Sort(metricsViewNames) // Map iteration order is random; sort for a deterministic prompt. + } + // If a specific dashboard is being explored, we pre-invoke some relevant tool calls for that dashboard. // TODO: This uses `first`, but that may not be safe if the user has navigated to another dashboard. We probably need some more sophisticated de-duplication here. - var metricsViewNames []string if first { - if args.Explore != "" { - _, metricsView, err := t.getValidExploreAndMetricsView(ctx, args.Explore) - if err != nil { - return nil, err - } - metricsViewNames = append(metricsViewNames, metricsView.Meta.Name.Name) - } else if args.Canvas != "" { - // Pre-invoke the get_canvas tool to get the canvas definition. + // Pre-invoke the get_canvas tool to get the canvas definition. + if args.Canvas != "" { _, err := s.CallTool(ctx, RoleAssistant, GetCanvasName, nil, &GetCanvasArgs{ Canvas: args.Canvas, }) if err != nil && errors.Is(err, ctx.Err()) { // Don't exit on non-context errors return nil, err } - - _, metricsViews, err := t.getValidCanvasAndMetricsViews(ctx, args.Canvas) - if err != nil { - return nil, err - } - - for _, res := range metricsViews { - metricsViewNames = append(metricsViewNames, res.Meta.Name.Name) - } } // Pre-invoke the query_metrics_view tool for each metrics view tied to the explore/canvas. @@ -150,13 +154,13 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An return nil, err } } - } - // If no specific dashboard is being explored, we pre-invoke the list_metrics_views tool. - if first && len(metricsViewNames) == 0 { - _, err := s.CallTool(ctx, RoleAssistant, ListMetricsViewsName, nil, &ListMetricsViewsArgs{}) - if err != nil && errors.Is(err, ctx.Err()) { // Don't exit on non-context errors - return nil, err + // If no specific dashboard is being explored, we pre-invoke the list_metrics_views tool. + if len(metricsViewNames) == 0 { + _, err := s.CallTool(ctx, RoleAssistant, ListMetricsViewsName, nil, &ListMetricsViewsArgs{}) + if err != nil && errors.Is(err, ctx.Err()) { // Don't exit on non-context errors + return nil, err + } } } @@ -170,8 +174,15 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An tools = append(tools, CreateChartName) } - // Build completion messages - systemPrompt, err := t.systemPrompt(ctx, metricsViewNames, args) + // Build completion messages. + // The system prompt contains only instructions that stay stable for the duration of a session; + // per-turn dynamic info (date, dashboard state, etc.) goes in a separate context message near the + // end, so the LLM prompt cache prefix stays valid across turns. + systemPrompt, err := t.systemPrompt(ctx, args) + if err != nil { + return nil, err + } + contextPrompt, err := t.contextPrompt(ctx, metricsViewNames, args) if err != nil { return nil, err } @@ -190,6 +201,11 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An } } + // Insert the context prompt just before the trailing user prompt. It is regenerated on every + // call and not stored in the session, so keeping it at the end preserves the cached prefix of + // the earlier messages. + messages = slices.Insert(messages, len(messages)-1, NewTextCompletionMessage(RoleUser, contextPrompt)) + // Run an LLM tool call loop var response string err = s.Complete(ctx, "Analyst loop", &response, &CompleteOptions{ @@ -205,20 +221,13 @@ func (t *AnalystAgent) Handler(ctx context.Context, args *AnalystAgentArgs) (*An return &AnalystAgentResult{Response: response}, nil } -func (t *AnalystAgent) systemPrompt(ctx context.Context, metricsViewNames []string, args *AnalystAgentArgs) (string, error) { +func (t *AnalystAgent) systemPrompt(ctx context.Context, args *AnalystAgentArgs) (string, error) { // Prepare template data. - // NOTE: All the template properties are optional and may be empty. + // NOTE: The system prompt must stay byte-identical across turns in a session so the LLM prompt + // cache prefix remains valid. Only use template properties that are stable for the duration of + // a session; per-turn dynamic values belong in the context prompt (see contextPrompt). session := GetSession(ctx) - instance, err := t.Runtime.Instance(ctx, session.InstanceID()) - if err != nil { - return "", fmt.Errorf("failed to get instance: %w", err) - } - instanceCfg, err := instance.Config() - if err != nil { - return "", fmt.Errorf("failed to get instance config: %w", err) - } - ff, err := t.Runtime.FeatureFlags(ctx, session.InstanceID(), session.Claims()) if err != nil { return "", fmt.Errorf("failed to get feature flags: %w", err) @@ -228,63 +237,13 @@ func (t *AnalystAgent) systemPrompt(ctx context.Context, metricsViewNames []stri ff["chat_charts"] = false } - metricsViewsQuoted := make([]string, len(metricsViewNames)) - for i, mv := range metricsViewNames { - metricsViewsQuoted[i] = fmt.Sprintf("`%s`", mv) - } - - dimensionsQuoted := make([]string, len(args.Dimensions)) - for i, dim := range args.Dimensions { - dimensionsQuoted[i] = fmt.Sprintf("`%s`", dim) - } - - measuresQuoted := make([]string, len(args.Measures)) - for i, measure := range args.Measures { - measuresQuoted[i] = fmt.Sprintf("`%s`", measure) - } - data := map[string]any{ - "ai_instructions": session.ProjectInstructions(), - "is_prompt": args.Prompt != "", - "metrics_views": strings.Join(metricsViewsQuoted, ", "), - "explore": args.Explore, - "canvas": args.Canvas, - "canvas_component": args.CanvasComponent, - "dimensions": strings.Join(dimensionsQuoted, ", "), - "measures": strings.Join(measuresQuoted, ", "), - "feature_flags": ff, - "forked": session.Forked(), - "is_report": args.IsReport, - "now": time.Now(), - "max_query_limit": instanceCfg.AIMaxQueryLimit, - } - - if !args.TimeStart.IsZero() && !args.TimeEnd.IsZero() { - data["time_start"] = args.TimeStart.Format(time.RFC3339) - data["time_end"] = args.TimeEnd.Format(time.RFC3339) - } - - if !args.ComparisonTimeStart.IsZero() && !args.ComparisonTimeEnd.IsZero() { - data["comparison_start"] = args.ComparisonTimeStart.Format(time.RFC3339) - data["comparison_end"] = args.ComparisonTimeEnd.Format(time.RFC3339) - } - - if args.Where != nil { - data["where"], err = metricsview.ExpressionToSQL(args.Where) - if err != nil { - return "", err - } - } - - if args.WherePerMetricsView != nil { - wherePerMetricsView := map[string]string{} - for metricsViewName, whereExpr := range args.WherePerMetricsView { - wherePerMetricsView[metricsViewName], err = metricsview.ExpressionToSQL(whereExpr) - if err != nil { - return "", err - } - } - data["where_per_metrics_view"] = wherePerMetricsView + "is_prompt": args.Prompt != "", + "feature_flags": ff, + "forked": session.Forked(), + "is_report": args.IsReport, + "has_comparison": !args.ComparisonTimeStart.IsZero() && !args.ComparisonTimeEnd.IsZero(), + "ai_instructions": session.ProjectInstructions(), // Set once at session creation, so stable for the session } // Generate the system prompt @@ -295,8 +254,6 @@ You systematically explore data using available metrics tools, then apply analyt You are operating in an automated scheduled insight report mode where you will come up with insights on your own without additional user input. {{ if .is_prompt }}The user has provided a custom prompt for this scheduled insight report. Tailor your analysis to address this prompt specifically. {{ end }} {{ end }} - -Today's date is {{ .now.Format "Monday, January 2, 2006" }} ({{ .now.Format "2006-01-02" }}). @@ -307,43 +264,17 @@ Today's date is {{ .now.Format "Monday, January 2, 2006" }} ({{ .now.Format "200 **Phase 1: discovery (setup)** -{{ if .explore }} -Your goal is to analyze the contents of the dashboard "{{ .explore }}", which is powered by the metrics view(s) {{ .metrics_views }}. -{{ if not .is_report }}The user is actively viewing this dashboard, and it's what you they refer to if they use expressions like "this dashboard", "the current view", etc. {{ end }} -The metrics view's definition and time range of available data has been provided in your tool calls. +A user message tagged provides details for the current request, such as today's date and, if the user is viewing a dashboard, the dashboard's metrics views and currently applied settings. -Here is an overview of the settings applied to the dashboard: -{{ if (and .time_start .time_end) }}Use time range: start={{.time_start}}, end={{.time_end}}{{ end }} -{{ if (and .comparison_start .comparison_end) }}Use comparison time range: start={{.comparison_start}}, end={{.comparison_end}}{{ end }} -{{ if .where }}Use where filters: "{{ .where }}"{{ end }} -{{ if .measures }}Use measures: {{ .measures }}{{ end }} -{{ if .dimensions }}Use dimensions: {{ .dimensions }}{{ end }} - -You should: -1. Carefully study the metrics view definition to understand the measures and dimensions available for analysis. +If the context describes a dashboard: +1. Carefully study the metrics view definition(s) provided in your tool calls to understand the measures and dimensions available for analysis. 2. Remember the time range of available data and use it to inform and filter your queries. -{{ else if .canvas }} -Your goal is to analyze the contents of the canvas "{{ .canvas }}", which is powered by the metrics view(s) {{ .metrics_views }}. -The user is actively viewing this dashboard, and it's what you they refer to if they use expressions like "this dashboard", "the current view", etc. -The metrics views and canvas definitions have been provided in your tool calls. - -Here is an overview of the settings the user has currently applied to the dashboard (Merge component's dimension_filters with "and"): -{{ if (and .time_start .time_end) }}Use time range: start={{.time_start}}, end={{.time_end}}{{ end }} -{{ if .where_per_metrics_view }}{{range $mv, $filter := .where_per_metrics_view}}Use where filters for metrics view "{{ $mv }}": "{{ $filter }}" -{{end}}{{ end }} +3. Apply the dashboard settings (time range, filters, measures, dimensions) listed in the context to your queries. -You should: -1. Carefully study the canvas and metrics view definition to understand the measures and dimensions available for analysis. -2. Remember the time range of available data and use it to inform and filter your queries. -{{ if .canvas_component }} -The user is looking at "{{ .canvas_component }}". Pay special attention to its definition and filters and use it to inform your analysis. -{{ end }} -{{ else }} -Follow these steps in order: +If no dashboard context is provided, follow these steps in order: 1. **Discover**: Use "list_metrics_views" to identify available datasets -2. **Understand**: Use "get_metrics_view" to understand measures and dimensions for the selected view +2. **Understand**: Use "get_metrics_view" to understand measures and dimensions for the selected view 3. **Scope**: Use "query_metrics_view_summary" to determine the span of available data -{{ end }} {{ if .forked }} Important instructions regarding access permissions: @@ -358,7 +289,7 @@ Execute a MINIMUM of 4-6 distinct analytical queries, building each query based Continue until you have sufficient insights for comprehensive analysis. Some analyses may require up to 20 queries. {{ if and .is_report (not .is_prompt) }} -{{ if (and .comparison_start .comparison_end) }} +{{ if .has_comparison }} You are doing comparative analysis between two time periods in scheduled insight report mode, your analysis should: 1. Compare current period to the comparison period for all key measures @@ -447,6 +378,13 @@ If a question seems unrelated, first inspect the available metrics views to see Decline to engage if the topic is clearly outside the scope of the data (e.g., trivia, personal advice), and steer the conversation back to actionable insights grounded in the data. +{{ if .ai_instructions }} + +The administrator has provided the following project-wide instructions, which may or may not be relevant to this task: +{{ .ai_instructions }} + +{{ end }} + After each query in Phase 2, think through: - What patterns or anomalies did this reveal? @@ -482,15 +420,113 @@ Based on the data analysis, here are the key insights: - Use descriptive text in sentence case (e.g. "This suggests Android is valuable ([Device breakdown](url))." or "Revenue increased 25%% ([Revenue by country](url)).") - When one paragraph contains multiple insights from the same query, cite once at the end of the paragraph +`, data) +} - -The system allows a max row limit of {{ .max_query_limit }} per query. +// contextPrompt generates a user message with dynamic context for the current request, such as +// today's date and the state of the dashboard the user is viewing. It is kept separate from the +// system prompt and inserted near the end of the messages, so that changes to it between turns +// don't invalidate the LLM prompt cache for the preceding messages. +func (t *AnalystAgent) contextPrompt(ctx context.Context, metricsViewNames []string, args *AnalystAgentArgs) (string, error) { + // Prepare template data. + // NOTE: All the template properties are optional and may be empty. + session := GetSession(ctx) -{{ if .ai_instructions }} -The administrator has provided the following project-wide instructions, which may or may not be relevant to this task: -{{ .ai_instructions }} + instance, err := t.Runtime.Instance(ctx, session.InstanceID()) + if err != nil { + return "", fmt.Errorf("failed to get instance: %w", err) + } + instanceCfg, err := instance.Config() + if err != nil { + return "", fmt.Errorf("failed to get instance config: %w", err) + } + + metricsViewsQuoted := make([]string, len(metricsViewNames)) + for i, mv := range metricsViewNames { + metricsViewsQuoted[i] = fmt.Sprintf("`%s`", mv) + } + + dimensionsQuoted := make([]string, len(args.Dimensions)) + for i, dim := range args.Dimensions { + dimensionsQuoted[i] = fmt.Sprintf("`%s`", dim) + } + + measuresQuoted := make([]string, len(args.Measures)) + for i, measure := range args.Measures { + measuresQuoted[i] = fmt.Sprintf("`%s`", measure) + } + + data := map[string]any{ + "metrics_views": strings.Join(metricsViewsQuoted, ", "), + "explore": args.Explore, + "canvas": args.Canvas, + "canvas_component": args.CanvasComponent, + "dimensions": strings.Join(dimensionsQuoted, ", "), + "measures": strings.Join(measuresQuoted, ", "), + "is_report": args.IsReport, + "now": time.Now(), + "max_query_limit": instanceCfg.AIMaxQueryLimit, + } + + if !args.TimeStart.IsZero() && !args.TimeEnd.IsZero() { + data["time_start"] = args.TimeStart.Format(time.RFC3339) + data["time_end"] = args.TimeEnd.Format(time.RFC3339) + } + + if !args.ComparisonTimeStart.IsZero() && !args.ComparisonTimeEnd.IsZero() { + data["comparison_start"] = args.ComparisonTimeStart.Format(time.RFC3339) + data["comparison_end"] = args.ComparisonTimeEnd.Format(time.RFC3339) + } + + if args.Where != nil { + data["where"], err = metricsview.ExpressionToSQL(args.Where) + if err != nil { + return "", err + } + } + + if args.WherePerMetricsView != nil { + wherePerMetricsView := map[string]string{} + for metricsViewName, whereExpr := range args.WherePerMetricsView { + wherePerMetricsView[metricsViewName], err = metricsview.ExpressionToSQL(whereExpr) + if err != nil { + return "", err + } + } + data["where_per_metrics_view"] = wherePerMetricsView + } + + // Generate the context prompt + return executeTemplate(` +Today's date is {{ .now.Format "Monday, January 2, 2006" }} ({{ .now.Format "2006-01-02" }}). + +{{ if .explore }} +Your goal is to analyze the contents of the dashboard "{{ .explore }}", which is powered by the metrics view(s) {{ .metrics_views }}. +{{ if not .is_report }}The user is actively viewing this dashboard, and it's what they refer to if they use expressions like "this dashboard", "the current view", etc. {{ end }} +The metrics view's definition and time range of available data has been provided in your tool calls. + +Here is an overview of the settings applied to the dashboard: +{{ if (and .time_start .time_end) }}Use time range: start={{.time_start}}, end={{.time_end}}{{ end }} +{{ if (and .comparison_start .comparison_end) }}Use comparison time range: start={{.comparison_start}}, end={{.comparison_end}}{{ end }} +{{ if .where }}Use where filters: "{{ .where }}"{{ end }} +{{ if .measures }}Use measures: {{ .measures }}{{ end }} +{{ if .dimensions }}Use dimensions: {{ .dimensions }}{{ end }} +{{ else if .canvas }} +Your goal is to analyze the contents of the canvas "{{ .canvas }}", which is powered by the metrics view(s) {{ .metrics_views }}. +The user is actively viewing this dashboard, and it's what they refer to if they use expressions like "this dashboard", "the current view", etc. +The metrics views and canvas definitions have been provided in your tool calls. + +Here is an overview of the settings the user has currently applied to the dashboard (Merge component's dimension_filters with "and"): +{{ if (and .time_start .time_end) }}Use time range: start={{.time_start}}, end={{.time_end}}{{ end }} +{{ if .where_per_metrics_view }}{{range $mv, $filter := .where_per_metrics_view}}Use where filters for metrics view "{{ $mv }}": "{{ $filter }}" +{{end}}{{ end }} +{{ if .canvas_component }} +The user is looking at "{{ .canvas_component }}". Pay special attention to its definition and filters and use it to inform your analysis. {{ end }} - +{{ end }} + +The system allows a max row limit of {{ .max_query_limit }} per query. + `, data) } diff --git a/runtime/drivers/admin/admin.go b/runtime/drivers/admin/admin.go index 04fd766e5721..3ec449509271 100644 --- a/runtime/drivers/admin/admin.go +++ b/runtime/drivers/admin/admin.go @@ -31,6 +31,11 @@ var spec = drivers.Spec{ Required: true, Secret: true, }, + { + Key: "max_input_tokens", + Type: drivers.NumberPropertyType, + Required: false, + }, }, } @@ -46,6 +51,8 @@ type configProperties struct { AdminURL string `mapstructure:"admin_url"` AccessToken string `mapstructure:"access_token"` ProjectID string `mapstructure:"project_id"` + // MaxInputTokens is the input token limit for AI completion requests proxied to the admin service. + MaxInputTokens int `mapstructure:"max_input_tokens"` } func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { diff --git a/runtime/drivers/admin/ai.go b/runtime/drivers/admin/ai.go index a587e278f82d..1e3fdb3d7f00 100644 --- a/runtime/drivers/admin/ai.go +++ b/runtime/drivers/admin/ai.go @@ -8,6 +8,16 @@ import ( "github.com/rilldata/rill/runtime/drivers" ) +// MaxInputTokens implements drivers.AIService. +// The admin service proxies to a Rill-managed provider whose configuration isn't known here, so +// this defaults to a conservative limit that fits all providers currently used for managed AI. +func (h *Handle) MaxInputTokens() int { + if h.config.MaxInputTokens > 0 { + return h.config.MaxInputTokens + } + return drivers.DefaultAIMaxInputTokens +} + func (h *Handle) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { var outputJSONSchema string if opts.OutputSchema != nil { @@ -22,6 +32,7 @@ func (h *Handle) Complete(ctx context.Context, opts *drivers.CompleteOptions) (* Messages: opts.Messages, Tools: opts.Tools, OutputJsonSchema: outputJSONSchema, + CacheKey: opts.CacheKey, }) if err != nil { return nil, err diff --git a/runtime/drivers/ai.go b/runtime/drivers/ai.go index d595386cef3a..903b3a41d8cc 100644 --- a/runtime/drivers/ai.go +++ b/runtime/drivers/ai.go @@ -7,14 +7,26 @@ import ( aiv1 "github.com/rilldata/rill/proto/gen/rill/ai/v1" ) +// DefaultAIMaxInputTokens is a conservative default input token limit for AI completion requests. +// It matches the smallest context window among the commonly used OpenAI, Gemini, and Claude models +// (Claude's 200k; GPT-5.x allows 272k and Gemini ~1M). Drivers can override it via their +// max_input_tokens config property. +const DefaultAIMaxInputTokens = 200_000 + type AIService interface { Complete(ctx context.Context, opts *CompleteOptions) (*CompleteResult, error) + // MaxInputTokens returns the maximum number of input tokens a completion request may contain. + // It is configurable per connector via the max_input_tokens property and defaults to DefaultAIMaxInputTokens. + MaxInputTokens() int } type CompleteOptions struct { Messages []*aiv1.CompletionMessage Tools []*aiv1.Tool OutputSchema *jsonschema.Schema + // CacheKey identifies a series of requests that share a prompt prefix (e.g. an AI session ID). + // Providers may use it to improve prompt cache routing. Drivers that don't support it ignore it. + CacheKey string } type CompleteResult struct { diff --git a/runtime/drivers/claude/claude.go b/runtime/drivers/claude/claude.go index 4b56eb5dd654..a1e7fcf7a0a2 100644 --- a/runtime/drivers/claude/claude.go +++ b/runtime/drivers/claude/claude.go @@ -53,6 +53,13 @@ var spec = drivers.Spec{ Description: "Maximum number of tokens in the response.", Default: "8192", }, + { + Key: "max_input_tokens", + Type: drivers.NumberPropertyType, + Required: false, + DisplayName: "Max Input Tokens", + Description: "Maximum number of input tokens a completion request may contain (default: 200000).", + }, { Key: "temperature", Type: drivers.NumberPropertyType, @@ -120,11 +127,12 @@ func (d driver) TertiarySourceConnectors(ctx context.Context, srcProps map[strin } type configProperties struct { - APIKey string `mapstructure:"api_key"` - Model string `mapstructure:"model"` - MaxTokens int `mapstructure:"max_tokens"` - Temperature *float64 `mapstructure:"temperature"` - BaseURL string `mapstructure:"base_url"` + APIKey string `mapstructure:"api_key"` + Model string `mapstructure:"model"` + MaxTokens int `mapstructure:"max_tokens"` + MaxInputTokens int `mapstructure:"max_input_tokens"` + Temperature *float64 `mapstructure:"temperature"` + BaseURL string `mapstructure:"base_url"` } func (c *configProperties) getModel() string { @@ -252,6 +260,14 @@ func (h *handle) Ping(ctx context.Context) error { return nil } +// MaxInputTokens implements drivers.AIService. +func (h *handle) MaxInputTokens() int { + if h.config.MaxInputTokens > 0 { + return h.config.MaxInputTokens + } + return drivers.DefaultAIMaxInputTokens +} + // Complete implements drivers.AIService. func (h *handle) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { system, msgs, err := convertMessages(opts.Messages) diff --git a/runtime/drivers/gemini/gemini.go b/runtime/drivers/gemini/gemini.go index 21cce1d5b032..4a7eee0dc113 100644 --- a/runtime/drivers/gemini/gemini.go +++ b/runtime/drivers/gemini/gemini.go @@ -68,6 +68,13 @@ var spec = drivers.Spec{ DisplayName: "Max Output Tokens", Description: "Maximum number of tokens in the response.", }, + { + Key: "max_input_tokens", + Type: drivers.NumberPropertyType, + Required: false, + DisplayName: "Max Input Tokens", + Description: "Maximum number of input tokens a completion request may contain (default: 200000).", + }, { Key: "temperature", Type: drivers.NumberPropertyType, @@ -143,6 +150,7 @@ type configProperties struct { IncludeThoughts bool `mapstructure:"include_thoughts"` ThinkingLevel string `mapstructure:"thinking_level"` MaxOutputTokens int `mapstructure:"max_output_tokens"` + MaxInputTokens int `mapstructure:"max_input_tokens"` Temperature *float64 `mapstructure:"temperature"` TopP float64 `mapstructure:"top_p"` TopK float64 `mapstructure:"top_k"` @@ -277,6 +285,14 @@ func (h *handle) Ping(ctx context.Context) error { return nil } +// MaxInputTokens implements drivers.AIService. +func (h *handle) MaxInputTokens() int { + if h.config.MaxInputTokens > 0 { + return h.config.MaxInputTokens + } + return drivers.DefaultAIMaxInputTokens +} + // Complete implements drivers.AIService. func (h *handle) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { // Convert Rill messages to Gemini format, extracting system instruction separately diff --git a/runtime/drivers/mock/ai/ai.go b/runtime/drivers/mock/ai/ai.go index 6d976710f8a0..9bbb1b91d73d 100644 --- a/runtime/drivers/mock/ai/ai.go +++ b/runtime/drivers/mock/ai/ai.go @@ -161,6 +161,11 @@ func (c *connection) AsNotifier(properties map[string]any) (drivers.Notifier, er } // Complete implements drivers.AIService. +// MaxInputTokens implements drivers.AIService. +func (c *connection) MaxInputTokens() int { + return drivers.DefaultAIMaxInputTokens +} + func (c *connection) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { if c.toolCallingMode { return &drivers.CompleteResult{ diff --git a/runtime/drivers/openai/openai.go b/runtime/drivers/openai/openai.go index 2a5d2691e665..dfee567e01df 100644 --- a/runtime/drivers/openai/openai.go +++ b/runtime/drivers/openai/openai.go @@ -55,6 +55,13 @@ var spec = drivers.Spec{ DisplayName: "Max Output Tokens", Description: "Maximum number of tokens to generate in the completion (default: 8192).", }, + { + Key: "max_input_tokens", + Type: drivers.NumberPropertyType, + Required: false, + DisplayName: "Max Input Tokens", + Description: "Maximum number of input tokens a completion request may contain (default: 200000).", + }, { Key: "reasoning_effort", Type: drivers.StringPropertyType, @@ -154,6 +161,7 @@ type configProperties struct { APIKey string `mapstructure:"api_key"` Model string `mapstructure:"model"` MaxOutputTokens int64 `mapstructure:"max_output_tokens"` + MaxInputTokens int `mapstructure:"max_input_tokens"` ReasoningEffort string `mapstructure:"reasoning_effort"` Temperature *float64 `mapstructure:"temperature"` BaseURL string `mapstructure:"base_url"` @@ -175,6 +183,13 @@ func (c *configProperties) getMaxOutputTokens() int64 { return 8192 } +func (c *configProperties) getMaxInputTokens() int { + if c.MaxInputTokens > 0 { + return c.MaxInputTokens + } + return drivers.DefaultAIMaxInputTokens +} + func (c *configProperties) getTemperature() *float64 { if c.Temperature != nil { return c.Temperature @@ -291,6 +306,11 @@ func (o *openaiHandle) Ping(ctx context.Context) error { return nil } +// MaxInputTokens implements drivers.AIService. +func (o *openaiHandle) MaxInputTokens() int { + return o.config.getMaxInputTokens() +} + // Complete implements drivers.AIService. func (o *openaiHandle) Complete(ctx context.Context, opts *drivers.CompleteOptions) (*drivers.CompleteResult, error) { // Convert Rill messages to OpenAI's message format @@ -326,6 +346,10 @@ func (o *openaiHandle) Complete(ctx context.Context, opts *drivers.CompleteOptio if o.config.ReasoningEffort != "" { params.ReasoningEffort = shared.ReasoningEffort(o.config.ReasoningEffort) } + // A stable per-session cache key improves prompt cache routing when many sessions share the same prompt prefix. + if opts.CacheKey != "" { + params.PromptCacheKey = openai.String(opts.CacheKey) + } // Set response format based on output schema if opts.OutputSchema != nil { diff --git a/web-admin/src/client/gen/index.schemas.ts b/web-admin/src/client/gen/index.schemas.ts index e66c609c5a62..8f4e9ac917df 100644 --- a/web-admin/src/client/gen/index.schemas.ts +++ b/web-admin/src/client/gen/index.schemas.ts @@ -305,6 +305,9 @@ export interface V1CompleteRequest { messages?: V1CompletionMessage[]; tools?: V1Tool[]; outputJsonSchema?: string; + /** Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). +Providers may use it to improve prompt cache routing. */ + cacheKey?: string; } export interface V1CompleteResponse { diff --git a/web-common/src/proto/gen/rill/admin/v1/ai_pb.ts b/web-common/src/proto/gen/rill/admin/v1/ai_pb.ts index 4f4d4c7d2be6..b05586184985 100644 --- a/web-common/src/proto/gen/rill/admin/v1/ai_pb.ts +++ b/web-common/src/proto/gen/rill/admin/v1/ai_pb.ts @@ -32,6 +32,14 @@ export class CompleteRequest extends Message { */ outputJsonSchema = ""; + /** + * Optional key identifying a series of requests that share a prompt prefix (e.g. an AI session ID). + * Providers may use it to improve prompt cache routing. + * + * @generated from field: string cache_key = 4; + */ + cacheKey = ""; + constructor(data?: PartialMessage) { super(); proto3.util.initPartial(data, this); @@ -43,6 +51,7 @@ export class CompleteRequest extends Message { { no: 1, name: "messages", kind: "message", T: CompletionMessage, repeated: true }, { no: 2, name: "tools", kind: "message", T: Tool, repeated: true }, { no: 3, name: "output_json_schema", kind: "scalar", T: 9 /* ScalarType.STRING */ }, + { no: 4, name: "cache_key", kind: "scalar", T: 9 /* ScalarType.STRING */ }, ]); static fromBinary(bytes: Uint8Array, options?: Partial): CompleteRequest {