#!/usr/bin/env python3 import fcntl import hashlib import os import time from contextlib import contextmanager from pathlib import Path import psycopg import requests from dotenv import load_dotenv from blacklist import ( EXCLUDE_DIR_NAMES, EXCLUDE_FILENAME_KEYWORDS, EXCLUDE_PATH_PARTS, SENSITIVE_LITERAL_MARKERS, SENSITIVE_REGEX_PATTERNS, ) load_dotenv(Path(__file__).parent.parent.parent / '.env.memory') VAULT_ROOT = Path(os.getenv('VAULT_DIR', '.')).resolve() def require_env(name: str) -> str: value = os.getenv(name, '').strip() if not value: raise RuntimeError(f'缺少必需环境变量: {name}') return value def get_conn(): return psycopg.connect(require_env('PG_DSN')) def normalize_rel(path: Path) -> str: return str(path.resolve().relative_to(VAULT_ROOT)).replace('\\', '/') def safe_rel(path: Path) -> str | None: try: return normalize_rel(path) except Exception: return None def sha256_text(text: str) -> str: return hashlib.sha256(text.encode('utf-8')).hexdigest() def vector_literal(values: list[float]) -> str: return '[' + ','.join(f'{v:.8f}' for v in values) + ']' def _sample_head_mid_tail(content: bytes, span: int = 1200) -> str: size = len(content) if size <= span * 3: return content.decode('utf-8', errors='ignore').lower() mid = max(0, (size // 2) - (span // 2)) sampled = content[:span] + content[mid : mid + span] + content[-span:] return sampled.decode('utf-8', errors='ignore').lower() def is_excluded(file_path: Path) -> bool: rel = safe_rel(file_path) if rel is None: return True parts = set(Path(rel).parts) if parts.intersection(EXCLUDE_DIR_NAMES): return True if parts.intersection(EXCLUDE_PATH_PARTS): return True if any(kw in file_path.name.lower() for kw in EXCLUDE_FILENAME_KEYWORDS): return True try: snippet = _sample_head_mid_tail(file_path.read_bytes()) if any(marker in snippet for marker in SENSITIVE_LITERAL_MARKERS): return True if any(pattern.search(snippet) for pattern in SENSITIVE_REGEX_PATTERNS): return True except Exception: return True return False @contextmanager def index_lock(): lock_file = Path(os.getenv('INDEX_LOCK_FILE', str(VAULT_ROOT / '.memory-index.lock'))) lock_file.parent.mkdir(parents=True, exist_ok=True) with open(lock_file, 'w', encoding='utf-8') as fh: fcntl.flock(fh, fcntl.LOCK_EX) try: yield finally: fcntl.flock(fh, fcntl.LOCK_UN) def embed_text(text: str) -> list[float]: api_key = require_env('OPENROUTER_API_KEY') base_url = os.getenv('OPENROUTER_BASE_URL', 'https://openrouter.ai/api/v1').rstrip('/') model = require_env('OPENROUTER_EMBED_MODEL') expected_dim = int(os.getenv('OPENROUTER_EMBED_DIM', '1536')) headers = { 'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json', } referer = os.getenv('OPENROUTER_HTTP_REFERER', '').strip() if referer: headers['HTTP-Referer'] = referer title = os.getenv('OPENROUTER_X_OPENROUTER_TITLE', '').strip() or os.getenv( 'OPENROUTER_X_TITLE', '' ).strip() if title: headers['X-OpenRouter-Title'] = title payload = {'model': model, 'input': text} last_error: Exception | None = None for attempt in range(1, 4): try: response = requests.post( f'{base_url}/embeddings', headers=headers, json=payload, timeout=30, ) response.raise_for_status() body = response.json() data = body.get('data') if not isinstance(data, list) or not data: raise RuntimeError(f'embedding 响应缺少 data 字段: keys={list(body.keys())}') embedding = data[0].get('embedding') if not isinstance(embedding, list): raise RuntimeError('embedding 响应结构异常: data[0].embedding 缺失') if len(embedding) != expected_dim: raise RuntimeError( f'embedding 维度不匹配: got={len(embedding)} expected={expected_dim}' ) return embedding except Exception as error: last_error = error if attempt < 3: time.sleep(0.8 * attempt) raise RuntimeError(f'OpenRouter embedding 请求失败: {last_error}') from last_error