1414 * limitations under the License.
1515 */
1616
17+ import { stripUndefinedProps } from '@genkit-ai/core' ;
1718import * as assert from 'assert' ;
1819import { beforeEach , describe , it } from 'node:test' ;
1920import { GenkitBeta , genkit } from '../src/beta' ;
@@ -24,14 +25,11 @@ describe('formats', () => {
2425
2526 beforeEach ( ( ) => {
2627 ai = genkit ( { } ) ;
27- defineEchoModel ( ai ) ;
28- } ) ;
29-
30- it ( 'lets you define and use a custom output format' , async ( ) => {
3128 ai . defineFormat (
3229 {
3330 name : 'banana' ,
3431 format : 'banana' ,
32+ constrained : true ,
3533 } ,
3634 ( schema ) => {
3735 let instructions : string | undefined ;
@@ -53,6 +51,10 @@ describe('formats', () => {
5351 } ;
5452 }
5553 ) ;
54+ } ) ;
55+
56+ it ( 'lets you define and use a custom output format with native constrained generation' , async ( ) => {
57+ defineEchoModel ( ai , { supports : { constrained : 'all' } } ) ;
5658
5759 const { output } = await ai . generate ( {
5860 model : 'echoModel' ,
@@ -73,5 +75,110 @@ describe('formats', () => {
7375 }
7476 assert . deepStrictEqual ( chunks , [ 'banana: 3' , 'banana: 2' , 'banana: 1' ] ) ;
7577 assert . strictEqual ( ( await response ) . output , 'banana: Echo: hi' ) ;
78+ assert . deepStrictEqual ( stripUndefinedProps ( ( await response ) . request ) , {
79+ config : { } ,
80+ messages : [
81+ {
82+ content : [ { text : 'hi' } ] ,
83+ role : 'user' ,
84+ } ,
85+ ] ,
86+ output : {
87+ constrained : true ,
88+ format : 'banana' ,
89+ schema : { } ,
90+ } ,
91+ tools : [ ] ,
92+ } ) ;
93+ } ) ;
94+
95+ it ( 'lets you define and use a custom output format with simulated constrained generation' , async ( ) => {
96+ defineEchoModel ( ai , { supports : { constrained : false } } ) ;
97+
98+ const { output } = await ai . generate ( {
99+ model : 'echoModel' ,
100+ prompt : 'hi' ,
101+ output : { format : 'banana' } ,
102+ } ) ;
103+
104+ assert . strictEqual (
105+ output ,
106+ 'banana: Echo: hi,Output should be in JSON format and conform to the following schema:\n' +
107+ '\n' +
108+ '```\n' +
109+ '{}\n' +
110+ '```\n'
111+ ) ;
112+
113+ const { response, stream } = await ai . generateStream ( {
114+ model : 'echoModel' ,
115+ prompt : 'hi' ,
116+ output : { format : 'banana' } ,
117+ } ) ;
118+ const chunks : string [ ] = [ ] ;
119+ for await ( const chunk of stream ) {
120+ chunks . push ( `${ chunk . output } ) ;
121+ }
122+ assert . deepStrictEqual ( chunks , [ 'banana: 3' , 'banana: 2' , 'banana: 1' ] ) ;
123+ assert . strictEqual (
124+ ( await response ) . output ,
125+ 'banana: Echo: hi,Output should be in JSON format and conform to the following schema:\n' +
126+ '\n' +
127+ '```\n' +
128+ '{}\n' +
129+ '```\n'
130+ ) ;
131+ assert . deepStrictEqual ( stripUndefinedProps ( ( await response ) . request ) , {
132+ config : { } ,
133+ messages : [
134+ {
135+ content : [ { text : 'hi' } ] ,
136+ role : 'user' ,
137+ } ,
138+ ] ,
139+ output : {
140+ constrained : true ,
141+ format : 'banana' ,
142+ schema : { } ,
143+ } ,
144+ tools : [ ] ,
145+ } ) ;
146+ } ) ;
147+
148+ it ( 'used explicitly specified output options overriding format options' , async ( ) => {
149+ defineEchoModel ( ai , { supports : { constrained : 'all' } } ) ;
150+ const response = await ai . generate ( {
151+ model : 'echoModel' ,
152+ prompt : 'hi' ,
153+ output : {
154+ format : 'banana' ,
155+ // Explicitly specified, should ignore whatever format sets
156+ constrained : false ,
157+ jsonSchema : { type : 'string' } ,
158+ } ,
159+ } ) ;
160+ assert . deepStrictEqual ( stripUndefinedProps ( response . request ) , {
161+ config : { } ,
162+ messages : [
163+ {
164+ content : [
165+ { text : 'hi' } ,
166+ {
167+ text : 'Output should be in banana format' ,
168+ metadata : { purpose : 'output' } ,
169+ } ,
170+ ] ,
171+ role : 'user' ,
172+ } ,
173+ ] ,
174+ output : {
175+ constrained : false ,
176+ format : 'banana' ,
177+ schema : {
178+ type : 'string' ,
179+ } ,
180+ } ,
181+ tools : [ ] ,
182+ } ) ;
76183 } ) ;
77184} ) ;
0 commit comments