feat(aiagent): handle custom memory_id (#7432)

* first draft

* fix

* nit

* cleaning

* fix on deploy if not set

* add off mode to memory
This commit is contained in:
centdix
2025-12-23 14:47:50 +01:00
committed by GitHub
parent 327c997dd0
commit 117bfdd4c2
11 changed files with 146 additions and 50 deletions
+14 -1
View File
@@ -2,6 +2,7 @@ use crate::ai::providers::openai::OpenAIToolCall;
use serde::{Deserialize, Serialize};
use serde_json::value::RawValue;
use std::collections::HashMap;
use uuid::Uuid;
use windmill_common::mcp_client::McpToolSource;
use windmill_common::{
ai_providers::AIProvider, db::DB, error::Error, flow_status::AgentAction, flows::FlowModule,
@@ -121,9 +122,12 @@ impl Default for OutputType {
#[derive(Deserialize, Debug, Clone)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub enum Memory {
Off,
Auto {
#[serde(default)]
context_length: usize,
#[serde(default)]
memory_id: Option<Uuid>,
},
Manual {
messages: Vec<OpenAIMessage>,
@@ -168,7 +172,16 @@ impl From<AIAgentArgsRaw> for AIAgentArgs {
// Backward compatibility: if messages_context_length is set, use auto mode
let memory = raw.memory.or_else(|| {
raw.messages_context_length
.map(|context_length| Memory::Auto { context_length })
.map(|context_length| Memory::Auto { context_length, memory_id: None })
});
// Backward compatibility: if context_length is 0, use off mode
let memory = memory.map(|memory| {
if let Memory::Auto { context_length: 0, .. } = memory {
Memory::Off
} else {
memory
}
});
AIAgentArgs {
+38 -38
View File
@@ -424,6 +424,19 @@ pub async fn run_agent(
let is_text_output = output_type == &OutputType::Text;
// Flow-level memory_id (from chat mode) takes precedence over step-level memory_id
let memory_id = flow_context
.flow_status
.as_ref()
.and_then(|fs| fs.memory_id)
.or_else(|| {
// Extract memory_id from Memory::Auto if present
match &args.memory {
Some(Memory::Auto { memory_id, .. }) => *memory_id,
_ => None,
}
});
// Load messages based on history mode
if matches!(output_type, OutputType::Text) {
match &args.memory {
@@ -433,14 +446,10 @@ pub async fn run_agent(
messages.extend(manual_messages.clone());
}
}
Some(Memory::Auto { context_length }) if *context_length > 0 => {
Some(Memory::Auto { context_length, .. }) => {
// Auto mode: load from memory
if let Some(step_id) = job.flow_step_id.as_deref() {
if let Some(memory_id) = flow_context
.flow_status
.as_ref()
.and_then(|fs| fs.memory_id)
{
if let Some(memory_id) = memory_id {
// Read messages from memory
match read_from_memory(db, &job.workspace_id, memory_id, step_id).await {
Ok(Some(loaded_messages)) => {
@@ -608,11 +617,6 @@ pub async fn run_agent(
.and_then(|fs| fs.chat_input_enabled)
.unwrap_or(false);
let memory_id = flow_context
.flow_status
.as_ref()
.and_then(|fs| fs.memory_id);
let step_name = get_step_name_from_flow(summary.as_deref(), job.flow_step_id.as_deref());
let max_iterations = args
@@ -978,37 +982,33 @@ pub async fn run_agent(
// Skip memory persistence if using manual messages (bypass memory entirely)
// final_messages contains the complete history (old messages + new ones)
if matches!(output_type, OutputType::Text) && !use_manual_messages {
if let Some(Memory::Auto { context_length }) = &args.memory {
if *context_length > 0 {
if let Some(step_id) = job.flow_step_id.as_deref() {
// Extract OpenAIMessages from final_messages
let all_messages: Vec<OpenAIMessage> =
final_messages.iter().map(|m| m.message.clone()).collect();
if let Some(Memory::Auto { context_length, .. }) = &args.memory {
if let Some(step_id) = job.flow_step_id.as_deref() {
// Extract OpenAIMessages from final_messages
let all_messages: Vec<OpenAIMessage> =
final_messages.iter().map(|m| m.message.clone()).collect();
if !all_messages.is_empty() {
// Keep only the last n messages
let start_idx = all_messages.len().saturating_sub(*context_length);
let messages_to_persist = all_messages[start_idx..].to_vec();
if !all_messages.is_empty() {
// Keep only the last n messages
let start_idx = all_messages.len().saturating_sub(*context_length);
let messages_to_persist = all_messages[start_idx..].to_vec();
if let Some(memory_id) =
flow_context.flow_status.and_then(|fs| fs.memory_id)
if let Some(memory_id) = memory_id {
if let Err(e) = write_to_memory(
db,
&job.workspace_id,
memory_id,
step_id,
&messages_to_persist,
)
.await
{
if let Err(e) = write_to_memory(
db,
&job.workspace_id,
memory_id,
tracing::error!(
"Failed to persist {} messages to memory for step {}: {}",
messages_to_persist.len(),
step_id,
&messages_to_persist,
)
.await
{
tracing::error!(
"Failed to persist {} messages to memory for step {}: {}",
messages_to_persist.len(),
step_id,
e
);
}
e
);
}
}
}
+2
View File
@@ -42,6 +42,7 @@ export interface SchemaProperty {
properties?: { [name: string]: SchemaProperty }
required?: string[]
showExpr?: string
hideWhenChatEnabled?: boolean
password?: boolean
order?: string[]
nullable?: boolean
@@ -52,6 +53,7 @@ export interface SchemaProperty {
originalType?: string
disabled?: boolean
'x-no-s3-storage-workspace-warning'?: string
'x-auto-generate'?: boolean
}
export interface ModalSchemaProperty {
+5 -1
View File
@@ -47,6 +47,7 @@
import AIProviderPicker from './AIProviderPicker.svelte'
import TextInput from './text_input/TextInput.svelte'
import FileInput from './common/fileInput/FileInput.svelte'
import { randomUUID } from './flows/conversations/FlowChatManager.svelte'
interface Props {
label?: string
@@ -122,6 +123,7 @@
| undefined
workspace?: string | undefined
s3StorageConfigured?: boolean
chatInputEnabled?: boolean
actions?: import('svelte').Snippet
innerBottomSnippet?: import('svelte').Snippet
fieldHeaderActions?: import('svelte').Snippet
@@ -182,6 +184,7 @@
computeS3ForceViewerPolicies = undefined,
workspace = undefined,
s3StorageConfigured = true,
chatInputEnabled = false,
actions,
innerBottomSnippet,
fieldHeaderActions,
@@ -255,7 +258,7 @@
nvalue = structuredClone($state.snapshot(defaultValue))
if (defaultValue === undefined || defaultValue === null) {
if (inputCat === 'string') {
nvalue = nullable ? null : ''
nvalue = nullable ? null : format === 'uuid' && extra?.['x-auto-generate'] ? randomUUID() : ''
} else if (inputCat == 'enum' && required) {
let firstV = enum_?.[0]
if (typeof firstV === 'string') {
@@ -1143,6 +1146,7 @@
{disablePortal}
{disabled}
{prettifyHeader}
{chatInputEnabled}
hiddenArgs={['label', 'kind']}
schema={{
properties: obj.properties,
@@ -69,6 +69,7 @@
helperScript?: DynamicInputTypes.HelperScript | undefined
isAgentTool?: boolean
s3StorageConfigured?: boolean
chatInputEnabled?: boolean
}
let {
@@ -94,7 +95,8 @@
otherArgs = {},
helperScript = undefined,
isAgentTool = false,
s3StorageConfigured = true
s3StorageConfigured = true,
chatInputEnabled = false
}: Props = $props()
let monaco: SimpleEditor | undefined = $state(undefined)
@@ -338,6 +340,19 @@
otherArgs: Record<string, any>
) {
const schemaProperty = schema?.properties?.[argName]
if (schemaProperty?.hideWhenChatEnabled && chatInputEnabled) {
if (!hidden) {
hidden = true
if (arg) {
arg.value = undefined
arg.expr = undefined
}
inputCheck = true
}
return
}
if (schemaProperty?.showExpr) {
// Build args object with current field value and other context
const currentValue = propertyType === 'static' ? arg?.value : arg?.expr
@@ -802,6 +817,7 @@
bind:placeholder={schema.properties[argName].placeholder}
{helperScript}
{s3StorageConfigured}
{chatInputEnabled}
otherArgs={Object.fromEntries(
Object.entries(otherArgs).map(([key, transform]) => [
key,
@@ -26,6 +26,7 @@
class?: string
helperScript?: DynamicInputTypes.HelperScript
isAgentTool?: boolean
chatInputEnabled?: boolean
}
let {
@@ -40,7 +41,8 @@
enableAi = false,
class: clazz = '',
helperScript = undefined,
isAgentTool = false
isAgentTool = false,
chatInputEnabled = false
}: Props = $props()
let inputCheck: { [id: string]: boolean } = $state({})
@@ -140,6 +142,7 @@
{helperScript}
{isAgentTool}
{s3StorageConfigured}
{chatInputEnabled}
otherArgs={Object.fromEntries(
Object.entries(args ?? {}).filter(([key]) => key !== argName)
)}
+22 -6
View File
@@ -70,6 +70,7 @@
| undefined)
| undefined
workspace?: string | undefined
chatInputEnabled?: boolean
actions?: import('svelte').Snippet<[{ item: { id: string; value: string } }]> | undefined
}
@@ -112,6 +113,7 @@
lightHeaderFont = false,
computeS3ForceViewerPolicies = undefined,
workspace = undefined,
chatInputEnabled = false,
actions: actions_render = undefined
}: Props = $props()
@@ -206,8 +208,17 @@
function handleHiddenFields(schema: Schema | any, args: Record<string, any>) {
for (const x of fields) {
if (schema?.properties?.[x.value]?.showExpr) {
if (computeShow(x.value, schema.properties?.[x.value]?.showExpr, args)) {
const prop = schema?.properties?.[x.value]
if (prop?.hideWhenChatEnabled && chatInputEnabled) {
if (!hidden[x.value]) {
hidden[x.value] = true
delete args[x.value]
inputCheck[x.value] = true
}
continue
}
if (prop?.showExpr) {
if (computeShow(x.value, prop.showExpr, args)) {
hidden[x.value] = false
} else if (!hidden[x.value]) {
hidden[x.value] = true
@@ -235,13 +246,18 @@
;[schema, args]
if (args && typeof args == 'object') {
let oneShowExpr = false
let hasShowExpr = false
let hasHideWhenChatEnabled = false
for (const key of fields) {
if (schema?.properties?.[key.value]?.showExpr) {
oneShowExpr = true
const prop = schema?.properties?.[key.value]
if (prop?.showExpr) {
hasShowExpr = true
}
if (prop?.hideWhenChatEnabled && chatInputEnabled) {
hasHideWhenChatEnabled = true
}
}
if (!oneShowExpr) {
if (!hasShowExpr && !hasHideWhenChatEnabled) {
return
}
for (const key in args) {
@@ -611,6 +611,7 @@
{enableAi}
{isAgentTool}
helperScript={retrieveDynCodeAndLang(flowModule.value)}
chatInputEnabled={flowStore.val.value?.chat_input_enabled ?? false}
/>
</PropPickerWrapper>
</div>
@@ -88,7 +88,6 @@ export class FlowChatManager {
}
focusInput() {
console.log('focusInput', this.inputElement)
this.inputElement?.focus()
}
@@ -38,6 +38,17 @@ export const AI_AGENT_SCHEMA: Schema = {
description:
'Configure how conversation memory is managed. Choose "auto" to let Windmill automatically store and load messages (up to N last messages), or "manual" to provide an explicit array of conversation messages. The system_prompt and user_message are added to the messages if provided.',
oneOf: [
{
type: 'object',
title: 'off',
properties: {
kind: {
type: 'string',
enum: ['off'],
description: 'Disable conversation memory'
}
}
},
{
type: 'object',
title: 'auto',
@@ -52,7 +63,15 @@ export const AI_AGENT_SCHEMA: Schema = {
type: 'number',
description:
'Number of most recent messages to store and load. Set to 0 to disable memory.',
default: 0
default: 5
},
memory_id: {
type: 'string',
format: 'uuid',
'x-auto-generate': true,
description:
'Custom memory identifier. Each unique ID maintains separate conversation history.',
hideWhenChatEnabled: true
}
},
required: ['kind'],
@@ -95,6 +95,7 @@ export function filteredContentForExport(flow: ExtendedOpenFlow) {
}
import { dfs as dfsApply } from './dfs'
import { randomUUID } from './conversations/FlowChatManager.svelte'
export function cleanFlow(flow: OpenFlow | any): OpenFlow & {
tag?: string
@@ -122,6 +123,28 @@ export function cleanFlow(flow: OpenFlow | any): OpenFlow & {
if (mod.value.type == 'rawscript' && mod.value.assets?.length == 0) {
mod.value.assets = undefined
}
// Generate memory_id for AI agents with auto memory if not already set
// Only if chat input is not enabled, as otherwise memory id is based on conversation id
if (!newFlow.value.chat_input_enabled && mod.value.type === 'aiagent') {
const memoryTransform = mod.value.input_transforms?.memory
if (memoryTransform?.type === 'static' && memoryTransform.value) {
const memoryValue = memoryTransform.value as {
kind: string
context_length: number
memory_id: string
}
if (
memoryValue.kind === 'auto' &&
memoryValue.context_length > 0 &&
!memoryValue.memory_id
) {
memoryTransform.value = {
...memoryValue,
memory_id: randomUUID()
}
}
}
}
})
if (newFlow.value.concurrency_key == '') {
newFlow.value.concurrency_key = undefined