mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
feat: add codex auth modes, background embedding/index repair, and qwenpaw logging (#371)
* feat(codex): add authentication mode support with thread-safe logging - Implement _CodexAuthConfig dataclass for resolved auth settings - Add auth_mode parameter with auto/api_key/oauth options - Separate API key and OAuth authentication flows - Force specific login method based on auth mode - Add explicit API key validation requirement - Serialize concurrent logger initialization in thread lock - Close logging handlers properly during cleanup - Update default config with auth_mode presets for codex and codex_oauth - Add comprehensive tests for authentication modes and concurrent logging * feat(file_store): implement background embedding backfill and keyword index repair - Add _after_embedding_backfill hook in FAISS local file store - Schedule startup embedding repair without delaying component readiness - Cancel and collect embedding backfill task during component shutdown - Log progress at fixed percentage boundaries for long-running operations - Process embedding backfill in configurable batch sizes with progress reporting - Rebuild keyword index in bounded batches with detailed mismatch diagnostics - Format stdlib logs consistently with QwenPaw console output using relative paths - Run embedding backfill as background task that doesn't block component startup - Add comprehensive tests for background embedding and keyword index repair scenarios * fix(file-store): repair graph-chunk consistency on load - Add _repair_graph_chunk_consistency method to detect and fix mismatched graph/chunk states - Clear torn graph/chunk state when missing or orphaned chunks are detected - Ensure keyword index sync handles empty chunks properly - Add comprehensive tests for graph-chunk consistency scenarios - Update test utilities to properly seed graph/chunk snapshots - Increment version to 0.4.1.2 * feat(file_io): enhance list step response format and add comprehensive logging - Format list output with bullet points for better readability - Add explicit "No files found" message when directory is empty - Include detailed timing information for file store startup phases - Add logging for chunk loading, graph consistency checks, and keyword indexing - Provide detailed metrics for embedding backfill operations - Add comprehensive test coverage for empty directory scenarios - Include batch processing statistics for embedding operations * feat(logger): add QwenPaw logging integration with forwarding mechanism - Introduce _ForwardToLoggerHandler to forward log records to target logger - Add qwenpaw logger integration that forwards ReMe logs to QwenPaw handlers - Maintain ReMe logger stability for modules that cache it at import time - Enable QwenPaw handlers to take effect without ReMe reconfiguration - Add comprehensive tests for stdlib forwarding to QwenPaw sinks - Support explicit REME_DISABLE_LOGURU=false to keep original Loguru backend - Preserve existing logging behavior when QwenPaw is not configured * fix(file_store): serialize concurrent FAISS dump operations to prevent corruption - Add asyncio lock to ensure only one FAISS dump operation runs at a time - Generate unique temporary filenames using UUID tokens for atomic replacement - Implement proper cleanup of temporary files in finally block - Add comprehensive test to verify concurrent dumps are serialized - Ensure atomic writes by replacing both index and idmap files together - Prevent partial state writes during concurrent access scenarios
This commit is contained in:
parent
1c08eaa559
commit
cf22ef3b1d
11 changed files with 915 additions and 87 deletions
|
|
@ -1,6 +1,6 @@
|
|||
"""ReMe CLI package."""
|
||||
|
||||
__version__ = "0.4.1.1"
|
||||
__version__ = "0.4.1.2"
|
||||
|
||||
from . import config
|
||||
from . import constants
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import suppress
|
||||
from dataclasses import fields, is_dataclass
|
||||
from dataclasses import dataclass, fields, is_dataclass
|
||||
from enum import Enum
|
||||
import hashlib
|
||||
import json
|
||||
|
|
@ -20,6 +20,15 @@ from ...schema import StreamChunk
|
|||
from ...utils.env_utils import load_env
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CodexAuthConfig:
|
||||
"""Resolved authentication settings for one Codex app-server."""
|
||||
|
||||
mode: str
|
||||
api_key: str = ""
|
||||
base_url: str = ""
|
||||
|
||||
|
||||
@R.register("codex")
|
||||
class CodexAgentWrapper(BaseAgentWrapper):
|
||||
"""Agent wrapper backed by the Codex Python SDK."""
|
||||
|
|
@ -30,6 +39,7 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
self._codex_home = codex_home
|
||||
self._codex: Any | None = None
|
||||
self._codex_config: Any | None = None
|
||||
self._codex_auth_config: _CodexAuthConfig | None = None
|
||||
self._client_lock = asyncio.Lock()
|
||||
self._turn_lock = asyncio.Lock()
|
||||
self._mcp_snapshot_path: Path | None = None
|
||||
|
|
@ -143,8 +153,13 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
def _mcp_config_source(self, kwargs: dict[str, Any]) -> str:
|
||||
return self._explicit_mcp_config(kwargs) or str(self._effective_config_snapshot())
|
||||
|
||||
def _build_client_config(self, kwargs: dict[str, Any]):
|
||||
from openai_codex import CodexConfig
|
||||
def _resolve_auth_config(self, kwargs: dict[str, Any]) -> _CodexAuthConfig:
|
||||
"""Resolve one explicit Codex auth mode without letting OAuth inherit API credentials."""
|
||||
requested_mode = str(kwargs.get("auth_mode") or "auto").lower()
|
||||
if requested_mode not in {"auto", "api_key", "oauth"}:
|
||||
raise ValueError("auth_mode must be one of: auto, api_key, oauth")
|
||||
if requested_mode == "oauth":
|
||||
return _CodexAuthConfig(mode="oauth")
|
||||
|
||||
credential = kwargs.get("credential") if isinstance(kwargs.get("credential"), dict) else {}
|
||||
default_credential = self._default_llm_credential()
|
||||
|
|
@ -156,6 +171,11 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
os.getenv("LLM_API_KEY"),
|
||||
default_credential.get("api_key"),
|
||||
)
|
||||
if requested_mode == "api_key" and not api_key:
|
||||
raise ValueError("auth_mode='api_key' requires a non-empty API key")
|
||||
if not api_key:
|
||||
return _CodexAuthConfig(mode="oauth")
|
||||
|
||||
base_url = self._first_non_empty(
|
||||
kwargs.get("base_url"),
|
||||
credential.get("base_url"),
|
||||
|
|
@ -164,17 +184,22 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
os.getenv("LLM_BASE_URL"),
|
||||
default_credential.get("base_url"),
|
||||
)
|
||||
return _CodexAuthConfig(mode="api_key", api_key=api_key, base_url=base_url)
|
||||
|
||||
def _build_client_config(self, kwargs: dict[str, Any], auth: _CodexAuthConfig | None = None):
|
||||
from openai_codex import CodexConfig
|
||||
|
||||
auth = auth or self._resolve_auth_config(kwargs)
|
||||
|
||||
project_env = self.project_path / ".env"
|
||||
env = load_env(project_env) if project_env.exists() else load_env()
|
||||
self.session_path.mkdir(parents=True, exist_ok=True)
|
||||
env["CODEX_HOME"] = str(self.session_path)
|
||||
if api_key:
|
||||
env["OPENAI_API_KEY"] = api_key
|
||||
|
||||
overrides = list(kwargs.get("config_overrides") or [])
|
||||
if base_url:
|
||||
overrides.append(f"openai_base_url={json.dumps(base_url)}")
|
||||
if auth.base_url:
|
||||
overrides.append(f"openai_base_url={json.dumps(auth.base_url)}")
|
||||
overrides.append(f"forced_login_method={json.dumps('api' if auth.mode == 'api_key' else 'chatgpt')}")
|
||||
return CodexConfig(
|
||||
codex_bin=kwargs.get("codex_bin"),
|
||||
config_overrides=tuple(overrides),
|
||||
|
|
@ -288,15 +313,27 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
"""Lazily start one app-server and reject launch-config changes while it is live."""
|
||||
from openai_codex import AsyncCodex
|
||||
|
||||
config = self._build_client_config(kwargs)
|
||||
auth = self._resolve_auth_config(kwargs)
|
||||
config = self._build_client_config(kwargs, auth)
|
||||
async with self._client_lock:
|
||||
if self._codex is not None:
|
||||
if config != self._codex_config:
|
||||
if config != self._codex_config or auth != self._codex_auth_config:
|
||||
raise RuntimeError("Codex client configuration changed; close the wrapper before reconfiguring it")
|
||||
return self._codex
|
||||
codex = AsyncCodex(config)
|
||||
try:
|
||||
if auth.mode == "api_key":
|
||||
await codex.login_api_key(auth.api_key)
|
||||
else:
|
||||
account = await codex.account()
|
||||
if account.account is None:
|
||||
raise RuntimeError(f"No ChatGPT OAuth login found in CODEX_HOME: {self.session_path}")
|
||||
except BaseException:
|
||||
await codex.close()
|
||||
raise
|
||||
self._codex = codex
|
||||
self._codex_config = config
|
||||
self._codex_auth_config = auth
|
||||
return codex
|
||||
|
||||
async def _close(self) -> None:
|
||||
|
|
@ -304,6 +341,7 @@ class CodexAgentWrapper(BaseAgentWrapper):
|
|||
async with self._client_lock:
|
||||
codex, self._codex = self._codex, None
|
||||
self._codex_config = None
|
||||
self._codex_auth_config = None
|
||||
self._thread_tool_contexts.clear()
|
||||
try:
|
||||
if codex is not None:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""FAISS-backed file store: chunk JSONL stays authoritative; FAISS replaces the linear vector scan."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from uuid import uuid4
|
||||
|
||||
import aiofiles
|
||||
import numpy as np
|
||||
|
|
@ -40,6 +42,7 @@ class FaissLocalFileStore(LocalFileStore):
|
|||
self._id_map: list[str] = [] # row -> chunk_id
|
||||
self._id_to_row: dict[str, int] = {} # chunk_id -> row (live entries only)
|
||||
self._tombstones: set[int] = set() # rows whose chunk_id was deleted
|
||||
self._faiss_dump_lock = asyncio.Lock()
|
||||
|
||||
@staticmethod
|
||||
def _import_faiss():
|
||||
|
|
@ -106,6 +109,10 @@ class FaissLocalFileStore(LocalFileStore):
|
|||
if len(self._tombstones) >= self.max_tombstones:
|
||||
self._rebuild_index()
|
||||
|
||||
async def _after_embedding_backfill(self) -> None:
|
||||
"""Make newly backfilled vectors visible to FAISS before persistence."""
|
||||
self._rebuild_index()
|
||||
|
||||
# -- persistence ------------------------------------------------------
|
||||
|
||||
async def load(self) -> None:
|
||||
|
|
@ -157,26 +164,33 @@ class FaissLocalFileStore(LocalFileStore):
|
|||
|
||||
async def dump(self) -> None:
|
||||
"""Persist chunks JSONL via the parent, then write the FAISS sidecar atomically."""
|
||||
await super().dump()
|
||||
if self._faiss_index is None or self.embedding_store is None:
|
||||
return
|
||||
try:
|
||||
self._compact_if_needed()
|
||||
await self._write_sidecar()
|
||||
self.logger.info(f"Saved FAISS index: {self._faiss_index.ntotal} vectors to {self.faiss_path}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write FAISS index: {e}")
|
||||
async with self._faiss_dump_lock:
|
||||
await super().dump()
|
||||
if self._faiss_index is None or self.embedding_store is None:
|
||||
return
|
||||
try:
|
||||
self._compact_if_needed()
|
||||
await self._write_sidecar()
|
||||
self.logger.info(f"Saved FAISS index: {self._faiss_index.ntotal} vectors to {self.faiss_path}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write FAISS index: {e}")
|
||||
|
||||
async def _write_sidecar(self) -> None:
|
||||
tmp_index = self.faiss_path.with_suffix(".tmp")
|
||||
self._faiss.write_index(self._faiss_index, str(tmp_index))
|
||||
tmp_index.replace(self.faiss_path)
|
||||
token = uuid4().hex
|
||||
tmp_index = self.faiss_path.with_name(f".{self.faiss_path.name}.{token}.tmp")
|
||||
tmp_idmap = self.faiss_idmap_path.with_name(f".{self.faiss_idmap_path.name}.{token}.tmp")
|
||||
payload = json.dumps({"id_map": list(self._id_map), "tombstones": sorted(self._tombstones)})
|
||||
try:
|
||||
self._faiss.write_index(self._faiss_index, str(tmp_index))
|
||||
async with aiofiles.open(tmp_idmap, "w", encoding=self.encoding) as f:
|
||||
await f.write(payload)
|
||||
|
||||
tmp_idmap = self.faiss_idmap_path.with_suffix(".tmp")
|
||||
payload = json.dumps({"id_map": self._id_map, "tombstones": sorted(self._tombstones)})
|
||||
async with aiofiles.open(tmp_idmap, "w", encoding=self.encoding) as f:
|
||||
await f.write(payload)
|
||||
tmp_idmap.replace(self.faiss_idmap_path)
|
||||
# Publish only after both parts of the sidecar have been written successfully.
|
||||
tmp_index.replace(self.faiss_path)
|
||||
tmp_idmap.replace(self.faiss_idmap_path)
|
||||
finally:
|
||||
tmp_index.unlink(missing_ok=True)
|
||||
tmp_idmap.unlink(missing_ok=True)
|
||||
|
||||
# -- CRUD overrides ---------------------------------------------------
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
"""In-memory file store with compressed JSONL persistence on close."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import datetime
|
||||
import heapq
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from contextlib import suppress
|
||||
|
||||
|
|
@ -23,6 +25,8 @@ CachedEmbedding = tuple[str, np.ndarray]
|
|||
_EMBEDDING_F16_B64_FIELD = "_embedding_f16_b64"
|
||||
_EMBEDDING_F16_DTYPE = np.dtype("<f2")
|
||||
_VECTOR_SEARCH_BATCH_SIZE = 1024
|
||||
_PROGRESS_LOG_PERCENT_STEP = 10
|
||||
_KEYWORD_REBUILD_BATCH_SIZE = 200
|
||||
|
||||
|
||||
@R.register("local")
|
||||
|
|
@ -62,15 +66,35 @@ class LocalFileStore(BaseFileStore):
|
|||
self.store_version = store_version
|
||||
self.file_chunks: dict[str, FileChunk] = {}
|
||||
self.chunks_path = self.component_metadata_path / f"file_chunks_{self.name}_{self.store_version}.jsonl.zst"
|
||||
self._embedding_backfill_task: asyncio.Task | None = None
|
||||
|
||||
# -- lifecycle ------------------------------------------------------------
|
||||
|
||||
async def _start(self) -> None:
|
||||
started_at = time.monotonic()
|
||||
self.component_metadata_path.mkdir(parents=True, exist_ok=True)
|
||||
await super()._start()
|
||||
|
||||
load_started_at = time.monotonic()
|
||||
await self.load()
|
||||
self.logger.info(
|
||||
f"{self.name}: file store load complete: chunks={len(self.file_chunks)}, "
|
||||
f"elapsed={time.monotonic() - load_started_at:.3f}s",
|
||||
)
|
||||
|
||||
backfill_started_at = time.monotonic()
|
||||
self._start_embedding_backfill()
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill scheduling complete: "
|
||||
f"scheduled={self._embedding_backfill_task is not None}, "
|
||||
f"elapsed={time.monotonic() - backfill_started_at:.3f}s",
|
||||
)
|
||||
self.logger.info(
|
||||
f"{self.name}: file store startup complete: " f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
|
||||
async def _close(self) -> None:
|
||||
await self._cancel_embedding_backfill()
|
||||
await self.dump()
|
||||
self.file_chunks.clear()
|
||||
await super()._close()
|
||||
|
|
@ -110,20 +134,69 @@ class LocalFileStore(BaseFileStore):
|
|||
|
||||
async def load(self) -> None:
|
||||
"""Load chunks from the JSONL file into memory; missing file is a no-op."""
|
||||
if not self.chunks_path.exists():
|
||||
return
|
||||
try:
|
||||
for line in read_jsonl_zst(self.chunks_path, self.encoding):
|
||||
line = line.strip()
|
||||
if line:
|
||||
chunk = self._deserialize_chunk(line)
|
||||
self.file_chunks[chunk.id] = chunk
|
||||
self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}")
|
||||
self._invalidate_stale_embeddings()
|
||||
await self._sync_keyword_index_from_chunks()
|
||||
await self._backfill_missing_embeddings()
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self.chunks_path}: {e}")
|
||||
started_at = time.monotonic()
|
||||
chunk_load_started_at = time.monotonic()
|
||||
if self.chunks_path.exists():
|
||||
try:
|
||||
for line in read_jsonl_zst(self.chunks_path, self.encoding):
|
||||
line = line.strip()
|
||||
if line:
|
||||
chunk = self._deserialize_chunk(line)
|
||||
self.file_chunks[chunk.id] = chunk
|
||||
self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self.chunks_path}: {e}")
|
||||
self.logger.info(
|
||||
f"{self.name}: chunk store load complete: chunks={len(self.file_chunks)}, "
|
||||
f"elapsed={time.monotonic() - chunk_load_started_at:.3f}s",
|
||||
)
|
||||
|
||||
graph_repair_started_at = time.monotonic()
|
||||
graph_repaired = await self._repair_graph_chunk_consistency()
|
||||
self.logger.info(
|
||||
f"{self.name}: graph consistency check complete: repaired={graph_repaired}, "
|
||||
f"elapsed={time.monotonic() - graph_repair_started_at:.3f}s",
|
||||
)
|
||||
|
||||
keyword_sync_started_at = time.monotonic()
|
||||
await self._sync_keyword_index_from_chunks()
|
||||
keyword_backend = type(self.keyword_index).__name__ if self.keyword_index is not None else "disabled"
|
||||
keyword_docs = getattr(self.keyword_index, "n_docs", 0) if self.keyword_index is not None else 0
|
||||
self.logger.info(
|
||||
f"{self.name}: BM25/keyword index sync complete: backend={keyword_backend}, docs={keyword_docs}, "
|
||||
f"elapsed={time.monotonic() - keyword_sync_started_at:.3f}s",
|
||||
)
|
||||
self.logger.info(
|
||||
f"{self.name}: file store load phases complete: " f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
|
||||
async def _repair_graph_chunk_consistency(self) -> bool:
|
||||
"""Clear torn graph/chunk state so the filesystem scan rebuilds it.
|
||||
|
||||
``InitChangesStep`` uses file-graph nodes as the indexed-file snapshot.
|
||||
A graph that survives a missing, truncated, or stale chunk store would
|
||||
otherwise make the source files look up to date and permanently hide
|
||||
the broken search index.
|
||||
"""
|
||||
assert self.file_graph is not None
|
||||
nodes = await self.file_graph.get_nodes()
|
||||
graph_chunk_ids = {chunk_id for node in nodes for chunk_id in node.chunk_ids}
|
||||
stored_chunk_ids = set(self.file_chunks)
|
||||
missing = graph_chunk_ids - stored_chunk_ids
|
||||
orphaned = stored_chunk_ids - graph_chunk_ids
|
||||
if not missing and not orphaned:
|
||||
return False
|
||||
|
||||
self.logger.warning(
|
||||
f"{self.name}: graph/chunk mismatch: nodes={len(nodes)}, graph_chunks={len(graph_chunk_ids)}, "
|
||||
f"stored_chunks={len(stored_chunk_ids)}, missing={len(missing)}, orphaned={len(orphaned)}; "
|
||||
"clearing derived index state for automatic rebuild",
|
||||
)
|
||||
# Clearing graph nodes is required: the next InitChangesStep scan will
|
||||
# then classify every watched source file as added and rebuild graph,
|
||||
# chunks, and search indexes from the user-owned files.
|
||||
await self.clear()
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _deserialize_chunk(line: str) -> FileChunk:
|
||||
|
|
@ -155,51 +228,203 @@ class LocalFileStore(BaseFileStore):
|
|||
return
|
||||
self._drop_stale_embeddings(self.file_chunks.values(), "load")
|
||||
|
||||
async def _backfill_missing_embeddings(self) -> None:
|
||||
"""Embed persisted chunks that predate embedding being enabled."""
|
||||
if not self.embedding_store or not self.file_chunks:
|
||||
def _start_embedding_backfill(self) -> None:
|
||||
"""Schedule startup embedding repair without delaying component readiness."""
|
||||
started_at = time.monotonic()
|
||||
if not self.embedding_store:
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill skipped: reason=embedding_disabled, "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
return
|
||||
|
||||
missing = [chunk for chunk in self.file_chunks.values() if chunk.text and chunk.embedding is None]
|
||||
if not missing:
|
||||
if not self.file_chunks:
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill skipped: reason=no_chunks, "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
return
|
||||
|
||||
self.logger.info(f"{self.name}: backfilling embeddings for {len(missing)} chunks")
|
||||
if not await self.embedding_store.health_check():
|
||||
self._disable_embedding("backfill health check failed")
|
||||
if self._embedding_backfill_task is not None and not self._embedding_backfill_task.done():
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill scheduling skipped: reason=already_running, "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
return
|
||||
self._embedding_backfill_task = asyncio.create_task(
|
||||
self._backfill_missing_embeddings(),
|
||||
name=f"embedding-backfill:{self.name}",
|
||||
)
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill scheduled: chunks={len(self.file_chunks)}, "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
|
||||
async def _cancel_embedding_backfill(self) -> None:
|
||||
"""Cancel and collect the startup repair task during component shutdown."""
|
||||
task = self._embedding_backfill_task
|
||||
self._embedding_backfill_task = None
|
||||
if task is None:
|
||||
return
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await self.embedding_store.get_node_embeddings(missing)
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
self.logger.exception(f"{self.name}: embedding backfill task failed during shutdown")
|
||||
|
||||
def _log_progress(self, operation: str, current: int, total: int, next_percent: int) -> int:
|
||||
"""Log progress at fixed percentage boundaries and return the next boundary."""
|
||||
if total <= 0:
|
||||
return 100
|
||||
percent = min(100, current * 100 // total)
|
||||
if current < total and percent < next_percent:
|
||||
return next_percent
|
||||
self.logger.info(f"{self.name}: {operation} progress: {current}/{total} ({percent}%)")
|
||||
while next_percent <= percent:
|
||||
next_percent += _PROGRESS_LOG_PERCENT_STEP
|
||||
return next_percent
|
||||
|
||||
async def _backfill_missing_embeddings(self) -> None:
|
||||
"""Background-repair persisted chunks that do not have usable vectors."""
|
||||
started_at = time.monotonic()
|
||||
if not self.embedding_store or not self.file_chunks:
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill finished without work: "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
return
|
||||
|
||||
scan_started_at = time.monotonic()
|
||||
self._invalidate_stale_embeddings()
|
||||
missing = [chunk for chunk in self.file_chunks.values() if chunk.text and chunk.embedding is None]
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill scan complete: chunks={len(self.file_chunks)}, "
|
||||
f"missing={len(missing)}, elapsed={time.monotonic() - scan_started_at:.3f}s",
|
||||
)
|
||||
if not missing:
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill complete: filled=0/0, "
|
||||
f"elapsed={time.monotonic() - started_at:.3f}s",
|
||||
)
|
||||
return
|
||||
|
||||
total = len(missing)
|
||||
batch_size = max(1, int(getattr(self.embedding_store, "max_batch_size", 10)))
|
||||
self.logger.info(f"{self.name}: embedding backfill started: total={total}, batch_size={batch_size}")
|
||||
try:
|
||||
health_check_started_at = time.monotonic()
|
||||
is_healthy = await self.embedding_store.health_check()
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding health check complete: healthy={is_healthy}, "
|
||||
f"elapsed={time.monotonic() - health_check_started_at:.3f}s",
|
||||
)
|
||||
if not is_healthy:
|
||||
self._disable_embedding("backfill health check failed")
|
||||
self.logger.warning(
|
||||
f"{self.name}: embedding backfill failed: processed=0/{total}, reason=health check failed",
|
||||
)
|
||||
return
|
||||
|
||||
processed = 0
|
||||
batch_count = 0
|
||||
embedding_started_at = time.monotonic()
|
||||
next_percent = _PROGRESS_LOG_PERCENT_STEP
|
||||
for start in range(0, total, batch_size):
|
||||
batch = missing[start : start + batch_size]
|
||||
await self.embedding_store.get_node_embeddings(batch)
|
||||
self._drop_stale_embeddings(batch, "backfill")
|
||||
processed += len(batch)
|
||||
batch_count += 1
|
||||
next_percent = self._log_progress("embedding backfill", processed, total, next_percent)
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding batches complete: processed={processed}/{total}, "
|
||||
f"batches={batch_count}, elapsed={time.monotonic() - embedding_started_at:.3f}s",
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
elapsed = time.monotonic() - started_at
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill cancelled: processed={processed if 'processed' in locals() else 0}/"
|
||||
f"{total}, elapsed={elapsed:.2f}s",
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
self._disable_embedding(f"backfill: {type(e).__name__}: {e}")
|
||||
elapsed = time.monotonic() - started_at
|
||||
self.logger.exception(
|
||||
f"{self.name}: embedding backfill failed: processed={processed if 'processed' in locals() else 0}/"
|
||||
f"{total}, elapsed={elapsed:.2f}s",
|
||||
)
|
||||
return
|
||||
|
||||
self._drop_stale_embeddings(missing, "backfill")
|
||||
filled = sum(1 for chunk in missing if chunk.embedding is not None)
|
||||
elapsed = time.monotonic() - started_at
|
||||
self.logger.info(
|
||||
f"{self.name}: embedding backfill complete: filled={filled}/{total}, elapsed={elapsed:.2f}s",
|
||||
)
|
||||
if filled:
|
||||
self.logger.info(f"{self.name}: backfilled embeddings for {filled}/{len(missing)} chunks")
|
||||
await self.dump()
|
||||
try:
|
||||
await self._after_embedding_backfill()
|
||||
await self.dump()
|
||||
except Exception:
|
||||
self.logger.exception(f"{self.name}: failed to persist completed embedding backfill")
|
||||
|
||||
async def _after_embedding_backfill(self) -> None:
|
||||
"""Backend hook for refreshing derived vector indexes after backfill."""
|
||||
|
||||
async def _sync_keyword_index_from_chunks(self) -> None:
|
||||
"""Repair keyword index when its persisted state does not match chunks."""
|
||||
if not self.keyword_index or not self.file_chunks:
|
||||
if not self.keyword_index:
|
||||
return
|
||||
|
||||
docs = {cid: chunk.text for cid, chunk in self.file_chunks.items() if chunk.text}
|
||||
if not docs:
|
||||
return
|
||||
|
||||
expected_ids = docs.keys()
|
||||
expected_ids = set(docs)
|
||||
live_ids = None
|
||||
with suppress(Exception):
|
||||
live_ids = self.keyword_index.document_ids
|
||||
live_ids = set(self.keyword_index.document_ids)
|
||||
|
||||
if live_ids == expected_ids:
|
||||
return
|
||||
|
||||
self.logger.warning(f"{self.name}: keyword index mismatch with chunks; rebuilding {len(docs)} docs")
|
||||
await self.keyword_index.reset_index(docs)
|
||||
missing_count = len(expected_ids - live_ids) if live_ids is not None else len(expected_ids)
|
||||
extra_count = len(live_ids - expected_ids) if live_ids is not None else -1
|
||||
indexed_count = len(live_ids) if live_ids is not None else getattr(self.keyword_index, "n_docs", -1)
|
||||
self.logger.warning(
|
||||
f"{self.name}: keyword index mismatch: indexed={indexed_count}, expected={len(expected_ids)}, "
|
||||
f"missing={missing_count}, extra={extra_count}; rebuilding",
|
||||
)
|
||||
await self._rebuild_keyword_index(docs)
|
||||
|
||||
async def _rebuild_keyword_index(self, docs: dict[str, str]) -> None:
|
||||
"""Synchronously rebuild keyword search in batches with progress logs."""
|
||||
assert self.keyword_index is not None
|
||||
total = len(docs)
|
||||
started_at = time.monotonic()
|
||||
self.logger.info(
|
||||
f"{self.name}: keyword index rebuild started: total={total}, batch_size={_KEYWORD_REBUILD_BATCH_SIZE}",
|
||||
)
|
||||
try:
|
||||
# All built-in keyword indexes support clear/add/dump. Keep a fallback
|
||||
# for third-party implementations that only expose reset_index.
|
||||
if not all(hasattr(self.keyword_index, method) for method in ("clear", "add_docs", "dump")):
|
||||
await self.keyword_index.reset_index(docs)
|
||||
self._log_progress("keyword index rebuild", total, total, _PROGRESS_LOG_PERCENT_STEP)
|
||||
else:
|
||||
await self.keyword_index.clear()
|
||||
items = list(docs.items())
|
||||
next_percent = _PROGRESS_LOG_PERCENT_STEP
|
||||
for start in range(0, total, _KEYWORD_REBUILD_BATCH_SIZE):
|
||||
batch = dict(items[start : start + _KEYWORD_REBUILD_BATCH_SIZE])
|
||||
await self.keyword_index.add_docs(batch)
|
||||
current = min(start + len(batch), total)
|
||||
next_percent = self._log_progress("keyword index rebuild", current, total, next_percent)
|
||||
await self.keyword_index.dump()
|
||||
except Exception:
|
||||
elapsed = time.monotonic() - started_at
|
||||
self.logger.exception(f"{self.name}: keyword index rebuild failed after {elapsed:.2f}s")
|
||||
raise
|
||||
elapsed = time.monotonic() - started_at
|
||||
self.logger.info(f"{self.name}: keyword index rebuild complete: total={total}, elapsed={elapsed:.2f}s")
|
||||
|
||||
async def dump(self) -> None:
|
||||
"""Atomically rewrite the JSONL, then cascade dump into keyword_index and file_graph."""
|
||||
|
|
|
|||
|
|
@ -675,6 +675,7 @@ components:
|
|||
system_prompt_mode: replace
|
||||
codex:
|
||||
backend: codex
|
||||
auth_mode: api_key
|
||||
model: ${CODEX_MODEL_NAME:-}
|
||||
api_key: ${CODEX_API_KEY:-}
|
||||
base_url: ${CODEX_BASE_URL:-}
|
||||
|
|
@ -682,6 +683,7 @@ components:
|
|||
sandbox: full-access
|
||||
codex_oauth:
|
||||
backend: codex
|
||||
auth_mode: oauth
|
||||
model: ${CODEX_MODEL_NAME:-}
|
||||
codex_home: ${CODEX_HOME:-~/.codex}
|
||||
approval_mode: auto_review
|
||||
|
|
|
|||
|
|
@ -92,7 +92,12 @@ class ListStep(BaseStep):
|
|||
items = self._format_relative(self._walk_files(target_dir, recursive, limit), workspace_dir)
|
||||
|
||||
self.context.response.success = True
|
||||
self.context.response.answer = f"Listed {len(items)} file(s) under {path or '.'}"
|
||||
location = path or "."
|
||||
if items:
|
||||
rendered_items = "\n".join(f"- {item}" for item in items)
|
||||
self.context.response.answer = f"Listed {len(items)} file(s) under {location}:\n{rendered_items}"
|
||||
else:
|
||||
self.context.response.answer = f"No files found under {location}."
|
||||
self.context.response.metadata.update({"items": items, "count": len(items)})
|
||||
self.logger.info(
|
||||
f"[{self.name}] listed dir={target_dir} recursive={recursive} count={len(items)} limit={limit}",
|
||||
|
|
|
|||
|
|
@ -3,14 +3,48 @@
|
|||
import logging
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from logging.handlers import TimedRotatingFileHandler
|
||||
|
||||
_logger = None
|
||||
_logger_lock = threading.RLock()
|
||||
|
||||
_LOGURU_FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}"
|
||||
_STDLIB_FORMAT = "%(asctime)s | %(levelname)s | %(filename)s:%(lineno)d | %(funcName)s | %(message)s"
|
||||
_STDLIB_FORMAT = "%(levelname)s %(source_path)s:%(lineno)d | %(asctime)s | %(message)s"
|
||||
_STDLIB_DATEFMT = "%Y-%m-%d %H:%M:%S"
|
||||
_QWENPAW_LOGGER_NAME = "qwenpaw"
|
||||
|
||||
|
||||
class _ForwardToLoggerHandler(logging.Handler):
|
||||
"""Forward records to a host logger without borrowing its handlers."""
|
||||
|
||||
def __init__(self, target_name: str) -> None:
|
||||
super().__init__()
|
||||
self.target_name = target_name
|
||||
|
||||
def emit(self, record: logging.LogRecord) -> None:
|
||||
target = logging.getLogger(self.target_name)
|
||||
if target.isEnabledFor(record.levelno):
|
||||
target.handle(record)
|
||||
|
||||
|
||||
class _QwenPawStdlibFormatter(logging.Formatter):
|
||||
"""Format stdlib records consistently with QwenPaw host logs."""
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
source_path = record.pathname
|
||||
cwd = os.getcwd()
|
||||
try:
|
||||
if os.path.commonpath([source_path, cwd]) == cwd:
|
||||
source_path = os.path.relpath(source_path, cwd)
|
||||
except ValueError:
|
||||
# Paths on different Windows drives cannot be compared.
|
||||
pass
|
||||
|
||||
# QwenPaw prefixes console records with a cwd-relative source path.
|
||||
record.source_path = source_path
|
||||
return super().format(record)
|
||||
|
||||
|
||||
def _enable_loguru() -> bool:
|
||||
|
|
@ -53,13 +87,25 @@ def _init_loguru(log_dir: str, level: str, log_to_console: bool, log_to_file: bo
|
|||
|
||||
def _init_stdlib(log_dir: str, level: str, log_to_console: bool, log_to_file: bool):
|
||||
logger = logging.getLogger("reme")
|
||||
logger.setLevel(level)
|
||||
logger.propagate = False
|
||||
|
||||
for handler in list(logger.handlers):
|
||||
logger.removeHandler(handler)
|
||||
handler.close()
|
||||
|
||||
formatter = logging.Formatter(_STDLIB_FORMAT, datefmt=_STDLIB_DATEFMT)
|
||||
qwenpaw_logger = logging.getLogger(_QWENPAW_LOGGER_NAME)
|
||||
if qwenpaw_logger.handlers:
|
||||
# QwenPaw owns the screen and file handlers. Forwarding keeps ReMe's
|
||||
# logger object stable for modules that cache it at import time, while
|
||||
# allowing future QwenPaw handlers (for example qwenpaw.log) to take
|
||||
# effect without another ReMe reconfiguration.
|
||||
logger.setLevel(logging.DEBUG)
|
||||
logger.addHandler(_ForwardToLoggerHandler(_QWENPAW_LOGGER_NAME))
|
||||
return logger
|
||||
|
||||
logger.setLevel(level)
|
||||
|
||||
formatter = _QwenPawStdlibFormatter(_STDLIB_FORMAT, datefmt=_STDLIB_DATEFMT)
|
||||
|
||||
if log_to_console:
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
|
|
@ -98,11 +144,17 @@ def get_logger(
|
|||
"""Return the global logger, initializing sinks on first call (or when force_init)."""
|
||||
global _logger
|
||||
|
||||
if _logger is not None and not force_init:
|
||||
return _logger
|
||||
# ReMe can be embedded multiple times in one process. Hosts may construct
|
||||
# those applications concurrently, while both logging backends reconfigure
|
||||
# a process-global logger via a remove-then-add sequence. Keep the whole
|
||||
# check/reconfigure/publish transaction atomic so concurrent force_init
|
||||
# calls cannot leave duplicate sinks or handlers behind.
|
||||
with _logger_lock:
|
||||
if _logger is not None and not force_init:
|
||||
return _logger
|
||||
|
||||
if _enable_loguru():
|
||||
_logger = _init_loguru(log_dir, level, log_to_console, log_to_file)
|
||||
else:
|
||||
_logger = _init_stdlib(log_dir, level, log_to_console, log_to_file)
|
||||
return _logger
|
||||
if _enable_loguru():
|
||||
_logger = _init_loguru(log_dir, level, log_to_console, log_to_file)
|
||||
else:
|
||||
_logger = _init_stdlib(log_dir, level, log_to_console, log_to_file)
|
||||
return _logger
|
||||
|
|
|
|||
|
|
@ -304,7 +304,7 @@ class _TurnResult:
|
|||
|
||||
|
||||
def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch):
|
||||
wrapper, _job = _wrapper(tmp_path)
|
||||
wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
|
||||
|
||||
class FakeThread:
|
||||
id = "thread-1"
|
||||
|
|
@ -327,6 +327,9 @@ def test_reply_returns_thread_id_and_structured_output(tmp_path, monkeypatch):
|
|||
async def close(self):
|
||||
return None
|
||||
|
||||
async def account(self):
|
||||
return SimpleNamespace(account=SimpleNamespace())
|
||||
|
||||
async def thread_start(self, **_kwargs):
|
||||
return FakeThread()
|
||||
|
||||
|
|
@ -553,7 +556,7 @@ def test_output_schema_normalizes_model_class_and_preserves_dict(tmp_path):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, monkeypatch):
|
||||
wrapper, _job = _wrapper(tmp_path)
|
||||
wrapper, _job = _wrapper(tmp_path, auth_mode="oauth")
|
||||
clients = []
|
||||
observed_schemas = []
|
||||
close_count = 0
|
||||
|
|
@ -576,6 +579,9 @@ async def test_reply_normalizes_schema_and_reuses_persistent_client(tmp_path, mo
|
|||
nonlocal close_count
|
||||
close_count += 1
|
||||
|
||||
async def account(self):
|
||||
return SimpleNamespace(account=SimpleNamespace())
|
||||
|
||||
async def thread_start(self, **_kwargs):
|
||||
return FakeThread()
|
||||
|
||||
|
|
@ -657,6 +663,9 @@ async def test_persistent_client_rejects_launch_config_changes(tmp_path, monkeyp
|
|||
def __init__(self, _config):
|
||||
pass
|
||||
|
||||
async def login_api_key(self, _api_key):
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
|
@ -671,6 +680,74 @@ async def test_persistent_client_rejects_launch_config_changes(tmp_path, monkeyp
|
|||
await wrapper.close()
|
||||
|
||||
|
||||
def test_oauth_mode_ignores_api_credentials_and_forces_chatgpt(tmp_path, monkeypatch):
|
||||
wrapper, _job = _wrapper(tmp_path)
|
||||
monkeypatch.setenv("CODEX_API_KEY", "ambient-key")
|
||||
monkeypatch.setenv("CODEX_BASE_URL", "https://ambient.example.test/v1")
|
||||
monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {})
|
||||
|
||||
auth = wrapper._resolve_auth_config( # pylint: disable=protected-access
|
||||
{
|
||||
"auth_mode": "oauth",
|
||||
"api_key": "explicit-key",
|
||||
"base_url": "https://explicit.example.test/v1",
|
||||
},
|
||||
)
|
||||
config = wrapper._build_client_config({}, auth) # pylint: disable=protected-access
|
||||
|
||||
assert auth.mode == "oauth"
|
||||
assert auth.api_key == ""
|
||||
assert auth.base_url == ""
|
||||
assert "OPENAI_API_KEY" not in config.env
|
||||
assert 'forced_login_method="chatgpt"' in config.config_overrides
|
||||
assert not any(value.startswith("openai_base_url=") for value in config.config_overrides)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_mode_logs_in_app_server_explicitly(tmp_path, monkeypatch):
|
||||
wrapper, _job = _wrapper(tmp_path)
|
||||
observed = {}
|
||||
|
||||
class FakeCodex:
|
||||
def __init__(self, config):
|
||||
observed["config"] = config
|
||||
|
||||
async def login_api_key(self, api_key):
|
||||
observed["api_key"] = api_key
|
||||
|
||||
async def close(self):
|
||||
observed["closed"] = True
|
||||
|
||||
monkeypatch.setattr("openai_codex.AsyncCodex", FakeCodex)
|
||||
monkeypatch.setattr("reme.components.agent_wrapper.codex_agent_wrapper.load_env", lambda *_args: {})
|
||||
|
||||
await wrapper.start()
|
||||
await wrapper._get_codex( # pylint: disable=protected-access
|
||||
{
|
||||
"auth_mode": "api_key",
|
||||
"api_key": "explicit-key",
|
||||
"base_url": "https://proxy.example.test/v1",
|
||||
},
|
||||
)
|
||||
await wrapper.close()
|
||||
|
||||
config = observed["config"]
|
||||
assert observed["api_key"] == "explicit-key"
|
||||
assert "OPENAI_API_KEY" not in config.env
|
||||
assert 'openai_base_url="https://proxy.example.test/v1"' in config.config_overrides
|
||||
assert 'forced_login_method="api"' in config.config_overrides
|
||||
assert observed["closed"] is True
|
||||
|
||||
|
||||
def test_api_key_mode_requires_key(tmp_path, monkeypatch):
|
||||
wrapper, _job = _wrapper(tmp_path)
|
||||
for name in ("CODEX_API_KEY", "OPENAI_API_KEY", "LLM_API_KEY"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
with pytest.raises(ValueError, match="requires a non-empty API key"):
|
||||
wrapper._resolve_auth_config({"auth_mode": "api_key"}) # pylint: disable=protected-access
|
||||
|
||||
|
||||
@pytest.mark.parametrize("review_status", ["approved", "denied"])
|
||||
def test_event_to_chunks_maps_approval_started_and_completed(review_status):
|
||||
action = {"type": "futureApprovalAction", "value": "preserved"}
|
||||
|
|
@ -724,7 +801,9 @@ def test_default_config_provides_codex_oauth_wrapper(monkeypatch):
|
|||
oauth = config["components"]["agent_wrapper"]["codex_oauth"]
|
||||
codex = config["components"]["agent_wrapper"]["codex"]
|
||||
assert oauth["backend"] == "codex"
|
||||
assert oauth["auth_mode"] == "oauth"
|
||||
assert oauth["codex_home"] == "~/.codex"
|
||||
assert oauth["sandbox"] == "full-access"
|
||||
assert "api_key" not in oauth
|
||||
assert codex["auth_mode"] == "api_key"
|
||||
assert codex["sandbox"] == "full-access"
|
||||
|
|
|
|||
|
|
@ -168,12 +168,32 @@ def test_list_lists_files():
|
|||
payload = _metadata(step)
|
||||
assert set(payload["items"]) == {"topics/a.md", "topics/b.md", "topics/sub/c.md"}
|
||||
assert payload["count"] == 3
|
||||
answer = step.context.response.answer
|
||||
assert answer.startswith("Listed 3 file(s) under topics:\n")
|
||||
assert all(f"- {item}" in answer for item in payload["items"])
|
||||
await store.close()
|
||||
print("✓ test_list_lists_files passed")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_list_empty_directory_has_explicit_answer():
|
||||
"""list tells an LLM explicitly when the target directory has no files."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = await _make_store()
|
||||
(Path(tmp) / "topics").mkdir()
|
||||
step = crud_list.ListStep(file_store=store)
|
||||
await step(path="topics")
|
||||
assert step.context.response.answer == "No files found under topics."
|
||||
assert _metadata(step) == {"items": [], "count": 0}
|
||||
await store.close()
|
||||
print("✓ test_list_empty_directory_has_explicit_answer passed")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_list_respects_limit_and_non_recursive():
|
||||
"""Non-recursive list ignores subdirs; limit caps the count."""
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ class FakeEmbeddingStore:
|
|||
"""Small deterministic embedding provider used by file-store tests."""
|
||||
|
||||
dimensions = 2
|
||||
max_batch_size = 10
|
||||
|
||||
def _embed(self, text: str) -> np.ndarray:
|
||||
if "beta" in text or "fresh" in text:
|
||||
|
|
@ -87,6 +88,19 @@ class HealthCountingEmbeddingStore(FakeEmbeddingStore):
|
|||
return True
|
||||
|
||||
|
||||
class BlockingEmbeddingStore(FakeEmbeddingStore):
|
||||
"""Fake provider that proves startup does not await remote backfill."""
|
||||
|
||||
def __init__(self):
|
||||
self.started = asyncio.Event()
|
||||
self.release = asyncio.Event()
|
||||
|
||||
async def get_node_embeddings(self, nodes: list[FileChunk], **kwargs) -> list[FileChunk]:
|
||||
self.started.set()
|
||||
await self.release.wait()
|
||||
return await super().get_node_embeddings(nodes, **kwargs)
|
||||
|
||||
|
||||
class WrongDimEmbeddingStore(FakeEmbeddingStore):
|
||||
"""Fake embedding store that returns vectors with the wrong dimension."""
|
||||
|
||||
|
|
@ -131,6 +145,20 @@ def chunk(chunk_id: str, path: str, text: str, **metadata) -> FileChunk:
|
|||
return FileChunk(id=chunk_id, path=path, text=text, start_line=1, end_line=1, metadata=metadata)
|
||||
|
||||
|
||||
async def set_chunks_with_graph(store: LocalFileStore, chunks: dict[str, FileChunk]) -> None:
|
||||
"""Seed a graph/chunk snapshot that satisfies the persistence invariant."""
|
||||
store.file_chunks = chunks
|
||||
chunk_ids_by_path: dict[str, list[str]] = {}
|
||||
for chunk_node in chunks.values():
|
||||
chunk_ids_by_path.setdefault(chunk_node.path, []).append(chunk_node.id)
|
||||
nodes = []
|
||||
for path, chunk_ids in chunk_ids_by_path.items():
|
||||
file_node = node(path)
|
||||
file_node.chunk_ids = chunk_ids
|
||||
nodes.append(file_node)
|
||||
await store.file_graph.upsert_nodes(nodes)
|
||||
|
||||
|
||||
def test_keyword_only_upsert_removes_old_chunks_and_docs():
|
||||
"""Keyword-only upsert removes stale chunks and keyword documents."""
|
||||
|
||||
|
|
@ -199,6 +227,77 @@ def test_load_rebuilds_keyword_index_from_persisted_chunks_when_missing():
|
|||
run(go())
|
||||
|
||||
|
||||
def test_load_clears_graph_when_persisted_chunks_are_missing():
|
||||
"""A surviving graph must not hide a missing chunk store from reindex."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
seed = LocalFileStore(name="t_missing_chunks", embedding_store="")
|
||||
await seed.start()
|
||||
indexed_node = node("memory.md")
|
||||
await seed.upsert(
|
||||
[(indexed_node, [chunk("memory-chunk", "memory.md", "remember this")])],
|
||||
)
|
||||
await seed.close()
|
||||
|
||||
seed.chunks_path.unlink()
|
||||
store = LocalFileStore(name="t_missing_chunks", embedding_store="")
|
||||
await store.start()
|
||||
|
||||
assert store.file_chunks == {}
|
||||
assert await store.get_nodes() == []
|
||||
assert set(store.keyword_index.document_ids) == set()
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_load_clears_graph_and_chunks_when_chunk_sets_partially_diverge():
|
||||
"""Missing and orphaned chunks invalidate the atomic derived snapshot."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = LocalFileStore(name="t_torn_chunks", embedding_store="")
|
||||
await store.start()
|
||||
indexed_node = node("memory.md")
|
||||
indexed_node.chunk_ids = ["kept", "missing"]
|
||||
await store.file_graph.upsert_nodes([indexed_node])
|
||||
store.file_chunks = {
|
||||
"kept": chunk("kept", "memory.md", "kept text"),
|
||||
"orphaned": chunk("orphaned", "old.md", "orphaned text"),
|
||||
}
|
||||
|
||||
repaired = await store._repair_graph_chunk_consistency()
|
||||
|
||||
assert repaired is True
|
||||
assert store.file_chunks == {}
|
||||
assert await store.get_nodes() == []
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_load_clears_stale_keyword_index_when_chunks_are_empty():
|
||||
"""An empty chunk store is still an exact state BM25 must mirror."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
seed = LocalFileStore(name="t_empty_chunk_keyword", embedding_store="")
|
||||
await seed.start()
|
||||
await seed.keyword_index.add_docs({"stale": "stale keyword document"})
|
||||
await seed.close()
|
||||
|
||||
store = LocalFileStore(name="t_empty_chunk_keyword", embedding_store="")
|
||||
await store.start()
|
||||
|
||||
assert store.file_chunks == {}
|
||||
assert set(store.keyword_index.document_ids) == set()
|
||||
assert not store.keyword_index.index_file.exists()
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_keyword_sync_rebuilds_when_backend_only_exposes_matching_count():
|
||||
"""Matching counts cannot prove that a backend contains the expected IDs."""
|
||||
|
||||
|
|
@ -217,6 +316,35 @@ def test_keyword_sync_rebuilds_when_backend_only_exposes_matching_count():
|
|||
run(go())
|
||||
|
||||
|
||||
def test_keyword_sync_rebuilds_in_progress_batches(monkeypatch):
|
||||
"""Foreground keyword repair uses bounded batches suitable for progress reporting."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = LocalFileStore(name="t_keyword_progress", embedding_store="")
|
||||
await store.start()
|
||||
store.file_chunks = {str(index): chunk(str(index), f"{index}.md", f"content {index}") for index in range(5)}
|
||||
await store.keyword_index.clear()
|
||||
|
||||
batch_sizes = []
|
||||
original_add_docs = store.keyword_index.add_docs
|
||||
|
||||
async def recording_add_docs(docs):
|
||||
batch_sizes.append(len(docs))
|
||||
await original_add_docs(docs)
|
||||
|
||||
monkeypatch.setattr(local_file_store_module, "_KEYWORD_REBUILD_BATCH_SIZE", 2)
|
||||
monkeypatch.setattr(store.keyword_index, "add_docs", recording_add_docs)
|
||||
|
||||
await store._sync_keyword_index_from_chunks()
|
||||
|
||||
assert batch_sizes == [2, 2, 1]
|
||||
assert set(store.keyword_index.document_ids) == set(store.file_chunks)
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_chunk_persistence_uses_compact_embedding_and_round_trips():
|
||||
"""Chunk persistence avoids JSON float lists while preserving float16 vectors."""
|
||||
|
||||
|
|
@ -226,7 +354,7 @@ def test_chunk_persistence_uses_compact_embedding_and_round_trips():
|
|||
await store.start()
|
||||
original = chunk("a", "a.md", "alpha text", source="test")
|
||||
original.embedding = np.array([0.25, -1.5, 3.0], dtype=np.float16)
|
||||
store.file_chunks[original.id] = original
|
||||
await set_chunks_with_graph(store, {original.id: original})
|
||||
await store.dump()
|
||||
|
||||
payload = json.loads(next(read_jsonl_zst(store.chunks_path)))
|
||||
|
|
@ -284,6 +412,10 @@ def test_chunk_persistence_loads_legacy_json_embedding_list():
|
|||
original = chunk("legacy", "legacy.md", "legacy text")
|
||||
original.embedding = np.array([0.5, 1.5], dtype=np.float16)
|
||||
write_jsonl_zst(store.chunks_path, [original.model_dump_json()])
|
||||
await set_chunks_with_graph(store, {})
|
||||
legacy_node = node("legacy.md")
|
||||
legacy_node.chunk_ids = [original.id]
|
||||
await store.file_graph.upsert_nodes([legacy_node])
|
||||
|
||||
await store.load()
|
||||
restored = store.file_chunks[original.id]
|
||||
|
|
@ -315,7 +447,7 @@ def test_same_chunk_id_with_changed_text_gets_new_embedding():
|
|||
|
||||
|
||||
def test_load_backfills_missing_embeddings_from_persisted_chunks():
|
||||
"""Loading old chunks after enabling embeddings backfills and persists vectors."""
|
||||
"""Startup backfills old chunks in the background and persists vectors."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
|
|
@ -330,9 +462,9 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks():
|
|||
await store.close()
|
||||
|
||||
store = LocalFileStore(name="t_embedding_backfill", embedding_store="")
|
||||
await store.start()
|
||||
store.embedding_store = FakeEmbeddingStore()
|
||||
await store.load()
|
||||
await store.start()
|
||||
await store._embedding_backfill_task
|
||||
|
||||
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
|
||||
assert store.file_chunks["b"].embedding.tolist() == [0.0, 1.0]
|
||||
|
|
@ -347,22 +479,78 @@ def test_load_backfills_missing_embeddings_from_persisted_chunks():
|
|||
run(go())
|
||||
|
||||
|
||||
def test_start_does_not_wait_for_embedding_backfill():
|
||||
"""Remote embedding repair runs after the file store becomes ready."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
seed = LocalFileStore(name="t_background_embedding", embedding_store="")
|
||||
await seed.start()
|
||||
await set_chunks_with_graph(seed, {"a": chunk("a", "a.md", "alpha text")})
|
||||
await seed.dump()
|
||||
await seed.close()
|
||||
|
||||
store = LocalFileStore(name="t_background_embedding", embedding_store="")
|
||||
fake = BlockingEmbeddingStore()
|
||||
store.embedding_store = fake
|
||||
await store.start()
|
||||
|
||||
await asyncio.wait_for(fake.started.wait(), timeout=1)
|
||||
assert store.is_started
|
||||
assert store.file_chunks["a"].embedding is None
|
||||
|
||||
fake.release.set()
|
||||
await store._embedding_backfill_task
|
||||
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_background_embedding_backfill_uses_provider_batch_size():
|
||||
"""Embedding repair reports progress over the provider's bounded batches."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
seed = LocalFileStore(name="t_embedding_batches", embedding_store="")
|
||||
await seed.start()
|
||||
await set_chunks_with_graph(
|
||||
seed,
|
||||
{str(index): chunk(str(index), f"{index}.md", f"content {index}") for index in range(5)},
|
||||
)
|
||||
await seed.dump()
|
||||
await seed.close()
|
||||
|
||||
store = LocalFileStore(name="t_embedding_batches", embedding_store="")
|
||||
fake = CountingFakeEmbeddingStore()
|
||||
fake.max_batch_size = 2
|
||||
store.embedding_store = fake
|
||||
await store.start()
|
||||
await store._embedding_backfill_task
|
||||
|
||||
assert [len(batch) for batch in fake.node_embedding_calls] == [2, 2, 1]
|
||||
assert all(chunk.embedding is not None for chunk in store.file_chunks.values())
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_load_skips_backfill_when_embedding_health_check_fails():
|
||||
"""Backfill disables embeddings before batching when the provider is unhealthy."""
|
||||
"""Background backfill disables embeddings before batching when the provider is unhealthy."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = LocalFileStore(name="t_embedding_backfill_unhealthy", embedding_store="")
|
||||
await store.start()
|
||||
store.file_chunks = {"a": chunk("a", "a.md", "alpha text")}
|
||||
await set_chunks_with_graph(store, {"a": chunk("a", "a.md", "alpha text")})
|
||||
await store.dump()
|
||||
await store.close()
|
||||
|
||||
store = LocalFileStore(name="t_embedding_backfill_unhealthy", embedding_store="")
|
||||
await store.start()
|
||||
fake = UnhealthyCountingEmbeddingStore()
|
||||
store.embedding_store = fake
|
||||
await store.load()
|
||||
await store.start()
|
||||
await store._embedding_backfill_task
|
||||
|
||||
assert not fake.node_embedding_calls
|
||||
assert store.embedding_store is None
|
||||
|
|
@ -381,15 +569,15 @@ def test_load_reembeds_persisted_chunks_with_stale_embedding_dimensions():
|
|||
await store.start()
|
||||
stale = chunk("a", "a.md", "alpha text")
|
||||
stale.embedding = np.array([1.0], dtype=np.float16)
|
||||
store.file_chunks = {"a": stale}
|
||||
await set_chunks_with_graph(store, {"a": stale})
|
||||
await store.dump()
|
||||
await store.close()
|
||||
|
||||
store = LocalFileStore(name="t_embedding_stale_dim", embedding_store="")
|
||||
await store.start()
|
||||
fake = CountingFakeEmbeddingStore()
|
||||
store.embedding_store = fake
|
||||
await store.load()
|
||||
await store.start()
|
||||
await store._embedding_backfill_task
|
||||
|
||||
assert fake.node_embedding_calls == [["a"]]
|
||||
assert store.file_chunks["a"].embedding.tolist() == [1.0, 0.0]
|
||||
|
|
@ -529,6 +717,56 @@ def test_faiss_rebuilds_stale_sidecar_and_updates_same_id_text():
|
|||
run(go())
|
||||
|
||||
|
||||
def test_faiss_concurrent_dumps_are_serialized(monkeypatch):
|
||||
"""Concurrent persistence must not interleave writes to the FAISS sidecar pair."""
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
try:
|
||||
store = FaissLocalFileStore(name="t_faiss_dump_lock", embedding_store="")
|
||||
except ImportError:
|
||||
pytest.skip("faiss is not installed")
|
||||
await store.start()
|
||||
store.embedding_store = FakeEmbeddingStore()
|
||||
store._faiss_index = store._new_index()
|
||||
|
||||
first_started = asyncio.Event()
|
||||
release_first = asyncio.Event()
|
||||
active_writers = 0
|
||||
max_active_writers = 0
|
||||
write_count = 0
|
||||
original_write_sidecar = store._write_sidecar
|
||||
|
||||
async def blocking_write_sidecar():
|
||||
nonlocal active_writers, max_active_writers, write_count
|
||||
active_writers += 1
|
||||
max_active_writers = max(max_active_writers, active_writers)
|
||||
write_count += 1
|
||||
if write_count == 1:
|
||||
first_started.set()
|
||||
await release_first.wait()
|
||||
active_writers -= 1
|
||||
|
||||
monkeypatch.setattr(store, "_write_sidecar", blocking_write_sidecar)
|
||||
|
||||
first = asyncio.create_task(store.dump())
|
||||
await first_started.wait()
|
||||
second = asyncio.create_task(store.dump())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert active_writers == 1
|
||||
assert max_active_writers == 1
|
||||
|
||||
release_first.set()
|
||||
await asyncio.gather(first, second)
|
||||
assert write_count == 2
|
||||
assert max_active_writers == 1
|
||||
monkeypatch.setattr(store, "_write_sidecar", original_write_sidecar)
|
||||
await store.close()
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_faiss_rebuild_skips_wrong_dimension_chunks():
|
||||
"""FAISS rebuild should ignore chunks whose embedding dimensions do not match."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,16 @@
|
|||
"""Tests for logging configuration handoff during app startup."""
|
||||
|
||||
import concurrent.futures
|
||||
import io
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from reme.application import Application
|
||||
from reme.config.config_parser import resolve_app_config
|
||||
from reme.utils import logger_utils
|
||||
|
||||
|
||||
class DummyLogger:
|
||||
|
|
@ -16,6 +25,108 @@ class DummyLogger:
|
|||
return None
|
||||
|
||||
|
||||
def test_stdlib_formatter_matches_qwenpaw_console_format(monkeypatch, tmp_path, capsys):
|
||||
"""Stdlib logs should use QwenPaw's level/path/time/message layout."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setenv("REME_DISABLE_LOGURU", "true")
|
||||
source_path = tmp_path / "src" / "qwenpaw" / "worker.py"
|
||||
record = logging.LogRecord(
|
||||
name="reme",
|
||||
level=logging.INFO,
|
||||
pathname=str(source_path),
|
||||
lineno=42,
|
||||
msg="Memory index loaded",
|
||||
args=(),
|
||||
exc_info=None,
|
||||
)
|
||||
record.created = 0
|
||||
|
||||
logger = logger_utils.get_logger(log_to_file=False, force_init=True)
|
||||
logger.handle(record)
|
||||
|
||||
formatted_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(record.created))
|
||||
assert capsys.readouterr().out == (f"INFO src/qwenpaw/worker.py:42 | {formatted_time} | Memory index loaded\n")
|
||||
logger_utils.get_logger(log_to_console=False, log_to_file=False, force_init=True)
|
||||
|
||||
|
||||
def test_stdlib_forwards_screen_and_file_logs_to_qwenpaw(monkeypatch, tmp_path):
|
||||
"""Embedded stdlib logging should reuse QwenPaw's active sinks."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setenv("REME_DISABLE_LOGURU", "true")
|
||||
monkeypatch.setattr(logger_utils, "_logger", None)
|
||||
|
||||
qwenpaw_logger = logging.getLogger("qwenpaw")
|
||||
original_handlers = list(qwenpaw_logger.handlers)
|
||||
original_level = qwenpaw_logger.level
|
||||
original_propagate = qwenpaw_logger.propagate
|
||||
for handler in original_handlers:
|
||||
qwenpaw_logger.removeHandler(handler)
|
||||
|
||||
console_stream = io.StringIO()
|
||||
console_handler = logging.StreamHandler(console_stream)
|
||||
file_path = tmp_path / "qwenpaw.log"
|
||||
file_handler = logging.FileHandler(file_path, encoding="utf-8")
|
||||
formatter = logging.Formatter("%(levelname)s | %(message)s")
|
||||
console_handler.setFormatter(formatter)
|
||||
file_handler.setFormatter(formatter)
|
||||
qwenpaw_logger.addHandler(console_handler)
|
||||
qwenpaw_logger.addHandler(file_handler)
|
||||
qwenpaw_logger.setLevel(logging.INFO)
|
||||
qwenpaw_logger.propagate = False
|
||||
|
||||
reme_logger = logging.getLogger("reme")
|
||||
try:
|
||||
logger = logger_utils.get_logger(
|
||||
log_to_console=True,
|
||||
log_to_file=True,
|
||||
force_init=True,
|
||||
)
|
||||
logger.info("Memory index loaded")
|
||||
file_handler.flush()
|
||||
|
||||
assert console_stream.getvalue() == "INFO | Memory index loaded\n"
|
||||
assert file_path.read_text(encoding="utf-8") == "INFO | Memory index loaded\n"
|
||||
assert not (tmp_path / "logs").exists()
|
||||
assert len(reme_logger.handlers) == 1
|
||||
assert reme_logger.handlers[0].target_name == "qwenpaw"
|
||||
finally:
|
||||
for handler in list(reme_logger.handlers):
|
||||
reme_logger.removeHandler(handler)
|
||||
handler.close()
|
||||
qwenpaw_logger.removeHandler(console_handler)
|
||||
qwenpaw_logger.removeHandler(file_handler)
|
||||
console_handler.close()
|
||||
file_handler.close()
|
||||
for handler in original_handlers:
|
||||
qwenpaw_logger.addHandler(handler)
|
||||
qwenpaw_logger.setLevel(original_level)
|
||||
qwenpaw_logger.propagate = original_propagate
|
||||
|
||||
|
||||
def test_explicit_loguru_enable_keeps_original_backend(monkeypatch):
|
||||
"""An explicit false value must continue to select Loguru unchanged."""
|
||||
sentinel_logger = object()
|
||||
calls = []
|
||||
|
||||
def fake_init(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return sentinel_logger
|
||||
|
||||
monkeypatch.setenv("REME_DISABLE_LOGURU", "false")
|
||||
monkeypatch.setattr(logger_utils, "_logger", None)
|
||||
monkeypatch.setattr(logger_utils, "_init_loguru", fake_init)
|
||||
monkeypatch.setattr(
|
||||
logger_utils,
|
||||
"_init_stdlib",
|
||||
lambda *_args, **_kwargs: pytest.fail("stdlib backend selected"),
|
||||
)
|
||||
|
||||
result = logger_utils.get_logger(force_init=True)
|
||||
|
||||
assert result is sentinel_logger
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_resolve_app_config_does_not_create_file_logger(monkeypatch):
|
||||
"""Config-loading messages should not create empty run log files."""
|
||||
calls = []
|
||||
|
|
@ -58,3 +169,47 @@ def test_application_reinitializes_logger_from_final_config(monkeypatch, tmp_pat
|
|||
"log_to_file": True,
|
||||
"force_init": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_loguru", [True, False])
|
||||
def test_concurrent_force_init_is_serialized(monkeypatch, use_loguru):
|
||||
"""Concurrent application startup must not overlap global logger resets."""
|
||||
state_lock = threading.Lock()
|
||||
active_initializers = 0
|
||||
max_active_initializers = 0
|
||||
initialization_count = 0
|
||||
sentinel_logger = object()
|
||||
|
||||
def fake_init(*_args, **_kwargs):
|
||||
nonlocal active_initializers
|
||||
nonlocal max_active_initializers
|
||||
nonlocal initialization_count
|
||||
|
||||
with state_lock:
|
||||
active_initializers += 1
|
||||
initialization_count += 1
|
||||
max_active_initializers = max(
|
||||
max_active_initializers,
|
||||
active_initializers,
|
||||
)
|
||||
time.sleep(0.01)
|
||||
with state_lock:
|
||||
active_initializers -= 1
|
||||
return sentinel_logger
|
||||
|
||||
monkeypatch.setattr(logger_utils, "_logger", None)
|
||||
monkeypatch.setattr(logger_utils, "_enable_loguru", lambda: use_loguru)
|
||||
monkeypatch.setattr(logger_utils, "_init_loguru", fake_init)
|
||||
monkeypatch.setattr(logger_utils, "_init_stdlib", fake_init)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
|
||||
results = list(
|
||||
pool.map(
|
||||
lambda _index: logger_utils.get_logger(force_init=True),
|
||||
range(32),
|
||||
),
|
||||
)
|
||||
|
||||
assert results == [sentinel_logger] * 32
|
||||
assert initialization_count == 32
|
||||
assert max_active_initializers == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue