import { NewAgentNetwork } from '@mastra/core/network/vNext';
import { Agent } from '@mastra/core/agent';
import { createStep, createWorkflow } from '@mastra/core/workflows';
import { Memory } from '@mastra/memory';
import { openai } from '@ai-sdk/openai';
import { LibSQLStore } from '@mastra/libsql';
import { z } from 'zod';
import { RuntimeContext } from '@mastra/core/runtime-context';
 
const memory = new Memory({
  storage: new LibSQLStore({
    url: 'file:../mastra.db', // Or your database URL
  }),
});
 
const agent1 = new Agent({
  name: 'agent1',
  instructions:
    'This agent is used to do research, but not create full responses. Answer in bullet points only and be concise.',
  description:
    'This agent is used to do research, but not create full responses. Answer in bullet points only and be concise.',
  model: openai('gpt-4o'),
});
 
const agent2 = new Agent({
  name: 'agent2',
  description: 'This agent is used to do text synthesis on researched material. It writes articles in full paragraphs.',
  instructions:
    'This agent is used to do text synthesis on researched material. Write a full report based on the researched material. Do not use bullet points. Write full paragraphs. There should not be a single bullet point in the final report. You write articles.',
  model: openai('gpt-4o'),
});
 
const agentStep1 = createStep({
  id: 'agent-step',
  description: 'This step is used to do research and text synthesis.',
  inputSchema: z.object({
    city: z.string().describe('The city to research'),
  }),
  outputSchema: z.object({
    text: z.string(),
  }),
  execute: async ({ inputData }) => {
    const resp = await agent1.generate(inputData.city, {
      output: z.object({
        text: z.string(),
      }),
    });
 
    return { text: resp.object.text };
  },
});
 
const agentStep2 = createStep({
  id: 'agent-step-two',
  description: 'This step is used to do research and text synthesis.',
  inputSchema: z.object({
    text: z.string().describe('The city to research'),
  }),
  outputSchema: z.object({
    text: z.string(),
  }),
  execute: async ({ inputData }) => {
    const resp = await agent2.generate(inputData.text, {
      output: z.object({
        text: z.string(),
      }),
    });
 
    return { text: resp.object.text };
  },
});
 
const workflow1 = createWorkflow({
  id: 'workflow1',
  description: 'This workflow is perfect for researching a specific city.',
  steps: [],
  inputSchema: z.object({
    city: z.string(),
  }),
  outputSchema: z.object({
    text: z.string(),
  }),
})
  .then(agentStep1)
  .then(agentStep2)
  .commit();
 
const network = new NewAgentNetwork({
  id: 'test-network',
  name: 'Test Network',
  instructions:
    'You can research cities. You can also synthesize research material. You can also write a full report based on the researched material.',
  model: openai('gpt-4o'),
  agents: {
    agent1,
    agent2,
  },
  workflows: {
    workflow1,
  },
  memory: memory,
});
 
const runtimeContext = new RuntimeContext();
 
// This will call agent1, as the workflow is meant to be used with individual cities. The best primitive according to the routing agent is thus agent1 which is a general research primitive.
console.log(await network.generate('What are the biggest cities in France? How are they like?', { runtimeContext }));
// This will call workflow1, as it is the most suitable primitive according to the routing agent when researching individual cities.
console.log(await network.generate('Tell me more about Paris', { runtimeContext }));