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', // またはデータベースの URL
  }),
});
 
const agent1 = new Agent({
  name: 'agent1',
  instructions:
    'このエージェントは調査に使用され、完全な回答は作成しません。箇条書きで簡潔に回答してください。',
  description:
    'このエージェントは調査に使用され、完全な回答は作成しません。箇条書きで簡潔に回答してください。',
  model: openai('gpt-4o'),
});
 
const agent2 = new Agent({
  name: 'agent2',
  description: 'このエージェントは調査資料のテキスト合成を行います。段落で構成された記事を執筆します。',
  instructions:
    'このエージェントは調査資料のテキスト合成を行います。調査結果に基づいて完全なレポートを書いてください。箇条書きは使わず、段落で書いてください。最終レポートには箇条書きを一切含めないでください。あなたは記事を書きます。',
  model: openai('gpt-4o'),
});
 
const agentStep1 = createStep({
  id: 'agent-step',
  description: 'このステップは調査とテキスト合成に使用されます。',
  inputSchema: z.object({
    city: z.string().describe('調査する都市'),
  }),
  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: 'このステップは調査とテキスト合成に使用されます。',
  inputSchema: z.object({
    text: z.string().describe('調査する都市'),
  }),
  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: 'このワークフローは特定の都市の調査に最適です。',
  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:
    '都市の調査ができます。調査資料の合成もできます。調査結果に基づいて完全なレポートを執筆することもできます。',
  model: openai('gpt-4o'),
  agents: {
    agent1,
    agent2,
  },
  workflows: {
    workflow1,
  },
  memory: memory,
});
 
const runtimeContext = new RuntimeContext();
 
// これは agent1 を呼び出します。ワークフローは個別の都市向けに設計されているため、ルーティングエージェントによる最適なプリミティブは汎用リサーチの agent1 になります。
console.log(await network.generate('What are the biggest cities in France? How are they like?', { runtimeContext }));
// これは workflow1 を呼び出します。個別の都市を調査する場合、ルーティングエージェントによれば最も適したプリミティブだからです。
console.log(await network.generate('Tell me more about Paris', { runtimeContext }));