Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(js/ai): make updateState patch state instead of replace it #1573

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions js/ai/src/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -85,16 +85,22 @@ export class Session<S = any> {
}

/**
* Update session state data.
* Update session state data by patching the existing state.
* @param data Partial state update that will be merged with existing state
*/
async updateState(data: S): Promise<void> {
async updateState(data: Partial<S>): Promise<void> {
let sessionData = this.sessionData;
if (!sessionData) {
sessionData = {} as SessionData<S>;
}
sessionData.state = data;
this.sessionData = sessionData;

// Merge the new data with existing state
sessionData.state = {
...sessionData.state,
...data,
} as S;

this.sessionData = sessionData;
await this.store.save(this.id, sessionData);
}

Expand Down
65 changes: 65 additions & 0 deletions js/genkit/tests/chat_test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -660,3 +660,68 @@ describe('preamble', () => {
]);
});
});

describe('session state', () => {
let ai: Genkit;

beforeEach(() => {
ai = genkit({
model: 'echoModel',
});
defineEchoModel(ai);
});

it('properly patches partial state updates', async () => {
const session = ai.createSession({
initialState: { userName: 'John', color: 'Blue', age: 25 },
});

// Update single property
await session.updateState({ color: 'Green' });
assert.deepStrictEqual(
session.state,
{ userName: 'John', color: 'Green', age: 25 },
'should preserve existing properties when updating a single field'
);

// Update multiple properties
await session.updateState({ userName: 'Jane', color: 'red', age: 26 });
assert.deepStrictEqual(
session.state,
{ userName: 'Jane', color: 'red', age: 26 },
'should preserve non-updated properties when updating multiple fields'
);

// Update with undefined state
const emptySession = ai.createSession();
await emptySession.updateState({ newProp: 'value' });
assert.deepStrictEqual(
emptySession.state,
{ newProp: 'value' },
'should handle updates when initial state is undefined'
);
});

it('maintains state across chat messages', async () => {
const session = ai.createSession({
initialState: { userName: 'Pavel' },
});

const chat = session.chat();
await chat.send('hi');

assert.deepStrictEqual(
session.state,
{ userName: 'Pavel' },
'should preserve state between chat messages'
);
await session.updateState({ userName: 'Jacob' });

await chat.send('bye');
assert.deepStrictEqual(
session.state,
{ userName: 'Jacob' },
'state should remain stable after chat messages'
);
});
});
Loading