Skip to content

fix: use stable cache dir for tool embeddings and remove exit() call#26

Open
rtmalikian wants to merge 1 commit into
mims-harvard:mainfrom
rtmalikian:fix/tool-embedding-path-and-error-handling
Open

fix: use stable cache dir for tool embeddings and remove exit() call#26
rtmalikian wants to merge 1 commit into
mims-harvard:mainfrom
rtmalikian:fix/tool-embedding-path-and-error-handling

Conversation

@rtmalikian

Copy link
Copy Markdown

What was broken

Issue #10: FileNotFoundError when the tool embedding .pt cache file is not in the current working directory. The original code saved embeddings to a CWD-relative path, so users who ran TxAgent from a different directory than their first run would never find the cached file.

Issue #6 (partial): The load_tool_desc_embedding method called exit() after generating embeddings on first run, forcing users to restart. Combined with vLLM's spawn multiprocessing requirement, this created confusing RuntimeError messages about bootstrapping.

Additionally, the original code used a bare except: clause that caught ALL exceptions (including KeyboardInterrupt and SystemExit), which masked real errors and made debugging difficult.

What the fix does

  1. Stable cache directory: Embeddings are now saved to ~/.cache/txagent/embeddings/ instead of the current working directory, so they are found regardless of launch location.

  2. Legacy migration: Existing CWD-relative .pt files are automatically moved to the new cache directory for backward compatibility.

  3. Specific exception handling: Replaced bare except: with except FileNotFoundError and except AssertionError — only catches the expected failure modes.

  4. No more exit(): After generating embeddings, the RAG encoder GPU memory is freed (del + gc.collect + torch.cuda.empty_cache()) and the encoder is re-loaded, allowing the main LLM to continue without a restart.

  5. Clear logging: Users now see exactly what is happening — cache miss, generation progress, migration, and memory cleanup.

Verification results

All 6 integration tests pass on macOS (no GPU required — uses mocked SentenceTransformer):

PASS: test_missing_file_generates_embeddings
PASS: test_cached_file_is_loaded
PASS: test_mismatched_tool_count_regenerates
PASS: test_legacy_path_migration
PASS: test_gpu_memory_freed_after_generation
PASS: test_no_exit_call_on_first_run

=== ALL TESTS PASSED ===

Tests verify:

  • Missing file triggers generation and saves to stable cache dir
  • Cached file is loaded without regeneration
  • Mismatched tool count triggers regeneration with clear error message
  • Legacy CWD-relative files are auto-migrated
  • GPU memory is freed after generation (gc.collect called)
  • No exit() call on first run (code continues normally)

Disclosure: This code was developed with assistance from mimo-2.5-pro (Xiaomi) via Hermes Agent (Nous Research). All changes were reviewed, tested against the actual codebase, and verified for correctness.

- Replace bare 'except:' with specific FileNotFoundError/AssertionError
  handling to avoid catching KeyboardInterrupt, SystemExit, etc.
- Use ~/.cache/txagent/embeddings/ instead of CWD-relative paths so the
  cached .pt file is found regardless of working directory
- Auto-migrate legacy CWD-relative .pt files to the new cache location
- Remove exit() call after embedding generation — instead, free the RAG
  encoder GPU memory (del + gc.collect + cuda.empty_cache) and re-load
  the encoder so the main LLM can continue without restart
- Add detailed logging for cache miss, generation, and migration

Fixes mims-harvard#10 (FileNotFoundError when embedding cache is in a different dir)
Fixes mims-harvard#6 (exit() call prevented users from running on first attempt)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant