refactor(memory): rename to memory, OpenAI env, chunking, batching, dedup (#75)
* refactor(memory): rename package second-brain to memory * refactor(memory): use OpenAI env naming and fix trailing slash * feat(memory): add chunking with env-configurable size and overlap * feat(memory): dedup search results by source in top-K * test(memory): add chunking, batching, dedup, live tests * docs(memory): update README and project map after rename * docs(handoff): add handoff and ADR-032 for memory refactor * docs(handoff): set PR number * fix(ci): reduce index.py complexity to rank A --------- Co-authored-by: opencode-agent <agent@opencode.local>
This commit is contained in:
Sergey
GitHub
parent
92d21a142c
commit
80a4be1d21
16 changed files with 492 additions and 67 deletions
|
|
@ -2,6 +2,14 @@
|
|||
AI_PROVIDER_BASE_URL=https://your-ai-provider.example.com/v1/
|
||||
AI_PROVIDER_API_KEY=your-api-key-here
|
||||
|
||||
# OpenAI Embeddings (Memory CLI)
|
||||
OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
OPENAI_API_KEY=your-openai-api-key
|
||||
OPENAI_EMBEDDING_MODEL=gemini-embedding-2-preview
|
||||
OPENAI_EMBEDDING_BATCH_SIZE=2048
|
||||
MEMORY_CHUNK_SIZE=512
|
||||
MEMORY_CHUNK_OVERLAP=64
|
||||
|
||||
# OpenCode Server
|
||||
OPENCODE_SERVER_PASSWORD=your-opencode-server-password
|
||||
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ opencode # .opencode/ auto-discovered
|
|||
- `app_data/workspaces/` — agent working directory
|
||||
- `app_data/ssh/` — SSH keys (not in git)
|
||||
- `app_data/opencode-memory/` — persistent memory (separate git repo)
|
||||
- `src/` — Python RAG CLI (second-brain)
|
||||
- `src/` — Python RAG CLI (memory)
|
||||
- `docs/` — handoffs, decisions (ADRs), project map
|
||||
- `.github/workflows/` — CI workflows (ubuntu-latest)
|
||||
|
||||
|
|
|
|||
23
docs/decisions/032-pr-75-rename-openai-env-chunking.md
Normal file
23
docs/decisions/032-pr-75-rename-openai-env-chunking.md
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
# ADR-032: Rename to memory, OpenAI env, chunking, batching, dedup (PR #75)
|
||||
|
||||
## Статус
|
||||
Accepted (2026-07-26)
|
||||
|
||||
## Контекст
|
||||
Python-пакет `src/memory/` (RAG CLI для opencode-memory plugin) имел 7 проблем: нейминг-путаница (3 разных имени для одного пакета), нестандартный env-нейминг (`AI_PROVIDER_API_URL` ≠ `AI_PROVIDER_BASE_URL` из `.env.example`/`opencode.json`), реальная 404-бага из-за trailing slash, грубый индекс (1 эмбеддинг на файл без чанков), отсутствие batching (риск упереться в OpenAI лимит input), отсутствие дедупа в search (засорение top-K одним файлом), захардкоженная модель эмбеддингов.
|
||||
|
||||
PR #1 в серии из 2 — этот PR чинит CLI, PR #2 интегрирует CLI в плагин через setup-memory.sh.
|
||||
|
||||
## Решение
|
||||
1. **Нейминг**: `pyproject.toml` name=`second-brain` → `memory`, `[project.scripts] rag = "src.memory.cli:main"` (script entry point, совместимость с wrapper). `cli.py` `prog="rag"` → `prog="memory"`. Модуль `src/memory/` НЕ переименован (src-layout, `packages=["src"]`).
|
||||
2. **Env-нейминг OpenAI-стандарт**: `OPENAI_BASE_URL`/`OPENAI_API_KEY` (вместо `AI_PROVIDER_API_URL`/`AI_PROVIDER_API_KEY`). `API_URL.rstrip("/")` фиксит trailing-slash. `EMBEDDING_MODEL` из env `OPENAI_EMBEDDING_MODEL` (default `gemini-embedding-2-preview`).
|
||||
3. **Chunking**: `_chunk_text(text, size, overlap)` в `index.py`, параметры из env `MEMORY_CHUNK_SIZE` (512)/`MEMORY_CHUNK_OVERLAP` (64). `file_map` с `chunk_idx`/`offset`/`text`. Graceful при `size<=overlap` (1 chunk, не infinite loop).
|
||||
4. **Batching**: `BATCH_SIZE = int(os.environ.get("OPENAI_EMBEDDING_BATCH_SIZE", "2048"))` в `embedder.py`. `embed_texts` режет texts > BATCH_SIZE, конкатенирует результаты. Default 2048 = OpenAI лимит для text-embedding-3-small.
|
||||
5. **Дедуп**: `search.py` после sort по score итерирует с `seen` set, оставляет highest score per source, останавливается на K уникальных source'ов.
|
||||
6. **Тесты**: `test_chunking.py` (7 edge cases), `test_embedder_live.py` (2 live, skip без RUN_LIVE), + 4 embedder теста (default/custom model, trailing slash, batches), + 3 index теста (long file, env override, .rag skip), + 1 search тест (dedup).
|
||||
|
||||
## Альтернативы
|
||||
- **Полный rename `src/memory/` → `src/rag/`** — отклонено: src-layout `packages=["src"]` в pyproject, rename модуля требует обновления всех imports (cli, embedder, index, search, tests, __main__). Script entry point `rag = "src.memory.cli:main"` даёт исполняемый `rag` без rename модуля. Module name `memory` (что package делает) + script name `rag` (как пользователь вызывает) — разумное разделение.
|
||||
- **Читать `OPENAI_BASE_URL`/`EMBEDDING_MODEL`/`BATCH_SIZE` в `embed_texts` каждый вызов** (вместо module-level) — отклонено: лишний overhead на каждый вызов (env lookup + int parse), несовместимость с существующим `API_URL` module-level паттерном. Module-level + `importlib.reload` в тестах — trade-off: простой код, хрупкие reload-тесты.
|
||||
- **Дропать empty files из индекса** (вместо 1 записи с пустым чанком) — отклонено: меняет существующее поведение (раньше empty files попадали в индекс с 1 embedding). Empty embedding даёт score 0, файл остаётся в index для completeness.
|
||||
- **Параметризовать `_chunk_text` через argparse вместо env** — отклонено: env-параметры позволяют настраивать без правки CLI invocation, совместимо с setup-memory.sh wrapper (env передаётся через process env, не флаги).
|
||||
46
docs/handoff/pr-75-rename-openai-env-chunking.md
Normal file
46
docs/handoff/pr-75-rename-openai-env-chunking.md
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
---
|
||||
pr: 75
|
||||
title: Rename to memory, OpenAI env, chunking, batching, dedup
|
||||
---
|
||||
|
||||
## Что сделано
|
||||
- `pyproject.toml`: `name = "second-brain"` → `name = "memory"`, добавлена секция `[project.scripts] rag = "src.memory.cli:main"`. `uv.lock` регенерирован (`uv lock`).
|
||||
- `src/memory/cli.py`: `prog="rag"` → `prog="memory"`.
|
||||
- `src/memory/embedder.py`: env нейминг `AI_PROVIDER_API_URL`/`AI_PROVIDER_API_KEY` → `OPENAI_BASE_URL`/`OPENAI_API_KEY` (OpenAI-совместимый). `API_URL.rstrip("/")` фиксит trailing-slash баг (`/v1//embeddings` → 404). `EMBEDDING_MODEL` из env `OPENAI_EMBEDDING_MODEL` (default `gemini-embedding-2-preview`). `BATCH_SIZE = int(os.environ.get("OPENAI_EMBEDDING_BATCH_SIZE", "2048"))` — `embed_texts` режет texts > BATCH_SIZE на батчи, конкатенирует результаты.
|
||||
- `src/memory/index.py`: добавлена `_chunk_text(text, size, overlap) -> list[tuple[str, int]]` (graceful при size<=overlap → 1 chunk). `run_index` использует чанки (`MEMORY_CHUNK_SIZE` default 512, `MEMORY_CHUNK_OVERLAP` default 64 из env). `file_map` поля: `source`, `chunk_idx`, `offset`, `text`. Пустые файлы → 1 запись с пустым чанком (не дропаются).
|
||||
- `src/memory/search.py`: дедуп по `source` в top-K — после sort по score, итерация с `seen` set, оставляет highest score per source, останавливается на K уникальных.
|
||||
- `.env.example`: добавлен блок "OpenAI Embeddings (Memory CLI)" с 6 переменными (`OPENAI_BASE_URL`, `OPENAI_API_KEY`, `OPENAI_EMBEDDING_MODEL`, `OPENAI_EMBEDDING_BATCH_SIZE`, `MEMORY_CHUNK_SIZE`, `MEMORY_CHUNK_OVERLAP`). `AI_PROVIDER_BASE_URL`/`AI_PROVIDER_API_KEY` оставлены (LLM провайдер в opencode.json, отдельный concern).
|
||||
- `tests/test_embedder.py`: env нейминг (`OPENAI_BASE_URL`), + 4 теста: `test_embed_texts_default_model`, `test_embed_texts_custom_model` (reload с env override), `test_embed_texts_trailing_slash` (URL без `//`), `test_embed_texts_batches` (3000 texts → 2 вызова [2048, 952]).
|
||||
- `tests/test_index.py`: + 3 теста: `test_chunking_long_file` (1500 chars → 3+ records, chunk_idx/offset), `test_chunking_size_env_override` (`MEMORY_CHUNK_SIZE=256`), `test_chunking_skip_rag_dir`. Существующий `test_index_creates_json` обновлён (проверка `chunk_idx`/`offset`).
|
||||
- `tests/test_search.py`: + `test_dedup_by_source` (2 чанка одного source → 1 в top-K, highest score).
|
||||
- `tests/test_chunking.py` (новый): 7 edge cases `_chunk_text` — empty, 1 char, ровно size, size+1, unicode emoji, size<overlap, size==overlap.
|
||||
- `tests/test_embedder_live.py` (новый): 2 live-теста (`@pytest.mark.skipif(not RUN_LIVE)`) — single + batch embed.
|
||||
- `README.md:40`, `docs/project-map/README.md:69`: `second-brain` → `memory`. Project-map обновлён: `embedder.py` → `OPENAI_BASE_URL`, `index.py` → chunking, `search.py` → dedup, `cli.py` → `prog="memory"`, + entries для `test_embedder_live.py`, `test_chunking.py`.
|
||||
- ADR-032 + этот handoff
|
||||
|
||||
## Почему
|
||||
7 проблем в `src/memory/` (RAG CLI для opencode-memory plugin), чинятся этим PR:
|
||||
|
||||
1. **Нейминг-путаница**: `pyproject.toml` name=`second-brain`, модуль `src/memory/`, CLI `prog="rag"` — три разных имени. Унифицировано под `memory` (пакет) + `rag` script entry point (совместимость с setup-memory.sh wrapper).
|
||||
2. **Env-нейминг нестандартен**: код читал `AI_PROVIDER_API_URL`, но `.env.example`/`opencode.json` используют `AI_PROVIDER_BASE_URL` — код не работал "из коробки". Перешли на OpenAI-стандарт `OPENAI_BASE_URL`/`OPENAI_API_KEY` (подтверждено openai-python docs).
|
||||
3. **Trailing-slash бага**: `f"{API_URL}/embeddings"` без `.rstrip("/")` → при `OPENAI_BASE_URL=https://api.openai.com/v1/` URL становился `/v1//embeddings` → HTTP 404 (воспроизведено smoke-test'ом).
|
||||
4. **Индекс 1-эмбеддинг-на-файл**: 1 embedding на весь файл — нет чанков. Для 74 файлов (1.2 MB) грубо, не находит фрагменты внутри длинных файлов. Добавлено чанкирование (default 512 chars, overlap 64).
|
||||
5. **Нет batching**: все texts одним POST — при 3396 чанках риск упереться в OpenAI лимит (2048 input для text-embedding-3-small). Добавлен batching (default 2048).
|
||||
6. **Нет дедупа в search**: 2 чанка одного файла в top-K засоряли выдачу одним файлом. Добавлен дедуп по source (highest score per source).
|
||||
7. **Модель захардкожена**: `EMBEDDING_MODEL = "gemini-embedding-2-preview"` — нельзя сменить без правки кода. Вынесена в env `OPENAI_EMBEDDING_MODEL`.
|
||||
|
||||
PR #1 в серии из 2 (PR #2 — интеграция в плагин через setup-memory.sh wrapper).
|
||||
|
||||
## Pending
|
||||
- PR #2: интеграция `memory` CLI в плагин `@mathew-cf/opencode-memory` через `setup-memory.sh` (замена Rust rag-cli на Python `memory index`/`memory search`). Требует: OPENAI_BASE_URL/OPENAI_API_KEY в `.env`, обновление wrapper в setup-memory.sh.
|
||||
- `opencode.json:28` всё ещё использует `AI_PROVIDER_API_KEY` для LLM провайдера — отдельный concern (LLM, не embeddings), НЕ трогался в этом PR. Возможный future PR для OpenAI-стандартизации LLM env.
|
||||
- Smoke-test (3396 чанков, 1.2 MB index.json) — требует реальный API key, не запускался в CI (только unit-тесты с моками). Запустить вручную перед merge: `OPENAI_BASE_URL=$AI_PROVIDER_BASE_URL OPENAI_API_KEY=$AI_PROVIDER_API_KEY uv run python -m src.memory index /root/.local/share/opencode/opencode-memory -o /tmp/.rag-test && ls /tmp/.rag-test/index.json`
|
||||
|
||||
## Watch out
|
||||
- **`EMBEDDING_MODEL`/`BATCH_SIZE` — module-level, читаются при импорте.** Тесты с env override (custom model, trailing slash) делают `importlib.reload(embedder_mod)` и восстанавливают env после теста. Если未来的 код сделает `EMBEDDING_MODEL` mutable через функцию — паттерн reload сломается. Альтернатива (читать env в `embed_texts` каждый вызов) отклонена: лишний overhead на каждый вызов + несовместимость с существующим `API_URL` module-level паттерном.
|
||||
- **Reload-тесты меняют module global** — `test_embed_texts_custom_model`/`test_embed_texts_trailing_slash` используют `monkeypatch.setenv` + `importlib.reload`. `monkeypatch` auto-restore env, но reload модуля НЕ auto-restore. Явный `reload` в конце каждого теста восстанавливает default state. Если тест упадёт ДО финального reload — следующий тест может получить stale module state. Mitigation: `monkeypatch.setenv` восстанавливает env, а финальный `reload` в `finally`-стиле (после assert) — нет. Если тест упал на assert — reload не выполнится. Риск минимальный (assert в конце), но для robustness можно обернуть в try/finally — НЕ сделано для простоты.
|
||||
- **Chunking `size<=overlap` → 1 chunk (не infinite loop).** `_chunk_text` возвращает `[(text, 0)]` если `size <= overlap` (step = size-overlap <= 0 → `range(0, len, 0)` = infinite). Это graceful degradation, не ошибка. Тест `test_size_less_than_overlap`/`test_size_equals_overlap` покрывают.
|
||||
- **`file_map` тип `dict[str, str | int]`** — `chunk_idx`/`offset` это `int`, `source`/`text` это `str`. mypy strict требует union. Существующий `dict[str, str]` расширен до `dict[str, str | int]`. JSON serialization сохраняет int как int (не string).
|
||||
- **Empty files → 1 запись с пустым чанком.** `_chunk_text("")` возвращает `[]`, но `run_index` заменяет на `[("", 0)]` чтобы не дропать файл из индекса (поиск по empty embedding даст score 0, но файл остаётся в index для completeness). Альтернатива (дропать empty files) отклонена — меняет существующее поведение (раньше empty files попадали в индекс с 1 embedding).
|
||||
- **`uv lock` обязателен после rename `pyproject.toml`** — `uv.lock:716` содержал `name = "second-brain"`. Без `uv lock` lockfile рассинхронизирован с pyproject. CI `uv sync` использует lockfile → несоответствие. `uv lock` регенерирован (confirmed: `Removed second-brain v0.1.0`, `Added memory v0.1.0`).
|
||||
- **ADR number = sequential (032), НЕ PR number.** Проверить ADR naming в handoff до push (эволюция паттерна PR#26 docs-reviewer typo).
|
||||
|
|
@ -66,14 +66,14 @@ opencode-config/
|
|||
│ ├── handoff/ # PR handoffs (pr-<N>-<slug>.md)
|
||||
│ ├── decisions/ # ADRs (NNN-pr-<N>-<slug>.md)
|
||||
│ └── project-map/ # This file — structure snapshot
|
||||
├── src/ # Python RAG CLI (second-brain) — PR#17
|
||||
├── src/ # Python RAG CLI (memory) — PR#17
|
||||
│ └── memory/
|
||||
│ ├── __init__.py
|
||||
│ ├── __main__.py # Entry point for `python -m memory`
|
||||
│ ├── cli.py # CLI commands
|
||||
│ ├── embedder.py # Embedding via AI_PROVIDER_API_URL (env-only)
|
||||
│ ├── index.py # Indexing
|
||||
│ └── search.py # Search
|
||||
│ ├── cli.py # CLI commands (prog="memory")
|
||||
│ ├── embedder.py # Embedding via OPENAI_BASE_URL (env-only, OpenAI-compatible)
|
||||
│ ├── index.py # Indexing with chunking (MEMORY_CHUNK_SIZE/OVERLAP env)
|
||||
│ └── search.py # Search with dedup by source in top-K
|
||||
├── tests/ # pytest + TS/MJS test suite — PR#17
|
||||
│ ├── _ts_loader.mjs # TS test loader (load/exec_stub/exec_stub_json/exec_real modes; relative import inlining via inlineShared()) — PR#38, PR#65
|
||||
│ ├── test_agent_frontmatter.py # Agent frontmatter validators (no top-level doom_loop, permission.doom_loop present, steps:150) — PR#49, PR#69
|
||||
|
|
@ -88,8 +88,10 @@ opencode-config/
|
|||
│ ├── test_create_issue_tool.ts # TS wrapper test (mjs loader; +repo cases) — PR#38, PR#65
|
||||
│ ├── test_create_pr_tool.py # .opencode/tools/create-pr.ts (via _ts_loader.mjs exec_stub_json; +repo explicit/omitted/invalid) — PR#38, PR#65
|
||||
│ ├── test_create_pr_tool.ts # TS wrapper test (mjs loader; +repo cases) — PR#38, PR#65
|
||||
│ ├── test_embedder.py # src/memory/embedder.py (mocks AI_PROVIDER_API_URL)
|
||||
│ ├── test_index.py # src/memory/index.py
|
||||
│ ├── test_embedder.py # src/memory/embedder.py (mocks OPENAI_BASE_URL)
|
||||
│ ├── test_embedder_live.py # Live embed tests (skip without RUN_LIVE=1)
|
||||
│ ├── test_index.py # src/memory/index.py (chunking + env override + .rag skip)
|
||||
│ ├── test_chunking.py # _chunk_text edge cases (empty, unicode, size<overlap)
|
||||
│ ├── test_memory_setup_tool.py # .opencode/tools/memory-setup.ts (via _ts_loader.mjs) — PR#36
|
||||
│ ├── test_memory_setup_tool.ts # TS wrapper test (mjs loader) — PR#36
|
||||
│ ├── test_merge_pr_tool.py # .opencode/tools/merge-pr.ts (via _ts_loader.mjs; 8 tests: base + repo explicit/omitted/invalid) — PR#65
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ requires = ["hatchling"]
|
|||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "second-brain"
|
||||
name = "memory"
|
||||
version = "0.1.0"
|
||||
description = "Docker-based AI coding assistant with persistent memory"
|
||||
readme = "README.md"
|
||||
|
|
@ -39,6 +39,9 @@ dev = [
|
|||
"pre-commit>=3.7",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
rag = "src.memory.cli:main"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/slaid098/opencode-config"
|
||||
Repository = "https://github.com/slaid098/opencode-config"
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from src.memory.search import run_search
|
|||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
parser = argparse.ArgumentParser(prog="rag")
|
||||
parser = argparse.ArgumentParser(prog="memory")
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
search_p = sub.add_parser("search")
|
||||
|
|
|
|||
|
|
@ -3,11 +3,13 @@ import os
|
|||
import httpx
|
||||
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential
|
||||
|
||||
API_URL = os.environ.get("AI_PROVIDER_API_URL")
|
||||
API_URL = os.environ.get("OPENAI_BASE_URL")
|
||||
if not API_URL:
|
||||
raise RuntimeError("AI_PROVIDER_API_URL env var not set")
|
||||
API_KEY = os.environ.get("AI_PROVIDER_API_KEY", "")
|
||||
EMBEDDING_MODEL = "gemini-embedding-2-preview"
|
||||
raise RuntimeError("OPENAI_BASE_URL env var not set")
|
||||
API_URL = API_URL.rstrip("/")
|
||||
API_KEY = os.environ.get("OPENAI_API_KEY", "")
|
||||
EMBEDDING_MODEL = os.environ.get("OPENAI_EMBEDDING_MODEL", "gemini-embedding-2-preview")
|
||||
BATCH_SIZE = int(os.environ.get("OPENAI_EMBEDDING_BATCH_SIZE", "2048"))
|
||||
|
||||
|
||||
def _is_retryable(exc: BaseException) -> bool:
|
||||
|
|
@ -42,7 +44,15 @@ def embed_texts(texts: list[str]) -> list[list[float]]:
|
|||
|
||||
url = f"{API_URL}/embeddings"
|
||||
headers = {"Authorization": f"Bearer {API_KEY}"}
|
||||
payload: dict[str, str | list[str]] = {"model": EMBEDDING_MODEL, "input": texts}
|
||||
|
||||
with httpx.Client(timeout=120.0) as client:
|
||||
if len(texts) <= BATCH_SIZE:
|
||||
payload: dict[str, str | list[str]] = {"model": EMBEDDING_MODEL, "input": texts}
|
||||
return _call_embedding_api(client, url, headers, payload)
|
||||
|
||||
results: list[list[float]] = []
|
||||
for i in range(0, len(texts), BATCH_SIZE):
|
||||
batch = texts[i : i + BATCH_SIZE]
|
||||
payload = {"model": EMBEDDING_MODEL, "input": batch}
|
||||
results.extend(_call_embedding_api(client, url, headers, payload))
|
||||
return results
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from src.memory.embedder import embed_texts
|
||||
|
|
@ -15,30 +16,53 @@ def _extract_text(content: str) -> str:
|
|||
return content.strip()
|
||||
|
||||
|
||||
def _chunk_text(text: str, size: int, overlap: int) -> list[tuple[str, int]]:
|
||||
if not text or size <= overlap or len(text) <= size:
|
||||
return [(text, 0)] if text else []
|
||||
step = size - overlap
|
||||
return [(text[i : i + size], i) for i in range(0, len(text), step)]
|
||||
|
||||
|
||||
def _file_entry(rel: Path, chunk_idx: int, offset: int, chunk: str) -> dict[str, str | int]:
|
||||
return {
|
||||
"source": str(rel),
|
||||
"chunk_idx": chunk_idx,
|
||||
"offset": offset,
|
||||
"text": chunk[:500],
|
||||
}
|
||||
|
||||
|
||||
def _build_file_map(
|
||||
memory_dir: Path, md_files: list[Path], chunk_size: int, chunk_overlap: int
|
||||
) -> tuple[list[str], list[dict[str, str | int]]]:
|
||||
texts: list[str] = []
|
||||
file_map: list[dict[str, str | int]] = []
|
||||
for fpath in md_files:
|
||||
rel = fpath.relative_to(memory_dir)
|
||||
content = fpath.read_text(encoding="utf-8")
|
||||
text = _extract_text(content)
|
||||
chunks = _chunk_text(text, chunk_size, chunk_overlap) or [("", 0)]
|
||||
for chunk_idx, (chunk, offset) in enumerate(chunks):
|
||||
texts.append(chunk)
|
||||
file_map.append(_file_entry(rel, chunk_idx, offset, chunk))
|
||||
return texts, file_map
|
||||
|
||||
|
||||
def run_index(args: argparse.Namespace) -> None:
|
||||
memory_dir = Path(args.memory_dir)
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
md_files = sorted(memory_dir.rglob("*.md"))
|
||||
md_files = [f for f in md_files if ".rag" not in f.parts]
|
||||
|
||||
md_files = [f for f in sorted(memory_dir.rglob("*.md")) if ".rag" not in f.parts]
|
||||
if not md_files:
|
||||
print("No .md files found")
|
||||
return
|
||||
|
||||
texts: list[str] = []
|
||||
file_map: list[dict[str, str]] = []
|
||||
|
||||
for fpath in md_files:
|
||||
rel = fpath.relative_to(memory_dir)
|
||||
content = fpath.read_text(encoding="utf-8")
|
||||
text = _extract_text(content)
|
||||
texts.append(text)
|
||||
file_map.append({"source": str(rel), "text": text[:500]})
|
||||
chunk_size = int(os.environ.get("MEMORY_CHUNK_SIZE", "512"))
|
||||
chunk_overlap = int(os.environ.get("MEMORY_CHUNK_OVERLAP", "64"))
|
||||
|
||||
texts, file_map = _build_file_map(memory_dir, md_files, chunk_size, chunk_overlap)
|
||||
embeddings = embed_texts(texts)
|
||||
|
||||
index = {
|
||||
"files": [{**fm, "embedding": emb} for fm, emb in zip(file_map, embeddings, strict=False)],
|
||||
}
|
||||
|
|
|
|||
|
|
@ -37,6 +37,16 @@ def run_search(args: argparse.Namespace) -> None:
|
|||
results.append({"source": f["source"], "score": sim, "text": f["text"]})
|
||||
|
||||
results.sort(key=lambda r: r["score"], reverse=True)
|
||||
top = results[: args.k]
|
||||
|
||||
seen: set[str] = set()
|
||||
top: list[dict[str, str | float]] = []
|
||||
for r in results:
|
||||
src = str(r["source"])
|
||||
if src in seen:
|
||||
continue
|
||||
seen.add(src)
|
||||
top.append(r)
|
||||
if len(top) >= args.k:
|
||||
break
|
||||
|
||||
print(json.dumps(top, ensure_ascii=False))
|
||||
|
|
|
|||
47
tests/test_chunking.py
Normal file
47
tests/test_chunking.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
import os
|
||||
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
||||
|
||||
from src.memory.index import _chunk_text
|
||||
|
||||
|
||||
class TestChunkText:
|
||||
def test_empty(self) -> None:
|
||||
assert _chunk_text("", 512, 64) == []
|
||||
|
||||
def test_single_char(self) -> None:
|
||||
assert _chunk_text("x", 512, 64) == [("x", 0)]
|
||||
|
||||
def test_exactly_size(self) -> None:
|
||||
text = "a" * 512
|
||||
assert _chunk_text(text, 512, 64) == [(text, 0)]
|
||||
|
||||
def test_size_plus_one(self) -> None:
|
||||
text = "a" * 513
|
||||
chunks = _chunk_text(text, 512, 64)
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0] == (text[:512], 0)
|
||||
assert chunks[1][1] == 512 - 64
|
||||
assert chunks[1][0] == text[448 : 448 + 512]
|
||||
|
||||
def test_unicode_emoji(self) -> None:
|
||||
text = "😀" * 100
|
||||
chunks = _chunk_text(text, 10, 2)
|
||||
assert len(chunks) >= 2
|
||||
offsets = [off for _, off in chunks]
|
||||
assert offsets == sorted(offsets)
|
||||
assert offsets[0] == 0
|
||||
for chunk, _ in chunks:
|
||||
assert chunk in text
|
||||
|
||||
def test_size_less_than_overlap(self) -> None:
|
||||
text = "a" * 100
|
||||
chunks = _chunk_text(text, 10, 20)
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == (text, 0)
|
||||
|
||||
def test_size_equals_overlap(self) -> None:
|
||||
text = "a" * 100
|
||||
chunks = _chunk_text(text, 10, 10)
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == (text, 0)
|
||||
|
|
@ -1,12 +1,14 @@
|
|||
import importlib
|
||||
import os
|
||||
|
||||
os.environ.setdefault("AI_PROVIDER_API_URL", "http://test/v1")
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
||||
|
||||
from typing import NoReturn
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import src.memory.embedder as embedder_mod
|
||||
from src.memory.embedder import embed_texts
|
||||
|
||||
|
||||
|
|
@ -60,3 +62,109 @@ def test_embed_texts_api_error() -> None:
|
|||
|
||||
with patch.object(httpx.Client, "post", mock_post), pytest.raises(httpx.HTTPStatusError):
|
||||
embed_texts(["test"])
|
||||
|
||||
|
||||
def test_embed_texts_default_model() -> None:
|
||||
captured: dict[str, str | list[str]] = {}
|
||||
|
||||
class FakeResponse:
|
||||
status_code = 200
|
||||
|
||||
def json(self):
|
||||
return {"data": [{"embedding": [0.1], "index": 0}]}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def mock_post(self, url, **kwargs):
|
||||
captured.update(kwargs["json"])
|
||||
return FakeResponse()
|
||||
|
||||
with patch.object(httpx.Client, "post", mock_post):
|
||||
embed_texts(["text"])
|
||||
|
||||
assert captured["model"] == "gemini-embedding-2-preview"
|
||||
|
||||
|
||||
def test_embed_texts_custom_model(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
|
||||
importlib.reload(embedder_mod)
|
||||
|
||||
captured: dict[str, str | list[str]] = {}
|
||||
|
||||
class FakeResponse:
|
||||
status_code = 200
|
||||
|
||||
def json(self):
|
||||
return {"data": [{"embedding": [0.1], "index": 0}]}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def mock_post(self, url, **kwargs):
|
||||
captured.update(kwargs["json"])
|
||||
return FakeResponse()
|
||||
|
||||
with patch.object(httpx.Client, "post", mock_post):
|
||||
embedder_mod.embed_texts(["text"])
|
||||
|
||||
assert captured["model"] == "text-embedding-3-small"
|
||||
|
||||
monkeypatch.delenv("OPENAI_EMBEDDING_MODEL", raising=False)
|
||||
importlib.reload(embedder_mod)
|
||||
|
||||
|
||||
def test_embed_texts_trailing_slash(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", "https://api.test.com/v1/")
|
||||
importlib.reload(embedder_mod)
|
||||
|
||||
captured_url: dict[str, str] = {}
|
||||
|
||||
class FakeResponse:
|
||||
status_code = 200
|
||||
|
||||
def json(self):
|
||||
return {"data": [{"embedding": [0.1], "index": 0}]}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def mock_post(self, url, **kwargs):
|
||||
captured_url["url"] = url
|
||||
return FakeResponse()
|
||||
|
||||
with patch.object(httpx.Client, "post", mock_post):
|
||||
embedder_mod.embed_texts(["text"])
|
||||
|
||||
assert captured_url["url"] == "https://api.test.com/v1/embeddings"
|
||||
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", "http://test/v1")
|
||||
importlib.reload(embedder_mod)
|
||||
|
||||
|
||||
def test_embed_texts_batches() -> None:
|
||||
calls: list[int] = []
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, count: int) -> None:
|
||||
self.count = count
|
||||
|
||||
status_code = 200
|
||||
|
||||
def json(self):
|
||||
return {"data": [{"embedding": [0.1], "index": i} for i in range(self.count)]}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def mock_post(self, url, **kwargs):
|
||||
count = len(kwargs["json"]["input"])
|
||||
calls.append(count)
|
||||
return FakeResponse(count)
|
||||
|
||||
with patch.object(httpx.Client, "post", mock_post):
|
||||
result = embed_texts(["text"] * 3000)
|
||||
|
||||
assert len(calls) == 2
|
||||
assert calls == [2048, 952]
|
||||
assert len(result) == 3000
|
||||
|
|
|
|||
21
tests/test_embedder_live.py
Normal file
21
tests/test_embedder_live.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
||||
|
||||
from src.memory.embedder import embed_texts
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get("RUN_LIVE"), reason="needs RUN_LIVE=1")
|
||||
def test_live_embed_single() -> None:
|
||||
r = embed_texts(["hello world"])
|
||||
assert len(r) == 1
|
||||
assert len(r[0]) > 100
|
||||
|
||||
|
||||
@pytest.mark.skipif(not os.environ.get("RUN_LIVE"), reason="needs RUN_LIVE=1")
|
||||
def test_live_embed_batch() -> None:
|
||||
r = embed_texts(["text one", "text two", "text three"])
|
||||
assert len(r) == 3
|
||||
assert all(len(emb) > 100 for emb in r)
|
||||
|
|
@ -1,8 +1,13 @@
|
|||
import json
|
||||
import os
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
||||
|
||||
from src.memory.index import _extract_text, run_index
|
||||
|
||||
|
||||
|
|
@ -39,6 +44,8 @@ class TestRunIndex:
|
|||
assert len(index["files"]) == 1
|
||||
assert index["files"][0]["source"] == "test.md"
|
||||
assert index["files"][0]["embedding"] == [0.1, 0.2, 0.3]
|
||||
assert index["files"][0]["chunk_idx"] == 0
|
||||
assert index["files"][0]["offset"] == 0
|
||||
|
||||
def test_no_md_files(self, tmp_path: Path, capsys) -> None:
|
||||
memory_dir = tmp_path / "empty"
|
||||
|
|
@ -50,3 +57,74 @@ class TestRunIndex:
|
|||
|
||||
captured = capsys.readouterr()
|
||||
assert "No .md files found" in captured.out
|
||||
|
||||
def test_chunking_long_file(self, tmp_path: Path) -> None:
|
||||
memory_dir = tmp_path / "memory"
|
||||
index_dir = tmp_path / ".rag"
|
||||
memory_dir.mkdir(parents=True)
|
||||
|
||||
long_body = "a" * 1500
|
||||
(memory_dir / "long.md").write_text(long_body)
|
||||
|
||||
def fake_embed_texts(texts):
|
||||
return [[0.1, 0.2, 0.3]] * len(texts)
|
||||
|
||||
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
||||
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
||||
run_index(args)
|
||||
|
||||
index = json.loads((index_dir / "index.json").read_text())
|
||||
assert len(index["files"]) >= 3
|
||||
chunk_idxs = [f["chunk_idx"] for f in index["files"]]
|
||||
assert chunk_idxs == sorted(chunk_idxs)
|
||||
assert chunk_idxs[0] == 0
|
||||
assert index["files"][0]["offset"] == 0
|
||||
offsets = [f["offset"] for f in index["files"]]
|
||||
assert offsets == sorted(offsets)
|
||||
|
||||
def test_chunking_size_env_override(
|
||||
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
memory_dir = tmp_path / "memory"
|
||||
index_dir = tmp_path / ".rag"
|
||||
memory_dir.mkdir(parents=True)
|
||||
|
||||
monkeypatch.setenv("MEMORY_CHUNK_SIZE", "256")
|
||||
monkeypatch.setenv("MEMORY_CHUNK_OVERLAP", "0")
|
||||
(memory_dir / "f.md").write_text("a" * 600)
|
||||
|
||||
def fake_embed_texts(texts):
|
||||
return [[0.1]] * len(texts)
|
||||
|
||||
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
||||
args = Namespace(memory_dir=str(memory_dir), output=str(index_dir))
|
||||
run_index(args)
|
||||
|
||||
index = json.loads((index_dir / "index.json").read_text())
|
||||
assert len(index["files"]) == 3
|
||||
for i, f in enumerate(index["files"]):
|
||||
assert f["chunk_idx"] == i
|
||||
assert f["offset"] == i * 256
|
||||
|
||||
def test_chunking_skip_rag_dir(self, tmp_path: Path) -> None:
|
||||
memory_dir = tmp_path / "memory"
|
||||
index_dir = memory_dir / ".rag"
|
||||
memory_dir.mkdir(parents=True)
|
||||
|
||||
(memory_dir / "real.md").write_text("real content")
|
||||
(index_dir).mkdir()
|
||||
(index_dir / "index.json").write_text('{"files": []}')
|
||||
(index_dir / "ignore.md").write_text("should be ignored")
|
||||
|
||||
def fake_embed_texts(texts):
|
||||
return [[0.1]] * len(texts)
|
||||
|
||||
output_dir = tmp_path / "out"
|
||||
with patch("src.memory.index.embed_texts", fake_embed_texts):
|
||||
args = Namespace(memory_dir=str(memory_dir), output=str(output_dir))
|
||||
run_index(args)
|
||||
|
||||
index = json.loads((output_dir / "index.json").read_text())
|
||||
sources = [f["source"] for f in index["files"]]
|
||||
assert "real.md" in sources
|
||||
assert all(".rag" not in s for s in sources)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
import json
|
||||
import os
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("OPENAI_BASE_URL", "http://test/v1")
|
||||
|
||||
import src.memory.search as search_mod
|
||||
from src.memory.search import _cosine_sim
|
||||
|
||||
|
|
@ -56,3 +60,44 @@ class TestSearchOutput:
|
|||
with patch.object(search_mod, "embed_texts", fake_embed_texts):
|
||||
args = Namespace(index_dir=str(index_dir), query="hello", k=5, json=True)
|
||||
search_mod.run_search(args)
|
||||
|
||||
def test_dedup_by_source(self, tmp_path: Path, capsys) -> None:
|
||||
index_dir = tmp_path / ".rag"
|
||||
index_dir.mkdir()
|
||||
index = {
|
||||
"files": [
|
||||
{
|
||||
"source": "doc.md",
|
||||
"text": "chunk 0",
|
||||
"embedding": [1.0, 0.0, 0.0],
|
||||
},
|
||||
{
|
||||
"source": "doc.md",
|
||||
"text": "chunk 1",
|
||||
"embedding": [0.9, 0.1, 0.0],
|
||||
},
|
||||
{
|
||||
"source": "other.md",
|
||||
"text": "chunk 0",
|
||||
"embedding": [0.5, 0.5, 0.0],
|
||||
},
|
||||
],
|
||||
}
|
||||
(index_dir / "index.json").write_text(json.dumps(index))
|
||||
|
||||
def fake_embed_texts(texts):
|
||||
return [[1.0, 0.0, 0.0]]
|
||||
|
||||
with patch.object(search_mod, "embed_texts", fake_embed_texts):
|
||||
args = Namespace(index_dir=str(index_dir), query="q", k=5, json=True)
|
||||
search_mod.run_search(args)
|
||||
|
||||
captured = capsys.readouterr()
|
||||
result = json.loads(captured.out)
|
||||
assert len(result) == 2
|
||||
sources = [r["source"] for r in result]
|
||||
assert "doc.md" in sources
|
||||
assert "other.md" in sources
|
||||
assert sources.count("doc.md") == 1
|
||||
doc_entry = next(r for r in result if r["source"] == "doc.md")
|
||||
assert doc_entry["text"] == "chunk 0"
|
||||
|
|
|
|||
72
uv.lock
generated
72
uv.lock
generated
|
|
@ -373,6 +373,42 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/d2/f0/834e479e47e499b6478e807fb57b31cc2db696c4db30557bb6f5aea4a90b/mando-0.7.1-py2.py3-none-any.whl", hash = "sha256:26ef1d70928b6057ee3ca12583d73c63e05c49de8972d620c278a7b206581a8a", size = 28149, upload-time = "2022-02-24T08:12:25.24Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "memory"
|
||||
version = "0.1.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
{ name = "numpy" },
|
||||
{ name = "tenacity" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
dev = [
|
||||
{ name = "mypy" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-cov" },
|
||||
{ name = "pytest-timeout" },
|
||||
{ name = "ruff" },
|
||||
{ name = "xenon" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "httpx" },
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10" },
|
||||
{ name = "numpy" },
|
||||
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
|
||||
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5.0" },
|
||||
{ name = "pytest-timeout", marker = "extra == 'dev'", specifier = ">=2.2" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.5" },
|
||||
{ name = "tenacity" },
|
||||
{ name = "xenon", marker = "extra == 'dev'", specifier = ">=0.9" },
|
||||
]
|
||||
provides-extras = ["dev"]
|
||||
|
||||
[[package]]
|
||||
name = "mypy"
|
||||
version = "2.3.0"
|
||||
|
|
@ -712,42 +748,6 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/57/c9/e69b1ff4c8b69093ef08b8919ab767af0569666865b39c30a8795d88d3c6/ruff-0.15.22-py3-none-win_arm64.whl", hash = "sha256:e1168075b72158510839f250027659cdd78476f40507dd517892304c41318661", size = 11298172, upload-time = "2026-07-16T15:14:10.51Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "second-brain"
|
||||
version = "0.1.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
{ name = "numpy" },
|
||||
{ name = "tenacity" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
dev = [
|
||||
{ name = "mypy" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-cov" },
|
||||
{ name = "pytest-timeout" },
|
||||
{ name = "ruff" },
|
||||
{ name = "xenon" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "httpx" },
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10" },
|
||||
{ name = "numpy" },
|
||||
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
|
||||
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5.0" },
|
||||
{ name = "pytest-timeout", marker = "extra == 'dev'", specifier = ">=2.2" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.5" },
|
||||
{ name = "tenacity" },
|
||||
{ name = "xenon", marker = "extra == 'dev'", specifier = ">=0.9" },
|
||||
]
|
||||
provides-extras = ["dev"]
|
||||
|
||||
[[package]]
|
||||
name = "six"
|
||||
version = "1.17.0"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue