kage-research/level1.ts

/**
 * Level 1 POC: Minimal Pi Shadow
 * 
 * This tests:
 * 1. Pi agent-core works
 * 2. OpenRouter integration
 * 3. Basic tool execution
 * 4. Memory usage
 */

import { Agent } from "@mariozechner/pi-agent-core";
import { registerBuiltInApiProviders, streamSimple } from "@mariozechner/pi-ai";
import type { Model } from "@mariozechner/pi-ai";
import * as fs from "fs";
import { exec } from "child_process";

// Set API key from environment
const OPENROUTER_API_KEY = process.env.OPENROUTER_API_KEY || "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";
process.env.OPENROUTER_API_KEY = OPENROUTER_API_KEY;

// Register the API providers
registerBuiltInApiProviders();

// Manually create model for OpenRouter - Free model
const model: Model<"openai-responses"> = {
  id: "stepfun/step-3.5-flash:free",
  name: "Step-3.5 Flash (Free)",
  api: "openai-responses",
  provider: "openrouter",
  baseUrl: "https://openrouter.ai/api/v1",
  reasoning: false,
  input: ["text"],
  output: ["text"],
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  contextWindow: 128000,
  maxTokens: 8192,
};

// Memory tracking
const startMemory = process.memoryUsage();
console.log("📊 Start Memory:", {
  heapUsed: Math.round(startMemory.heapUsed / 1024 / 1024) + " MB",
  heapTotal: Math.round(startMemory.heapTotal / 1024 / 1024) + " MB",
  rss: Math.round(startMemory.rss / 1024 / 1024) + " MB",
});

// Basic tools similar to what OpenCode provides
const tools = [
  {
    name: "read",
    label: "Read File",
    description: "Read the contents of a file",
    parameters: {
      type: "object",
      properties: {
        path: { type: "string", description: "Path to the file to read" },
      },
      required: ["path"],
    } as const,
    execute: async (toolCallId: string, params: { path: string }) => {
      try {
        const content = fs.readFileSync(params.path, "utf-8");
        return {
          content: [{ type: "text" as const, text: content }],
          details: { path: params.path, lines: content.split("\n").length },
        };
      } catch (error: any) {
        throw new Error(`Failed to read file: ${error.message}`);
      }
    },
  },
  {
    name: "bash",
    label: "Run Command",
    description: "Run a shell command",
    parameters: {
      type: "object",
      properties: {
        command: { type: "string", description: "Command to run" },
      },
      required: ["command"],
    } as const,
    execute: async (toolCallId: string, params: { command: string }) => {
      return new Promise((resolve, reject) => {
        exec(params.command, { cwd: process.cwd() }, (error, stdout, stderr) => {
          if (error) {
            resolve({
              content: [{ type: "text" as const, text: stderr || error.message }],
              details: { command: params.command, exitCode: error.code },
              isError: true,
            });
          } else {
            resolve({
              content: [{ type: "text" as const, text: stdout }],
              details: { command: params.command, exitCode: 0 },
            });
          }
        });
      });
    },
  },
];

// Create the agent
const agent = new Agent({
  initialState: {
    systemPrompt: `You are a helpful coding assistant. You have access to tools:
- read: Read file contents
- bash: Run shell commands

Use these tools to help the user. Be concise and practical.`,
    model: model,
    tools: tools as any,
    messages: [],
  },
  convertToLlm: (messages) => {
    // Filter to only user, assistant, toolResult roles
    return messages
      .filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult")
      .map((m) => ({
        role: m.role,
        content: m.content,
      }));
  },
});

// Track events
const events: string[] = [];
agent.subscribe((event) => {
  events.push(event.type);
  
  switch (event.type) {
    case "agent_start":
      console.log("🤖 Agent started");
      break;
    case "turn_start":
      console.log("🔄 Turn started");
      break;
    case "message_start":
      if ('message' in event && event.message) {
        const msg = event.message as any;
        if (msg.role === 'assistant') {
          console.log("\n💬 Assistant:");
        }
      }
      break;
    case "message_update":
      if ("assistantMessageEvent" in event) {
        const ev = event as any;
        if (ev.assistantMessageEvent.type === "text_delta") {
          const text = ev.assistantMessageEvent.delta || '';
          process.stdout.write(text);
        }
        if (ev.assistantMessageEvent.type === "content_block_delta") {
          // Handle content block updates
          const content = ev.assistantMessageEvent.delta?.content?.[0];
          if (content?.type === 'text' && content?.text) {
            process.stdout.write(content.text);
          }
        }
      }
      break;
    case "tool_execution_start":
      console.log(`\n🔧 Tool: ${event.toolName}`);
      break;
    case "tool_execution_end":
      console.log(`   → Done (error: ${event.isError})`);
      break;
    case "turn_end":
      console.log("\n✅ Turn ended");
      break;
    case "agent_end":
      console.log("\n🏁 Agent finished");
      
      // Log final messages
      if (event.messages && event.messages.length > 0) {
        console.log("\n📝 Final messages:");
        event.messages.slice(-3).forEach((msg: any, i: number) => {
          console.log(`  [${i}] ${msg.role}:`, (msg.content?.[0]?.text || '').substring(0, 100));
        });
      }
      
      // Final memory
      const endMemory = process.memoryUsage();
      console.log("\n📊 End Memory:", {
        heapUsed: Math.round(endMemory.heapUsed / 1024 / 1024) + " MB",
        heapTotal: Math.round(endMemory.heapTotal / 1024 / 1024) + " MB",
        rss: Math.round(endMemory.rss / 1024 / 1024) + " MB",
      });
      
      console.log("\n📋 Event sequence:", events.join(" → "));
      break;
  }
});

async function main() {
  console.log("\n🚀 Starting Pi agent with OpenRouter...\n");

  // Run a simple task
  try {
    console.log("\n📝 Prompt: Say hello and tell me the current time using bash command 'date'.\n");
    await agent.prompt("Say hello and tell me the current time using bash command 'date'.");
  } catch (error) {
    console.error("❌ Error:", error);
  }
  
  // Check if there's an error message
  if (agent.state.errorMessage) {
    console.log("❌ Agent error:", agent.state.errorMessage);
  }
}

main().catch(console.error);