loki-mode 11.7.1 → 11.9.1

package/memory/rag_injector.py

@@ -1,161 +0,0 @@
-#!/usr/bin/env python3
-"""RAG context injector -- builds knowledge context for RARV prompts.
-
-CLI usage: python3 -m memory.rag_injector --query "description" --max-tokens 2000
-
-Reads from the organization knowledge graph and returns relevant context
-for injection into the RARV prompt cycle.
-"""
-
-import argparse
-import json
-import re
-import sys
-from pathlib import Path
-
-
-# Per-field hard cap so a single oversized/garbage memory entry cannot dominate
-# the injected context or push real signal out of the token budget.
-_MAX_FIELD_CHARS = 600
-
-# Control characters (excluding ordinary whitespace handled separately) that
-# must never reach the prompt: they can corrupt rendering or smuggle hidden
-# directives. Covers C0 controls and DEL.
-_CONTROL_CHARS_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
-
-
-def _sanitize_field(value, max_chars=_MAX_FIELD_CHARS):
-    """Neutralize a single field from stored memory before it enters a prompt.
-
-    Memory entries are untrusted input: a malicious or garbage pattern could
-    contain prompt-injection payloads (fake instructions, fabricated section
-    headers, code-fence breakers, or directives that try to override the RARV
-    prompt). This collapses each field to a single safe inline line so it can
-    only ever read as data, never as structure or instructions.
-
-    Steps:
-    - Coerce non-strings to str (a dict/list value would otherwise break
-      formatting or interpolate unexpected structure).
-    - Drop control characters (keeps the visible text intact).
-    - Collapse ALL newlines/tabs/carriage-returns to single spaces so the
-      field cannot add new lines, fake "### " section headers, close the
-      knowledge block, or open/close a markdown code fence.
-    - Defang leading structural markdown markers (#, >, -, *, backticks) so the
-      field cannot pose as a heading/list/fence even after collapsing.
-    - Truncate to a bounded length.
-    """
-    if value is None:
-        return ''
-    if not isinstance(value, str):
-        value = str(value)
-    # Remove control characters first.
-    value = _CONTROL_CHARS_RE.sub('', value)
-    # Collapse any vertical whitespace and tabs into single spaces. This is the
-    # core injection defense: without newlines the field cannot introduce new
-    # markdown structure or standalone instruction lines.
-    value = re.sub(r'[\r\n\t]+', ' ', value)
-    # Collapse runs of spaces left behind.
-    value = re.sub(r' {2,}', ' ', value).strip()
-    # Defang leading structural markdown so the field cannot masquerade as a
-    # heading, blockquote, list item, or code fence.
-    value = re.sub(r'^[#>\-*`]+\s*', '', value)
-    # Defang markdown header markers anywhere in the (now single-line) value so
-    # a collapsed payload cannot reintroduce a "### " section marker inline.
-    # Backticks are neutralized too so a field cannot open/close a code fence.
-    value = re.sub(r'#{1,6}\s', '', value)
-    value = value.replace('`', '')
-    if len(value) > max_chars:
-        value = value[:max_chars].rstrip() + '...'
-    return value
-
-
-def build_rag_context(query, max_tokens=2000, knowledge_dir=None):
-    """Build RAG context from the knowledge graph.
-
-    Args:
-        query: Search query (usually PRD summary or task description)
-        max_tokens: Maximum approximate tokens (chars / 4)
-        knowledge_dir: Override knowledge directory
-
-    Returns:
-        Formatted context string for prompt injection, or empty string
-        if no matching patterns are found.
-    """
-    from .knowledge_graph import OrganizationKnowledgeGraph
-
-    kg = OrganizationKnowledgeGraph(knowledge_dir)
-    patterns = kg.query_patterns(query, max_results=10)
-
-    if not patterns:
-        return ''
-
-    max_chars = max_tokens * 4  # Rough chars-to-tokens ratio
-
-    sections = []
-    total_chars = 0
-
-    for p in patterns:
-        # Support both 'name'/'pattern' and 'description' fields.
-        # Every field is sanitized: memory entries are untrusted input and must
-        # not be able to inject instructions or break the prompt structure.
-        # `or` chaining, not .get(a, .get(b, default)): a present-but-empty
-        # 'description' would otherwise shadow a real 'correct_approach', and a
-        # row missing BOTH name keys would resolve to the literal string
-        # 'Unknown Pattern' -- truthy, so the substance check below could never
-        # fire on it.
-        name = _sanitize_field(p.get('name') or p.get('pattern') or '')
-        desc = _sanitize_field(p.get('description') or p.get('correct_approach') or '')
-        category = _sanitize_field(p.get('category', ''))
-
-        # Substance check: a pattern that matched only on its category (or
-        # whose fields sanitized away to nothing) has no memory to inject.
-        # Rendering it as "### Unknown Pattern" under the "patterns were found
-        # in the organization knowledge base" header presents a fabricated
-        # memory as a real one and burns prompt budget. Name OR description is
-        # enough -- category alone is a bucket label, not a memory, and a named
-        # pattern with no description is still genuine signal.
-        if not name and not desc:
-            continue
-        source_raw = p.get('_source_project', '')
-        # Path().name strips any directory traversal; sanitize the basename too.
-        source = _sanitize_field(Path(str(source_raw)).name) if source_raw else ''
-
-        # name may be empty after sanitization (e.g. a field that was only
-        # markdown markers); fall back so the heading is never blank.
-        if not name:
-            name = 'Unknown Pattern'
-
-        section = '### ' + name
-        if category:
-            section += ' (' + category + ')'
-        section += '\n'
-        if desc:
-            section += desc + '\n'
-        if source:
-            section += '_Source: ' + source + '_\n'
-
-        if total_chars + len(section) > max_chars:
-            break
-        sections.append(section)
-        total_chars += len(section)
-
-    if not sections:
-        return ''
-
-    return 'The following patterns were found in the organization knowledge base:\n\n' + '\n'.join(sections)
-
-
-def main():
-    parser = argparse.ArgumentParser(description='RAG context injector for RARV prompts')
-    parser.add_argument('--query', required=True, help='Search query')
-    parser.add_argument('--max-tokens', type=int, default=2000, help='Max tokens for context')
-    parser.add_argument('--knowledge-dir', help='Override knowledge directory')
-    args = parser.parse_args()
-
-    context = build_rag_context(args.query, args.max_tokens, args.knowledge_dir)
-    if context:
-        print(context)
-
-
-if __name__ == '__main__':
-    main()