From 80a4be1d21ffab8ecb5fa0bcbd97f1085bdedfdd Mon Sep 17 00:00:00 2001 From: Sergey <93754860+slaid098@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:13:05 +0300 Subject: [PATCH] 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 --- .env.example | 8 ++ README.md | 2 +- .../032-pr-75-rename-openai-env-chunking.md | 23 ++++ .../pr-75-rename-openai-env-chunking.md | 46 ++++++++ docs/project-map/README.md | 16 +-- pyproject.toml | 5 +- src/memory/cli.py | 2 +- src/memory/embedder.py | 22 +++- src/memory/index.py | 50 +++++--- src/memory/search.py | 12 +- tests/test_chunking.py | 47 ++++++++ tests/test_embedder.py | 110 +++++++++++++++++- tests/test_embedder_live.py | 21 ++++ tests/test_index.py | 78 +++++++++++++ tests/test_search.py | 45 +++++++ uv.lock | 72 ++++++------ 16 files changed, 492 insertions(+), 67 deletions(-) create mode 100644 docs/decisions/032-pr-75-rename-openai-env-chunking.md create mode 100644 docs/handoff/pr-75-rename-openai-env-chunking.md create mode 100644 tests/test_chunking.py create mode 100644 tests/test_embedder_live.py diff --git a/.env.example b/.env.example index 08749a2..9215310 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/README.md b/README.md index 65ea6f4..a0fafc2 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/docs/decisions/032-pr-75-rename-openai-env-chunking.md b/docs/decisions/032-pr-75-rename-openai-env-chunking.md new file mode 100644 index 0000000..6b9fc61 --- /dev/null +++ b/docs/decisions/032-pr-75-rename-openai-env-chunking.md @@ -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, не флаги). \ No newline at end of file diff --git a/docs/handoff/pr-75-rename-openai-env-chunking.md b/docs/handoff/pr-75-rename-openai-env-chunking.md new file mode 100644 index 0000000..c097ace --- /dev/null +++ b/docs/handoff/pr-75-rename-openai-env-chunking.md @@ -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-.md) │ ├── decisions/ # ADRs (NNN-pr--.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=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" diff --git a/src/memory/cli.py b/src/memory/cli.py index 0fbf9ad..d014fce 100644 --- a/src/memory/cli.py +++ b/src/memory/cli.py @@ -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") diff --git a/src/memory/embedder.py b/src/memory/embedder.py index 9c552bb..1732e8e 100644 --- a/src/memory/embedder.py +++ b/src/memory/embedder.py @@ -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: - return _call_embedding_api(client, url, headers, payload) + 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 diff --git a/src/memory/index.py b/src/memory/index.py index b47cd8a..7e39aad 100644 --- a/src/memory/index.py +++ b/src/memory/index.py @@ -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)], } diff --git a/src/memory/search.py b/src/memory/search.py index af52d39..cc77810 100644 --- a/src/memory/search.py +++ b/src/memory/search.py @@ -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)) diff --git a/tests/test_chunking.py b/tests/test_chunking.py new file mode 100644 index 0000000..d8aa721 --- /dev/null +++ b/tests/test_chunking.py @@ -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) diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 6c4956e..9df9021 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -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 diff --git a/tests/test_embedder_live.py b/tests/test_embedder_live.py new file mode 100644 index 0000000..d1c385f --- /dev/null +++ b/tests/test_embedder_live.py @@ -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) diff --git a/tests/test_index.py b/tests/test_index.py index b6a2df0..2d3ab47 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -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) diff --git a/tests/test_search.py b/tests/test_search.py index 367e5b4..276aa72 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -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" diff --git a/uv.lock b/uv.lock index 64a55c7..8f9e90f 100644 --- a/uv.lock +++ b/uv.lock @@ -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"