class MyWorkflow extends RootConstruct {
  constructor(name: string) {
    super(name);

    // Create a model
    const model = new OllamaModel(this, 'Model', {
      host: 'localhost:11434',
      modelName: 'llama3.1:8b'
    });

    // Create an MCP capability
    const mcpCapability = new MCPCapability(this, 'MCPCapability', {
      transport: {
        type: 'http', // also supports 'stdio'
        uri: 'http://localhost:7000/mcp'
      }
    });

    // Create a capability parser
    const capabilityParser = new StructuredOutputCapabilityParser(this, 'CapabilityParser', {});

    // Create a capable model
    const capableModel = new CapableModel(this, 'CapableModel', {
      model: model,
      capabilities: [mcpCapability],
      capabilityParser
    });

    // Create a workflow with the capable model as the entry point
    const workflow = new Workflow(this, 'Workflow', {
      definition: capableModel
    });
  }
}

// Store the Construct tree
const rootConstruct = new MyWorkflow('Root');

// Create a journal, provide the runtime implementations (modules):
const journal = new Journal([createCoreConstructsModule()], {
  rootConstruct
});

async function runWorkflow() {
  const workflowName = "Root/Workflow"; // full node path
  
  for await (const event of journal.executeWorkflow(workflowName, {
    messages: [
      { role: 'user', content: 'What is the capital of France?' }
    ]
  })) {
    console.log('Event:', event);
  }
}

runWorkflow().catch(err => console.error("Error:", err));

// Creating a model
const model = new OllamaModel(rootConstruct, 'Model', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

// Creating an MCP capability
const mcpCapability = new MCPCapability(rootConstruct, 'MCPCapability', {
  transport: {
    type: 'http', // or 'stdio'
    uri: 'http://localhost:7000/mcp'
  }
});

// Creating a capability parser
const capabilityParser = new TagCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Creating a capability parser
const templateParser = new DotTemplateParser(rootConstruct, 'CapabilityParser', {
  template: `Hello {{=it.user.name}}!`
});

const so = new StructuredOutput(this, "StructuredOutput", {
    capableModel,
    outputType: z.strictObject({
        accountholderName: z.string(),
        accountType: z.string(),
        currentBalance: z.number(),
        mostRecentDepositAmount: z.number().min(0),
        mostRecentWithdrawalAmount: z.number().max(0)
    })
})

// Creating a chain
const chain = new Chain(rootConstruct, 'Chain');
chain.pushLink(new EditMessagesTask(rootConstruct, 'FirstPrompt', {
  appendToLatestMessage: "Summarize the following text:"
}));
chain.pushLink(capableModel);

// Execute the workflow
async function runWorkflow() {
  const workflowName = "Root/Workflow";
  
  for await (const event of journal.executeWorkflow(workflowName, {
    messages: [
      { role: 'user', content: 'What is the capital of France?' }
    ]
  })) {
    // Handle events (model responses, capability calls, etc.)
    console.log('Event:', event);
  }
}

// Create an MCP capability
const mcpCapability = new MCPCapability(rootConstruct, 'MCPCapability', {
  transport: {
    type: 'http',
    uri: 'http://localhost:7000/mcp'
  }
});

// Create a capability parser
const capabilityParser = new TagCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Create a capable model with access to capabilities
const capableModel = new CapableModel(rootConstruct, 'CapableModel', {
  model: model,
  capabilities: [mcpCapability],
  capabilityParser
});

// Create the workflow
const workflow = new Workflow(rootConstruct, 'CapabilityWorkflow', {
  definition: capableModel
});

// HTTP transport
const httpMcpCapability = new MCPCapability(rootConstruct, 'HttpMCPCapability', {
  transport: {
    type: 'http',
    uri: 'http://localhost:7000/mcp'
  }
});

// Stdio transport
const stdioMcpCapability = new MCPCapability(rootConstruct, 'StdioMCPCapability', {
  transport: {
    type: 'stdio',
    command: 'npx',
    arguments: ['-y', '@modelcontextprotocol/server-sequential-thinking']
  }
});

// Add capabilities to a capable model
const capableModel = new CapableModel(rootConstruct, 'CapableModel', {
  model: model,
  capabilities: [httpMcpCapability, stdioMcpCapability],
  capabilityParser
});

// Create a model and capability parser
const model = new OllamaModel(rootConstruct, 'Model', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

const capabilityParser = new StructuredOutputCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Create a capable model
const capableModel = new CapableModel(rootConstruct, 'CapableModel', {
  model: model,
  capabilities: [],
  capabilityParser
});

// Create a structured output with a Zod schema
const structuredOutput = new StructuredOutput(rootConstruct, 'StructuredOutput', {
  capableModel,
  outputType: z.strictObject({
    name: z.string(),
    age: z.number(),
    interests: z.array(z.string())
  })
});

// Create a workflow with the structured output
const workflow = new Workflow(rootConstruct, 'StructuredOutputWorkflow', {
  definition: structuredOutput
});

// Create a model and capability parser
const model = new OllamaModel(rootConstruct, 'Model', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

const capabilityParser = new StructuredOutputCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Create a capable model
const capableModel = new CapableModel(rootConstruct, 'CapableModel', {
  model: model,
  capabilities: [],
  capabilityParser
});

// Create a range-based gate
const rangeGate = new LLMGate(rootConstruct, 'RangeGate', {
  capableModel: capableModel,
  prompt: "Please analyze the sentiment and provide a score from 1-10.",
  condition: {
    type: "pass_if_in_range",
    gte: 7,      // Pass if score >= 7
    lte: 10,     // Pass if score <= 10
    min: 1,      // Minimum valid score
    max: 10      // Maximum valid score
  }
});

// Create a workflow with the gate
const workflow = new Workflow(rootConstruct, 'GateWorkflow', {
  definition: rangeGate
});

// Create a model and capability parser
const model = new OllamaModel(rootConstruct, 'Model', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

const capabilityParser = new TagCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Create capable models
const capableModel1 = new CapableModel(rootConstruct, 'CapableModel1', {
  model: model,
  capabilities: [],
  capabilityParser
});

const capableModel2 = new CapableModel(rootConstruct, 'CapableModel2', {
  model: model,
  capabilities: [],
  capabilityParser
});

// Create a chain
const chain = new Chain(rootConstruct, 'Chain');

// Add links to the chain
chain.pushLink(new EditMessagesTask(rootConstruct, 'FirstPrompt', {
  appendToLatestMessage: "Summarize the following text:"
}));
chain.pushLink(capableModel1);
chain.pushLink(new EditMessagesTask(rootConstruct, 'SecondPrompt', {
  messagesPush: [{ role: "user", content: "Translate to French:" }]
}));
chain.pushLink(capableModel2);

// Create a workflow with the chain
const workflow = new Workflow(rootConstruct, 'ChainWorkflow', {
  definition: chain
});

// Create models for different routes
const greetingModel = new OllamaModel(rootConstruct, 'GreetingModel', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

const weatherModel = new OllamaModel(rootConstruct, 'WeatherModel', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

const mathModel = new OllamaModel(rootConstruct, 'MathModel', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

// Create a capability parser
const capabilityParser = new StructuredOutputCapabilityParser(rootConstruct, 'CapabilityParser', {});

// Create capable models for different routes
const greetingCapableModel = new CapableModel(rootConstruct, 'GreetingCapableModel', {
  model: greetingModel,
  capabilities: [],
  capabilityParser: new StructuredOutputCapabilityParser(rootConstruct, 'GreetingCapabilityParser', {})
});

const weatherCapableModel = new CapableModel(rootConstruct, 'WeatherCapableModel', {
  model: weatherModel,
  capabilities: [],
  capabilityParser: new StructuredOutputCapabilityParser(rootConstruct, 'WeatherCapabilityParser', {})
});

const mathCapableModel = new CapableModel(rootConstruct, 'MathCapableModel', {
  model: mathModel,
  capabilities: [],
  capabilityParser: new StructuredOutputCapabilityParser(rootConstruct, 'MathCapabilityParser', {})
});

// Create a model for the router
const routerModel = new OllamaModel(rootConstruct, 'RouterModel', {
  host: 'localhost:11434',
  modelName: 'llama3.1:8b'
});

// Create a capable model for the router
const routerCapableModel = new CapableModel(rootConstruct, 'RouterCapableModel', {
  model: routerModel,
  capabilities: [],
  capabilityParser
});

// Create a router with the tasks
const router = new Router(rootConstruct, 'Router', {
  capableModel: routerCapableModel,
  routes: [
    {
      name: 'greeting',
      description: 'Greetings, introductions, and general pleasantries',
      task: greetingCapableModel
    },
    {
      name: 'weather',
      description: 'Weather forecasts, conditions, and related questions',
      task: weatherCapableModel
    },
    {
      name: 'math',
      description: 'Mathematical calculations, equations, and problems',
      task: mathCapableModel
    }
  ],
  defaultRoute: 'greeting' // Optional default route if no match is found
});

// Create a workflow with the router
const workflow = new Workflow(rootConstruct, 'RouterWorkflow', {
  definition: router
});

// Create models
const optimizerModel = new OllamaModel(this, 'OptimizerModel', {
  host: "ollama:11434",
  modelName: "llama3.1:8b",
});

const evaluatorModel = new OllamaModel(this, 'EvaluatorModel', {
  host: "ollama:11434",
  modelName: "llama3.1:8b",
});

// Create capability parsers
const optimizerCapabilityParser = new StructuredOutputCapabilityParser(this, "OptimizerCapabilityParser", {});
const evaluatorCapabilityParser = new StructuredOutputCapabilityParser(this, "EvaluatorCapabilityParser", {});

// Create capable models
const optimizerCapableModel = new CapableModel(this, "OptimizerCapableModel", {
  model: optimizerModel,
  capabilities: [],
  capabilityParser: optimizerCapabilityParser
});

const evaluatorCapableModel = new CapableModel(this, "EvaluatorCapableModel", {
  model: evaluatorModel,
  capabilities: [],
  capabilityParser: evaluatorCapabilityParser
});

// Create template parsers
const evaluatorTemplate = new DotTemplateParser(this, "EvaluatorTemplate", {
  template: `
    You are an expert evaluator. Your task is to evaluate the quality of the response to the given prompt.

    Original prompt:
    {{=it.originalPrompt}}

    Response to evaluate:
    {{=it.response}}

    Provide a score from 1-10 where:
    1-3: Poor quality, major issues
    4-6: Average quality, some issues
    7-8: Good quality, minor issues
    9-10: Excellent quality, no significant issues

    Also provide specific, actionable feedback on how to improve the response.

    Return your evaluation as a JSON object with the following fields:
    - score: A number between 1 and 10
    - feedback: A string with specific, actionable feedback
    - shouldContinue: A boolean indicating whether the response needs further improvement (true) or is good enough (false)
  `,
  stripWhitespace: false
});

const optimizerTemplate = new DotTemplateParser(this, "OptimizerTemplate", {
  template: `
    You are tasked with generating a high-quality response to the following prompt:

    {{=it.originalPrompt}}

    {{? it.feedback}}
    Here is feedback on your previous attempt:
    Score: {{=it.score}}/10
    Feedback: {{=it.feedback}}

    Please improve your response based on this feedback.
    {{?}}

    Provide a comprehensive, well-structured response that addresses all aspects of the prompt.
  `,
  stripWhitespace: false
});

// Create evaluator optimizer
const evaluatorOptimizer = new EvaluatorOptimizer(this, 'EvaluatorOptimizer', {
  optimizerTask: optimizerCapableModel,
  evaluatorTask: evaluatorCapableModel,
  evaluatorTemplate: evaluatorTemplate,
  optimizerTemplate: optimizerTemplate,
  iterationHardLimit: 3,
  targetScore: 8
});

// Create workflow
const workflow = new Workflow(this, 'Workflow', {
  definition: evaluatorOptimizer
});

// Create a journal with the necessary modules
const journal = new Journal([createCoreConstructsModule()], {
  rootConstruct
});

// Execute a workflow
async function executeWorkflow(input) {
  const workflowName = 'YourWorkflowName';
  
  for await (const event of journal.executeWorkflow(workflowName, input)) {
    // Process events (agent responses, tool calls, etc.)
    handleEvent(event);
  }
  
  // Get the final state
  const finalState = journal.toSavedState();
  return finalState;
}

// Save and restore state
function saveWorkflowState() {
  const state = journal.toSavedState();
  localStorage.setItem('workflowState', JSON.stringify(state));
}

function restoreWorkflowState() {
  const stateJson = localStorage.getItem('workflowState');
  if (stateJson) {
    const state = JSON.parse(stateJson);
    return Journal.fromSavedState(state, [createCoreConstructsModule()]);
  }
  return null;
}

try {
  for await (const event of journal.executeWorkflow(workflowName, input)) {
    if (event.type === 'task_error') {
      // Handle error events
      console.error('Workflow error:', event.error);
      // Implement recovery strategy
    } else {
      // Process normal events
      handleEvent(event);
    }
  }
} catch (error) {
  // Handle unexpected errors
  console.error('Unexpected error:', error);
  // Implement fallback strategy
}

npm install --save-peer @ferment-ai/runtime-common zod

// In your-lib/src/lib/task-defs.ts
import { z } from 'zod';
import { TaskDef } from '@ferment-ai/runtime-common';

const MyCustomInputSchema = z.object({
  param1: z.string(),
  param2: z.number().optional()
});

const MyCustomOutputSchema = z.object({
  result: z.string()
});

export const MY_CUSTOM_TASK_DEF: TaskDef<typeof MyCustomInputSchema, typeof MyCustomOutputSchema> = {
  taskDefId: 'YourLib::MyCustomTaskDef',
  inputType: MyCustomInputSchema,
  outputType: MyCustomOutputSchema
};

// In your-lib/src/lib/my-custom-construct.ts
import { Construct } from 'constructs';
import { WorkflowTask } from '@ferment-ai/runtime-common';
import { MY_CUSTOM_TASK_DEF } from './task-defs';

export interface MyCustomConstructProps {
  name: string;
  description: string;
  // Other properties
}

export class MyCustomConstruct extends WorkflowTask<typeof MY_CUSTOM_TASK_DEF.inputType, typeof MY_CUSTOM_TASK_DEF.outputType> {
  public readonly props: MyCustomConstructProps;
  
  public override readonly taskDef = MY_CUSTOM_TASK_DEF;
  
  constructor(scope: Construct, id: string, props: MyCustomConstructProps) {
    super(scope, id, {});
    this.props = props;
    // Initialize other properties
  }
  
  // Add methods as needed. Setters for props, etc.
}

// In your-runtime/src/lib/my-custom-task.ts
import { MyCustomConstruct, MY_CUSTOM_TASK_DEF } from 'your-lib';
import { TaskCtx } from '@ferment-ai/runtime-common';

export function createMyCustomTaskImpl(construct: MyCustomConstruct): TaskImpl<typeof MY_CUSTOM_TASK_DEF.inputType, typeof MY_CUSTOM_TASK_DEF.outputType> {
  return {
    def: MY_CUSTOM_TASK_DEF,
    nodePath: construct.node.path,
    // for a generator function (this should be your default):
    execute: async function* (ctx: TaskCtx<typeof MY_CUSTOM_TASK_DEF.inputType, typeof MY_CUSTOM_TASK_DEF.outputType>) {
      // return type: TaskCallResult
    }

    // or, if you do not need to call other tools, and wish to use a promise:
    execute: convertPromiseToGenerator(async (ctx: TaskCtx<typeof MY_CUSTOM_TASK_DEF.inputType, typeof MY_CUSTOM_TASK_DEF.outputType>) => {
      // return type: TaskCallResult
    })
  }
}

// In your-runtime/src/lib/module.ts
import { Construct } from 'constructs';
import { Module } from '@ferment-ai/runtime-common';
import { MyCustomConstruct } from 'your-lib';
import { executeMyCustomTask } from './my-custom-task';

export function createYourLibModule(): Module {
  return (construct: Construct) => {
    if (construct instanceof WorkflowTask) {
      switch(construct.taskDef.taskDefId) {
        // we cast the constructs here instead of using instanceof so that reimplementors of the `-lib` works
        case INVOKE_MODEL_TASK_DEF.taskDefId:
          return createMyCustomTaskImpl(construct as MyCustomConstruct);
        default:
          //fallthrough
      }
    }
    
    // No task implementation for this construct from this module
    return undefined;
  };
}

// Example of an L2 construct for a RAG system
export class RAGSystem extends Construct {
  public readonly vectorStore: VectorStore;
  public readonly retriever: Retriever;
  public readonly capableModel: CapableModel;
  
  constructor(scope: Construct, id: string, props: RAGSystemProps) {
    super(scope, id, {});
    
    // Create the vector store
    this.vectorStore = new VectorStore(this, 'VectorStore', {
      documents: props.documents,
      embeddingModel: props.embeddingModel
    });
    
    // Create the retriever
    this.retriever = new Retriever(this, 'Retriever', {
      vectorStore: this.vectorStore,
      topK: props.topK ?? 5
    });
    
    // Create the capability parser
    const capabilityParser = new TagCapabilityParser(this, 'CapabilityParser', {});
    
    // Create the retriever capability
    const retrieverCapability = new CustomCapability(this, 'RetrieverCapability', {
      retriever: this.retriever
    });
    
    // Create the capable model with access to the retriever
    this.capableModel = new CapableModel(this, 'CapableModel', {
      model: props.model,
      capabilities: [retrieverCapability],
      capabilityParser
    });
  }
}