File size: 5,850 Bytes
9b9eafc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
// Vòng lặp agent tối giản (thay cho Pi agent-core của bản gốc):
//   user → model → (tool calls → kết quả → model)… → câu trả lời cuối
// Tool chạy tuần tự; tối đa 12 lượt model mỗi yêu cầu.

import type { WorkerEvent } from '../shared/protocol';
import type { AgentMessage, AssistantMessage, ToolResultMessage } from '../shared/types';
import type { Engine } from './engine';
import { TOOLS, type Workspace } from './workspace';

export const MAX_AGENT_TURNS = 12;

export const SYSTEM_PROMPT = [
  'You are Pi, a coding assistant powered by MiniCPM5-2B, a language model from OpenBMB.',
  "The model, inference, tools, and workspace run entirely in the user's browser. Workspace files are stored in this browser, not on the user's operating-system filesystem.",
  'Use your tools to inspect and change actual files. Never claim a command ran or a file changed without a successful tool result.',
  'When the user explicitly asks you to use a tool, call that tool even if you already know the answer.',
  'Working directory: /workspace. Use read, write, edit, and bash. Keep tool calls and answers concise.',
  'The bash tool supports cat, ls, find, grep, sed, awk, jq, sort, wc, printf, pipes, redirects, and shell scripts.',
  'Use jq for JSON and awk for CSV and calculations; python, python3, node, npm, native binaries, and network access are unavailable.',
  'You can create and edit text files, including HTML, CSS, JavaScript, and shell scripts. HTML may reference external assets by URL; writing such a reference does not require network access from the shell. Do not claim an external asset was fetched or tested.',
  'Interpret short follow-ups using the conversation and inspect the relevant files when needed.',
  'Inspect files before editing them. After changing files, verify the result with a tool. Finish with a short factual answer.',
].join('\n');

export interface AgentDeps {
  engine: Engine;
  workspace: Workspace;
  emit: (event: WorkerEvent) => void;
  /** Gọi sau mỗi tool chạy xong và khi kết thúc yêu cầu (để lưu workspace + hội thoại). */
  persist: () => Promise<void>;
  streamIntervalMs?: number;
}

export class Agent {
  messages: AgentMessage[];
  private streamTimer: ReturnType<typeof setTimeout> | undefined;
  private streamDirty: AssistantMessage | null = null;

  constructor(
    private deps: AgentDeps,
    initialMessages: AgentMessage[] = [],
  ) {
    this.messages = [...initialMessages];
  }

  reset() {
    this.messages = [];
    this.deps.emit({ type: 'messages', messages: [] });
  }

  private publish() {
    this.deps.emit({ type: 'messages', messages: this.messages });
  }

  private push(message: AgentMessage) {
    this.messages.push(message);
    this.publish();
  }

  /** Gộp các cập nhật streaming để không gửi quá dày (mặc định ~20 lần/giây). */
  private streamPartial = (message: AssistantMessage) => {
    this.streamDirty = message;
    if (this.streamTimer) return;
    this.streamTimer = setTimeout(() => {
      this.streamTimer = undefined;
      if (this.streamDirty) this.deps.emit({ type: 'stream', message: this.streamDirty });
    }, this.deps.streamIntervalMs ?? 50);
  };

  private clearStream() {
    if (this.streamTimer) clearTimeout(this.streamTimer);
    this.streamTimer = undefined;
    this.streamDirty = null;
    this.deps.emit({ type: 'stream', message: null });
  }

  async prompt(text: string, signal?: AbortSignal): Promise<void> {
    const { engine, workspace, emit, persist } = this.deps;
    let turns = 0;
    this.push({ role: 'user', content: text, timestamp: Date.now() });

    try {
      for (;;) {
        const assistant = await engine.generate({
          systemPrompt: SYSTEM_PROMPT,
          messages: this.messages,
          tools: TOOLS,
          signal,
          emit,
          onPartial: this.streamPartial,
        });
        this.clearStream();
        this.push(assistant);

        if (assistant.stopReason === 'error' || assistant.stopReason === 'aborted') break;
        const calls = assistant.content.filter((c) => c.type === 'toolCall');
        if (!calls.length) break;

        const results: ToolResultMessage[] = [];
        for (const call of calls) {
          let result: ToolResultMessage;
          if (signal?.aborted) {
            result = this.toolResult(call.id, call.name, 'Skipped: the request was stopped.', {}, true);
          } else {
            emit({ type: 'tool_start', name: call.name });
            try {
              const out = await workspace.runTool(call.name, call.arguments, signal);
              result = this.toolResult(call.id, call.name, out.content[0].text, out.details, out.isError);
            } catch (err) {
              const msg = signal?.aborted ? 'Stopped.' : String((err as Error)?.message ?? err);
              result = this.toolResult(call.id, call.name, msg, {}, true);
            }
            emit({ type: 'tool_end', name: call.name });
          }
          results.push(result);
          this.push(result);
          await persist();
        }

        turns++;
        if (signal?.aborted) break;
        if (turns >= MAX_AGENT_TURNS && results.length > 0) {
          emit({ type: 'notice', text: `Paused after ${MAX_AGENT_TURNS} agent turns. Send a follow-up to continue.` });
          break;
        }
      }
    } finally {
      this.clearStream();
      await persist();
    }
  }

  private toolResult(
    toolCallId: string,
    toolName: string,
    text: string,
    details: ToolResultMessage['details'],
    isError: boolean,
  ): ToolResultMessage {
    return {
      role: 'toolResult',
      toolCallId,
      toolName,
      content: [{ type: 'text', text }],
      details,
      isError,
      timestamp: Date.now(),
    };
  }
}