CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes48downloads
extractionAgentPlanner.ts308 linesDownload Raw Back to memory
1/**2 * @license3 * Copyright 2026 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7import type { Config } from '../config/config.js';8import { createDebugLogger } from '../utils/debugLogger.js';9import { runForkedAgent, getCacheSafeParams } from '../utils/forkedAgent.js';10import { buildFunctionResponseParts } from '../tools/agent/fork-subagent.js';11import type { Content } from '@google/genai';12import {13  MEMORY_FRONTMATTER_EXAMPLE,14  TYPES_SECTION_INDIVIDUAL,15  WHAT_NOT_TO_SAVE_SECTION,16} from './prompt.js';17import {18  AUTO_MEMORY_INDEX_FILENAME,19  getAutoMemoryRoot,20  getUserAutoMemoryRoot,21} from './paths.js';22import type { AutoMemoryType } from './types.js';23import {24  scanAutoMemoryTopicDocuments,25  scanUserAutoMemoryTopicDocuments,26} from './scan.js';27import { ToolNames } from '../tools/tool-names.js';28import { createMemoryScopedAgentConfig } from './memory-scoped-agent-config.js';29 30const MAX_TOPIC_SUMMARY_CHARS = 280;31 32const debugLogger = createDebugLogger('AUTO_MEMORY_EXTRACTION_AGENT');33 34const EXTRACTION_AGENT_SYSTEM_PROMPT = [35  'You are now acting as the managed memory extraction subagent for an AI coding assistant.',36  '',37  'The recent conversation history is already in your context. Analyze only that recent conversation and use it to update persistent managed memory.',38  '',39  'Rules:',40  '- Read existing memory files first to avoid creating duplicates.',41  '- Extract only durable facts stated by the user.',42  '- Ignore temporary, session-specific, speculative, or question content.',43  '- If the user explicitly asks the assistant to remember something durable, preserve it.',44  '- Use one of the allowed topics: user, feedback, project, reference.',45  '- Keep entries concise and suitable for bullet points. No leading bullet markers.',46  '- Do not investigate repository code, git history, or unrelated files.',47  '- Work only from the conversation history in your context and the existing memory files.',48  '- If nothing durable should be saved, make no file changes.',49  '',50  ...TYPES_SECTION_INDIVIDUAL,51  ...WHAT_NOT_TO_SAVE_SECTION,52  '',53  'Memory file format reference:',54  ...MEMORY_FRONTMATTER_EXAMPLE,55].join('\n');56 57export interface AutoMemoryExtractionExecutionResult {58  touchedTopics: AutoMemoryType[];59  /** True when at least one file inside the project-level memory root was written/edited. */60  touchedProjectScope: boolean;61  /** True when at least one file inside the user-level memory root was written/edited. */62  touchedUserScope: boolean;63  systemMessage?: string;64  hasToolActivity: boolean;65}66 67/**68 * Ensure the history slice ends with a `model` text message so that69 * agent-headless can send the task prompt as the first user turn without70 * creating consecutive user messages (Gemini API constraint).71 *72 * - Trailing `user` message: drop it.73 * - Last `model` message has open function calls: close them with placeholder74 *   responses and append a model ack so the sequence stays valid.75 * - Otherwise: return a shallow copy as-is.76 */77function buildAgentHistory(history: Content[]): Content[] {78  if (history.length === 0) return [];79  const last = history[history.length - 1];80  if (last.role !== 'model') {81    return history.slice(0, -1);82  }83  const openCalls = (last.parts ?? []).filter((p) => p.functionCall);84  if (openCalls.length === 0) {85    return [...history];86  }87  const toolResponses = buildFunctionResponseParts(88    last,89    'Background extraction started.',90  );91  return [92    ...history,93    { role: 'user' as const, parts: toolResponses },94    { role: 'model' as const, parts: [{ text: 'Acknowledged.' }] },95  ];96}97 98function truncate(text: string, maxChars: number): string {99  const normalized = text.replace(/\s+/g, ' ').trim();100  if (normalized.length <= maxChars) {101    return normalized;102  }103  return `${normalized.slice(0, maxChars).trimEnd()}…`;104}105 106async function buildTopicSummaryBlock(projectRoot: string): Promise<string> {107  // User-level scan is best-effort: a read failure on `~/.qwen/memories/`108  // must not deny the extraction agent its view of existing project-level109  // memories (which it uses to avoid creating duplicates).110  const [projectDocs, userDocs] = await Promise.all([111    scanAutoMemoryTopicDocuments(projectRoot),112    scanUserAutoMemoryTopicDocuments().catch((error: unknown) => {113      debugLogger.warn(114        `User-level auto-memory scan failed; extraction agent will see project-level summaries only: ${error instanceof Error ? error.message : String(error)}`,115      );116      return [];117    }),118  ]);119 120  const renderDoc = (doc: (typeof projectDocs)[number], scope: string) => {121    const body = truncate(122      doc.body === '_No entries yet._' ? '' : doc.body,123      MAX_TOPIC_SUMMARY_CHARS,124    );125    return [126      `- [${doc.title}](${doc.relativePath}) — ${doc.description || '(no description)'}`,127      `  scope=${scope}`,128      `  topic=${doc.type}`,129      `  path=${doc.filePath}`,130      `  current=${body || '(empty)'}`,131    ].join('\n');132  };133 134  const blocks = [135    ...userDocs.map((doc) => renderDoc(doc, 'user')),136    ...projectDocs.map((doc) => renderDoc(doc, 'project')),137  ];138 139  return blocks.join('\n\n');140}141 142function buildTaskPrompt(143  projectMemoryRoot: string,144  userMemoryRoot: string,145  topicSummaries: string,146): string {147  return [148    'Managed memory has TWO directories. Choose which one to write each memory into using the per-type `<scope>` guidance in your system instructions:',149    `- USER memory (cross-project, durable knowledge about who the user is): \`${userMemoryRoot}\``,150    `- PROJECT memory (this project only): \`${projectMemoryRoot}\``,151    '',152    'Scan the recent conversation history in your context and update durable managed memory in whichever directory each memory belongs.',153    '',154    'Available tools in this run: `read_file`, `grep_search`, `glob`, `list_directory`, read-only `run_shell_command`, and `write_file`/`edit` for paths inside EITHER managed memory directory above.',155    '- Do not use any other tools.',156    '- You have a limited turn budget. `edit` requires a prior `read_file` of the same file, so the efficient strategy is: first issue all reads in parallel for every file you might update; then issue all `write_file`/`edit` calls in parallel. Do not interleave reads and writes across multiple turns.',157    '- You MUST only use content from the recent conversation history in your context plus the current managed memory files.',158    '- Do not inspect repository code, git history, or unrelated files.',159    '- Prefer updating an existing memory file over creating a duplicate. Check both directories for an existing entry before creating a new one.',160    '- Keep one durable memory per file under `user/`, `feedback/`, `project/`, or `reference/` inside the chosen directory.',161    '',162    '## How to save memories',163    '',164    '**Step 1** — write or update the memory file itself, in the directory chosen by the type `<scope>`, using the required frontmatter format.',165    `**Step 2** — update the \`${AUTO_MEMORY_INDEX_FILENAME}\` in the SAME directory where you wrote the file (\`${userMemoryRoot}/${AUTO_MEMORY_INDEX_FILENAME}\` for USER memory, \`${projectMemoryRoot}/${AUTO_MEMORY_INDEX_FILENAME}\` for PROJECT memory). The index is one line per entry: \`- [Title](relative/path.md) — one-line hook\`. Never write memory content directly into the index.`,166    '- If you create or delete a memory file, also update the managed memory index in the SAME directory.',167    '- If nothing durable should be saved, make no file changes.',168    '',169    '## Existing memory files (across both directories)',170    '',171    topicSummaries || '(none yet)',172  ].join('\n');173}174 175/**176 * Derive which memory topics + scopes were touched from the list of file177 * paths written during the agent run. Avoids requiring JSON output from178 * the agent.179 */180function touchedTopicsFromFilePaths(181  filePaths: string[],182  projectRoot: string,183): {184  topics: AutoMemoryType[];185  touchedProjectScope: boolean;186  touchedUserScope: boolean;187} {188  // Use startsWith against the directly-retrieved roots (rather than the189  // isAutoMemPath helper, which calls into paths.ts internals and would190  // bypass module-level mocks in extractionAgentPlanner.test.ts). This191  // also keeps the routing decision symmetric across both scopes.192  const projectRootDir = getAutoMemoryRoot(projectRoot);193  const userRootDir = getUserAutoMemoryRoot();194  // Canonicalize separators to `/` on BOTH sides before the prefix check.195  // On Windows the roots are backslash-native (`C:\Users\foo\...\memory`)196  // while filesTouched (populated from raw model tool-call arguments)197  // commonly comes back forward-slash-normalized — `startsWith` against198  // the raw roots would miss those writes entirely. Also guards against199  // the inverse direction and the historical `/foo/memory` vs200  // `/foo/memory-other/...` collision: the character after the root must201  // be `/` so files inside (never AT) the root match exactly one prefix.202  const canon = (s: string): string => s.replace(/\\/g, '/');203  const isUnderRoot = (canonP: string, canonRoot: string): boolean => {204    if (!canonP.startsWith(canonRoot)) return false;205    return canonP.charAt(canonRoot.length) === '/';206  };207  const canonProject = canon(projectRootDir);208  const canonUser = canon(userRootDir);209  const topicSet = new Set<AutoMemoryType>();210  let touchedProjectScope = false;211  let touchedUserScope = false;212 213  for (const p of filePaths) {214    const canonP = canon(p);215    let canonRoot: string | undefined;216    if (isUnderRoot(canonP, canonProject)) {217      canonRoot = canonProject;218      touchedProjectScope = true;219    } else if (isUnderRoot(canonP, canonUser)) {220      canonRoot = canonUser;221      touchedUserScope = true;222    } else {223      continue;224    }225    // +1 to also strip the `/` we just checked for.226    const rel = canonP.slice(canonRoot.length + 1);227    const segment = rel.split('/')[0] as AutoMemoryType;228    if (229      segment === 'user' ||230      segment === 'feedback' ||231      segment === 'project' ||232      segment === 'reference'233    ) {234      topicSet.add(segment);235    }236  }237  return {238    topics: [...topicSet],239    touchedProjectScope,240    touchedUserScope,241  };242}243 244export async function runAutoMemoryExtractionByAgent(245  config: Config,246  projectRoot: string,247): Promise<AutoMemoryExtractionExecutionResult> {248  const cacheSafe = getCacheSafeParams();249  if (!cacheSafe) {250    throw new Error(251      'runAutoMemoryExtractionByAgent: no cache-safe params available; ' +252        'extraction must run after a completed main turn.',253    );254  }255  const extraHistory = buildAgentHistory(cacheSafe.history);256 257  const topicSummaries = await buildTopicSummaryBlock(projectRoot);258  const projectMemoryRoot = getAutoMemoryRoot(projectRoot);259  const userMemoryRoot = getUserAutoMemoryRoot();260  const scopedConfig = createMemoryScopedAgentConfig(config, projectRoot, {261    allowShell: true,262  });263 264  const result = await runForkedAgent({265    name: 'managed-auto-memory-extractor',266    config: scopedConfig,267    taskPrompt: buildTaskPrompt(268      projectMemoryRoot,269      userMemoryRoot,270      topicSummaries,271    ),272    systemPrompt: EXTRACTION_AGENT_SYSTEM_PROMPT,273    maxTurns: 5,274    maxTimeMinutes: 2,275    tools: [276      ToolNames.READ_FILE,277      ToolNames.GREP,278      ToolNames.GLOB,279      ToolNames.LS,280      ToolNames.SHELL,281      ToolNames.WRITE_FILE,282      ToolNames.EDIT,283    ],284    extraHistory,285  });286 287  if (result.status !== 'completed') {288    throw new Error(289      result.terminateReason ||290        'Extraction agent did not complete successfully',291    );292  }293 294  const { topics, touchedProjectScope, touchedUserScope } =295    touchedTopicsFromFilePaths(result.filesWritten ?? [], projectRoot);296 297  return {298    touchedTopics: topics,299    touchedProjectScope,300    touchedUserScope,301    hasToolActivity: result.filesTouched.length > 0,302    systemMessage:303      topics.length > 0304        ? `Managed auto-memory updated: ${topics.map((t) => `${t}.md`).join(', ')}`305        : undefined,306  };307}308 
basant307/AI_Governance_Project · CoolFace