diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..ee00f21 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,44 @@ +name: CI + +on: + push: + branches: ["main", "review", "review-1"] + pull_request: + +env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Check out + uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install toolchain + run: | + pip install uv + uv pip install -e . + uv pip install ruff pyright typeguard toml-sort yamllint + - name: Lint and format checks + run: make fmt-check + + tests: + runs-on: ubuntu-latest + steps: + - name: Check out + uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install dependencies + run: | + pip install uv + uv pip install -e . + uv pip install pytest + - name: Run pytest + run: make test diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..834aab7 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,11 @@ +# Agent Instructions + +## Documentation Workflow +- Update `CHANGELOG.md` with a new entry every time code changes are committed. +- Maintain `README_LATEST.md` so it always reflects the current implementation; refresh it alongside major feature updates. +- After each iteration, revise `ISSUES.md`, `SOT.md`, `PLAN.md`, `RECOMMENDATIONS.md`, and `TODO.md` to stay in sync with the codebase. + +## Style Guidelines +- Use descriptive Markdown headings starting at level 1 for top-level documents. +- Keep lines to 120 characters or fewer when practical. +- Prefer bullet lists for enumerations instead of inline commas. diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..fa41ae9 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,20 @@ +# Changelog + +## [Unreleased] - 2025-02-14 +### Added +- Vector-only, regex, exact-match, and optional LLM rerank retrieval helpers with reranker utilities and exports. +- MeshMind client wrappers for hybrid, vector, regex, and exact searches plus driver accessors. +- Example script demonstrating triplet storage and diverse retrieval flows. +- Pytest fixtures for encoder and memory factories alongside new retrieval tests that avoid external services. +- Makefile targets for linting, formatting, type checks, and tests, plus a GitHub Actions workflow running lint and pytest. +- README_LATEST.md capturing the current implementation and CHANGELOG.md for release notes. + +### Changed +- Updated `SearchConfig` to support rerank models and refreshed MeshMind documentation across PROJECT, PLAN, SOT, FINDINGS, + DISCREPANCIES, RECOMMENDATIONS, ISSUES, TODO, and NEEDED_FOR_TESTING files. +- Revised `meshmind.retrieval.search` to apply filters centrally, expose new search helpers, and integrate reranking. +- Exposed graph driver access on MeshMind and refreshed retrieval-facing examples and docs. + +### Fixed +- Example ingestion script now uses MeshMind APIs correctly and illustrates relationship persistence. +- Tests rely on fixtures rather than deprecated hooks, improving portability across environments without Memgraph/OpenAI. diff --git a/DISCREPANCIES.md b/DISCREPANCIES.md new file mode 100644 index 0000000..876af91 --- /dev/null +++ b/DISCREPANCIES.md @@ -0,0 +1,55 @@ +# README vs Implementation Discrepancies + +## Overview +- The legacy README still promises a fully featured memory graph with multi-level APIs, relationship storage, and diverse + retrieval methods. Many of those features now exist, but the document remains outdated and should be replaced by + `README_LATEST.md`. +- The current codebase delivers extraction, preprocessing, triplet persistence, CRUD helpers, and expanded retrieval strategies + that were missing when the README was written. +- Remaining gaps primarily involve graph-backed retrieval, observability, and automated infrastructure provisioning. + +## API Surface +- ✅ `MeshMind` now exposes CRUD helpers (`create_memory`, `update_memory`, `delete_memory`, `list_memories`, triplet helpers) + that the README referenced implicitly. +- ✅ Triplet storage routes through `store_triplets` and `MemoryManager.add_triplet`, calling `GraphDriver.upsert_edge`. +- ⚠️ The README still references `register_entity`, `register_allowed_predicates`, and `add_predicate`; predicate management is + handled automatically but there is no public API matching those method names. +- ⚠️ README snippets showing `mesh_mind.store_memory(memory)` should be updated to call `store_memories([memory])` or the new + CRUD helpers. + +## Retrieval Capabilities +- ✅ Vector-only, regex, exact-match, hybrid, BM25, fuzzy, and optional LLM rerank searches exist in `meshmind.retrieval.search` + and are surfaced through `MeshMind` helpers. +- ⚠️ README implies retrieval queries the graph directly. Current search helpers operate on in-memory lists supplied by the + caller; Memgraph-backed retrieval remains future work. +- ⚠️ Named helpers like `search_facts` or `search_procedures` never existed; the README should reference the dispatcher plus + specialized helpers now available. + +## Data & Relationship Modeling +- ✅ Predicates are persisted automatically when storing triplets and tracked in `PredicateRegistry`. +- ⚠️ README examples that look up subjects/objects by name still do not match the implementation, which expects UUIDs. Add + documentation explaining how to resolve names to UUIDs before storing edges. +- ⚠️ Consolidation and expiry remain limited to Celery jobs; README narratives about integrated maintenance still overstate the + current persistence story. + +## Configuration & Dependencies +- ✅ `README_LATEST.md` and `NEEDED_FOR_TESTING.md` document required environment variables, dependency guards, and setup steps. +- ⚠️ The legacy README omits optional tooling now required by the Makefile/CI (ruff, pyright, typeguard, toml-sort, yamllint). +- ⚠️ Python version support in `pyproject.toml` (3.13) still diverges from what many dependencies officially support; update the + documentation or relax the requirement. + +## Example Code Paths +- ✅ Updated example scripts demonstrate extraction, triplet creation, and multiple retrieval strategies. +- ⚠️ Legacy README code that instantiates custom Pydantic entities remains inaccurate; extraction returns `Memory` objects and + validates `entity_label` names only. +- ⚠️ Search examples should be updated to show the new helper functions and optional rerank usage instead of nonexistent + `search_facts`/`search_procedures` calls. + +## Tooling & Operations +- ✅ Makefile and CI workflows now exist, aligning with README promises about automation once the README is refreshed. +- ⚠️ Docker Compose still lacks service definitions for Memgraph/Redis; README setup sections should call this out explicitly. +- ⚠️ Celery tasks remain best-effort shims; README should clarify that maintenance requires the optional infrastructure. + +## Documentation State +- Promote `README_LATEST.md` as the authoritative guide, archive the legacy README, and ensure future updates propagate to + supporting docs (`SOT.md`, `PLAN.md`, `NEEDED_FOR_TESTING.md`). diff --git a/FINDINGS.md b/FINDINGS.md new file mode 100644 index 0000000..658fcb9 --- /dev/null +++ b/FINDINGS.md @@ -0,0 +1,39 @@ +# Findings + +## General Observations +- Core modules are now wired through the `MeshMind` client, including CRUD, triplet storage, and retrieval helpers. Remaining + integration work centers on graph-backed retrieval and maintenance persistence. +- Optional dependencies are largely guarded behind lazy imports or factory functions, improving portability. Environments still + need to install tooling referenced by the Makefile and CI (ruff, pyright, typeguard, toml-sort, yamllint). +- Documentation artifacts (`README_LATEST.md`, `SOT.md`, `NEEDED_FOR_TESTING.md`) stay current when updated with each iteration; + the legacy README should be archived. + +## Dependency & Environment Notes +- `MeshMind` defers Memgraph driver creation until persistence is required, enabling limited workflows without `pymgclient`. +- Encoder registration occurs during bootstrap, but custom deployments must ensure compatible models are registered before + extraction or hybrid search. +- The OpenAI embedding adapter still expects dictionary-like responses; adapting to SDK objects remains on the backlog. +- Celery tasks initialize lazily, yet Redis/Memgraph services are still required at runtime. Docker Compose lacks concrete + service definitions. + +## Data Flow & Persistence +- Triplet storage now persists relationships and tracks predicates automatically, closing an earlier data-loss gap. +- Consolidation and compression utilities operate in memory; persistence of maintenance results is still pending. +- Importance scoring remains a constant fallback; improved heuristics will raise retrieval quality once implemented. + +## CLI & Tooling +- CLI ingestion bootstraps encoders and entities automatically but still assumes Memgraph and OpenAI credentials are configured. +- The Makefile introduces lint, format, type-check, and test targets, plus a Docker helper. External tooling installation is + required before targets succeed. +- GitHub Actions now run formatting checks and pytest on push/PR, providing basic CI coverage. + +## Testing & Quality +- Pytest suites rely on fixtures (`memory_factory`, `dummy_encoder`) to run without external services. Additional coverage is + needed for Celery workflows and graph-backed retrieval when implemented. +- Type checking via `pyright` and runtime checks via `typeguard` are exposed in the Makefile; dependency installation is + necessary for full validation. + +## Documentation +- `README_LATEST.md` supersedes the legacy README and documents setup, pipelines, retrieval, and tooling. +- Supporting docs (`ISSUES.md`, `PLAN.md`, `RECOMMENDATIONS.md`, `SOT.md`) reflect the latest capabilities and highlight remaining + gaps, aiding onboarding and future planning. diff --git a/ISSUES.md b/ISSUES.md new file mode 100644 index 0000000..81b27e5 --- /dev/null +++ b/ISSUES.md @@ -0,0 +1,30 @@ +# Issues Checklist + +## Blockers +- [ ] MeshMind client fails without `mgclient`; introduce lazy driver initialization or documented in-memory fallback. +- [ ] Register a default embedding encoder (OpenAI or sentence-transformers) during startup so extraction and hybrid search can run. +- [ ] Update OpenAI integration to match the current SDK (Responses API payload, embeddings API response structure). +- [ ] Replace eager `tiktoken` imports in `meshmind.core.utils` and `meshmind.pipeline.compress` with guarded, optional imports. +- [ ] Align declared Python requirement with supported dependencies (currently set to Python 3.13 despite ecosystem gaps). + +## High Priority +- [x] Implement relationship persistence (`GraphDriver.upsert_edge`) within the storage pipeline and expose triplet APIs. +- [x] Restore high-level API methods promised in README (`register_entity`, predicate management, `add_memory`, `update_memory`, `delete_memory`). +- [x] Ensure CLI ingestion registers entity models and embedding encoders or fails fast with actionable messaging. +- [x] Provide configuration documentation and examples for Memgraph, Redis, and OpenAI environment variables. +- [x] Add automated tests or smoke checks that run without external services (mock OpenAI, stub Memgraph driver). +- [ ] Create real docker-compose services for Memgraph and Redis or remove the placeholder file. + +## Medium Priority +- [ ] Persist results from consolidation and compression tasks back to the database (currently in-memory only). +- [ ] Refine `Memory.importance` scoring to reflect actual ranking heuristics instead of a constant. +- [x] Add vector, regex, and exact-match search helpers to match stated feature set or update documentation to demote them. +- [ ] Harden Celery tasks to initialize dependencies lazily and log failures when the driver is unavailable. (In progress: lazy driver initialization added, persistence pending) +- [ ] Reconcile tests that depend on `Memory.pre_init` and outdated OpenAI interfaces with the current implementation. +- [x] Add linting, formatting, and type-checking tooling to improve code quality. + +## Low Priority / Nice to Have +- [ ] Offer alternative storage backends (in-memory driver, SQLite, etc.) for easier local development. +- [ ] Provide an administrative dashboard or CLI commands for listing namespaces, counts, and maintenance statistics. +- [ ] Publish onboarding guides and troubleshooting FAQs for contributors. +- [ ] Explore plugin registration for embeddings and retrieval strategies to reduce manual wiring. diff --git a/Makefile b/Makefile index 6693190..7acad1a 100644 --- a/Makefile +++ b/Makefile @@ -1,17 +1,31 @@ -.PHONY: install lint fmt test docker +.PHONY: install lint fmt fmt-check typecheck test check docker clean install: pip install -e . lint: - ruff . + ruff check . fmt: - isort . - black . + ruff format . + +fmt-check: + ruff format --check . + ruff check . + toml-sort --check pyproject.toml + yamllint .github/workflows + +typecheck: + pyright + python -m typeguard --check meshmind test: pytest +check: fmt-check lint typecheck test + +clean: + rm -rf .pytest_cache .ruff_cache + docker: - docker-compose up \ No newline at end of file + docker compose up diff --git a/NEEDED_FOR_TESTING.md b/NEEDED_FOR_TESTING.md new file mode 100644 index 0000000..1594ce0 --- /dev/null +++ b/NEEDED_FOR_TESTING.md @@ -0,0 +1,45 @@ +# Needed for Testing MeshMind + +## Python Runtime +- Python 3.11 or 3.12 is recommended; project metadata currently states 3.13+, but several dependencies (`pymgclient`, + `sentence-transformers`) do not yet publish wheels for 3.13. +- Use a virtual environment (`uv`, `venv`, or `conda`) to isolate dependencies. + +## Python Dependencies +- Install the project editable: `pip install -e .` from the repository root. +- Required packages declared in `pyproject.toml` include `openai`, `pydantic`, `pydantic-settings`, `numpy`, `scikit-learn`, + `rapidfuzz`, `python-dotenv`, `celery[redis]`, `sentence-transformers`, `tiktoken`, and `pymgclient`. +- Development tooling referenced by the Makefile and CI: + - `ruff` for linting and formatting. + - `pyright` for static type checks. + - `typeguard` for runtime type enforcement (`python -m typeguard --check meshmind`). + - `toml-sort` and `yamllint` for configuration validation. +- Optional helpers for local workflows: `pytest-cov`, `pre-commit`, `httpx` (for future service interfaces). + +## External Services and Infrastructure +- **Memgraph** (or compatible Bolt graph database) reachable via `MEMGRAPH_URI` with credentials exported in + `MEMGRAPH_USERNAME`/`MEMGRAPH_PASSWORD`. +- **Redis** for Celery task queues, referenced through `REDIS_URL`. +- **OpenAI API access** for extraction, embeddings, and LLM reranking (`OPENAI_API_KEY`). +- Recommended: Docker Compose or equivalent orchestration to run Memgraph and Redis together when developing locally. + +## Environment Variables +- `OPENAI_API_KEY` — required for extraction, embeddings, and reranking. +- `MEMGRAPH_URI` — e.g., `bolt://localhost:7687`. +- `MEMGRAPH_USERNAME` and `MEMGRAPH_PASSWORD` — credentials for the graph database. +- `REDIS_URL` — optional Redis connection URI (defaults to `redis://localhost:6379/0`). +- `EMBEDDING_MODEL` — encoder key registered with `EncoderRegistry` (defaults to `text-embedding-3-small`). +- Optional overrides for Celery broker/backend if using hosted services. + +## Local Configuration Steps +- Ensure an embedding encoder is registered before extraction or hybrid search. The bootstrap utilities invoked by the CLI and + `MeshMind` constructor handle this, but custom scripts must call `bootstrap_encoders()`. +- Seed demo data as needed using the `examples/extract_preprocess_store_example.py` script after configuring environment + variables. +- Create a `.env` file storing the environment variables above for consistent local configuration. + +## Current Blockers in This Environment +- Neo4j/Memgraph binaries and Docker are unavailable in this workspace, preventing local graph provisioning. +- Redis cannot be installed without container or host-level access; Celery tasks remain untestable locally until a remote + instance is provisioned. +- External network restrictions may limit installation of proprietary packages or access to OpenAI endpoints. diff --git a/NEW_README.md b/NEW_README.md new file mode 100644 index 0000000..0e934f7 --- /dev/null +++ b/NEW_README.md @@ -0,0 +1,136 @@ +# MeshMind + +MeshMind is an experimental memory orchestration service that pairs large language models with a property graph. It extracts +structured `Memory` records from unstructured text, enriches them with embeddings and metadata, and stores both nodes and +relationships via a Memgraph driver. Retrieval helpers operate on in-memory collections today, offering hybrid, vector-only, +regex, exact-match, fuzzy, and BM25 scoring with optional LLM reranking. + +## Status at a Glance +- ✅ `meshmind.client.MeshMind` orchestrates extraction, preprocessing, storage, CRUD helpers, and retrieval wrappers. +- ✅ Pipelines deduplicate memories, normalize importance, compress metadata, and persist nodes and triplets. +- ✅ Retrieval helpers expose hybrid, vector-only, regex, exact-match, BM25, fuzzy, and rerank workflows with namespace/entity + filters. +- ✅ Celery tasks for expiry, consolidation, and compression initialize lazily and run when Redis and Memgraph are configured. +- ✅ Makefile and GitHub Actions provide linting, formatting, type checking, and pytest automation. +- ⚠️ Retrieval still consumes caller-provided lists; graph-backed queries and observability remain on the roadmap. +- ⚠️ Docker Compose does not yet provision Memgraph/Redis services; manual setup is required. + +## Requirements +- Python 3.11 or 3.12 recommended (metadata targets 3.13; verify third-party support before upgrading). +- Memgraph instance reachable via Bolt and the `pymgclient` Python package. +- OpenAI API key for extraction, embeddings, and optional reranking. +- Optional: Redis and Celery for scheduled maintenance tasks. +- Install project dependencies with `pip install -e .`; see `pyproject.toml` for the full list. + +## Installation +1. Create and activate a virtual environment using Python 3.11/3.12 (e.g., `uv venv`, `python -m venv .venv`). +2. Install MeshMind: + ```bash + pip install -e . + ``` +3. Install optional dependencies as needed: + ```bash + pip install pymgclient tiktoken sentence-transformers celery[redis] ruff pyright typeguard toml-sort yamllint + ``` +4. Export required environment variables: + ```bash + export OPENAI_API_KEY=sk-... + export MEMGRAPH_URI=bolt://localhost:7687 + export MEMGRAPH_USERNAME=neo4j + export MEMGRAPH_PASSWORD=secret + export REDIS_URL=redis://localhost:6379/0 + export EMBEDDING_MODEL=text-embedding-3-small + ``` + +## Encoder Registration +`MeshMind` bootstraps encoders and entities during initialization, but custom scripts can register additional encoders: +```python +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder + +if not EncoderRegistry.is_registered("text-embedding-3-small"): + EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder("text-embedding-3-small")) +``` +You may register deterministic or local encoders (e.g., sentence-transformers) for offline testing. + +## Quick Start +```python +from meshmind.client import MeshMind +from meshmind.core.types import Memory, Triplet + +mm = MeshMind() +texts = ["Python is a programming language created by Guido van Rossum."] +memories = mm.extract_memories( + instructions="Extract key facts as Memory objects.", + namespace="demo", + entity_types=[Memory], + content=texts, +) +memories = mm.deduplicate(memories) +memories = mm.score_importance(memories) +memories = mm.compress(memories) +mm.store_memories(memories) + +if len(memories) >= 2: + relation = Triplet( + subject=str(memories[0].uuid), + predicate="RELATED_TO", + object=str(memories[1].uuid), + namespace="demo", + entity_label="Knowledge", + ) + mm.store_triplets([relation]) +``` + +## Retrieval +`MeshMind` exposes multiple retrieval helpers that operate on lists of `Memory` objects (e.g., fetched via +`mm.list_memories(namespace="demo")`). +```python +from meshmind.core.types import SearchConfig + +memories = mm.list_memories(namespace="demo") +config = SearchConfig(encoder=mm.embedding_model, top_k=5, rerank_model="gpt-4o-mini") + +hybrid = mm.search("Python", memories, namespace="demo", config=config, use_llm_rerank=True) +vector_only = mm.search_vector("programming", memories, namespace="demo") +regex_hits = mm.search_regex(r"Guido", memories, namespace="demo") +exact_hits = mm.search_exact("Python", memories, namespace="demo") +``` +The in-memory search helpers support namespace/entity filters and optional reranking via the OpenAI Responses API. Graph-backed +retrieval is planned for a future release. + +## Command-Line Ingestion +```bash +meshmind ingest \ + --namespace demo \ + --instructions "Extract key facts as Memory objects." \ + ./path/to/text/files +``` +The CLI bootstraps encoders/entities automatically. Ensure environment variables are set and Memgraph is reachable. + +## Maintenance Tasks +Celery tasks in `meshmind.tasks.scheduled` provide expiry, consolidation, and compression maintenance. +```bash +celery -A meshmind.tasks.celery_app.app worker -B +``` +Tasks instantiate the driver lazily; provide valid environment variables and ensure Memgraph/Redis are running. + +## Tooling +- **Makefile** – `make fmt`, `make lint`, `make typecheck`, `make test`, `make check`, `make docker`, `make clean`. +- **CI** – `.github/workflows/ci.yml` runs formatting checks (ruff, toml-sort, yamllint) and pytest on push/PR. +- **Examples** – `examples/extract_preprocess_store_example.py` demonstrates ingestion, triplet creation, and multiple retrieval + strategies. + +## Testing +- Run `pytest` to execute the suite; tests rely on fixtures and do not require external services. +- `make typecheck` invokes `pyright` and `typeguard`; install the tooling listed above beforehand. +- See `NEEDED_FOR_TESTING.md` for environment requirements and known blockers (Docker/Memgraph/Redis availability). + +## Known Limitations +- Retrieval operates on in-memory lists; direct Memgraph queries are not yet implemented. +- Consolidation/compression tasks do not persist results back into the graph. +- Observability (structured logging, metrics) is minimal. +- Docker Compose lacks service definitions; infrastructure setup is manual for now. + +## Roadmap Snapshot +Consult `PROJECT.md`, `PLAN.md`, and `RECOMMENDATIONS.md` for prioritized enhancements: graph-backed retrieval, maintenance +persistence, observability, alternative drivers, and service interfaces. diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..670cb9a --- /dev/null +++ b/PLAN.md @@ -0,0 +1,35 @@ +# Plan of Action + +## Phase 1 – Stabilize Runtime Basics ✅ +1. **Dependency Guards** – Implemented lazy driver factories, optional imports, and clear ImportErrors for missing packages. +2. **Default Encoder Registration** – Bootstraps register encoders/entities automatically and the CLI invokes them on startup. +3. **OpenAI SDK Compatibility** – Extraction and embedding adapters align with the Responses API; remaining polish tracked in + `ISSUES.md`. +4. **Configuration Clarity** – `README_LATEST.md` and `NEEDED_FOR_TESTING.md` document environment variables and service setup. + +## Phase 2 – Restore Promised API Surface ✅ +1. **Entity & Predicate Registry Wiring** – `MeshMind` now boots registries and storage persists predicates automatically. +2. **CRUD & Triplet Support** – CRUD helpers and triplet APIs live on `MeshMind` and `MemoryManager`, storing relationships via + `GraphDriver.upsert_edge`. +3. **Relationship-Aware Examples** – Updated example script demonstrates triplet creation and retrieval flows. + +## Phase 3 – Retrieval & Maintenance Enhancements (In Progress) +1. **Search Coverage** – Hybrid, vector-only, regex, exact-match, fuzzy, and LLM rerank helpers are implemented and exposed. + Next: wire graph-backed retrieval queries once Memgraph search endpoints are available. +2. **Maintenance Tasks** – Tasks initialize lazily but still return in-memory results. Persist consolidation outputs and improve + logging. +3. **Importance Scoring Improvements** – Placeholder scoring remains; design data-driven or LLM-assisted heuristics. + +## Phase 4 – Developer Experience & Tooling (In Progress) +1. **Testing Overhaul** – Pytest suites rely on local fixtures with no external services. Extend coverage for Celery workflows + and graph-backed retrieval once implemented. +2. **Automation & CI** – Makefile provides lint/format/type/test targets and CI runs fmt-check + pytest. Add caching and matrix + builds when dependencies stabilize. +3. **Environment Provisioning** – Documented manual setup in `NEEDED_FOR_TESTING.md`. Next step: ship docker-compose services + or scripts for Memgraph and Redis. + +## Phase 5 – Strategic Enhancements (Planned) +1. **Pluggable Storage Backends** – Design in-memory/Neo4j drivers that satisfy `GraphDriver` for easier local development. +2. **Service Interfaces** – Build REST/gRPC surfaces for ingestion and retrieval to support remote agents. +3. **Operational Observability** – Layer structured logging, metrics, and dashboards across pipelines and Celery tasks. +4. **Onboarding & Documentation** – Promote the updated README stack, maintain SOT diagrams, and provide troubleshooting guides. diff --git a/PROJECT.md b/PROJECT.md new file mode 100644 index 0000000..53a8736 --- /dev/null +++ b/PROJECT.md @@ -0,0 +1,79 @@ +# MeshMind Project Overview + +## Vision and Scope +- Transform unstructured text into graph-backed `Memory` records enriched with embeddings and metadata. +- Offer pipelines for extraction, preprocessing, storage, and retrieval that can be orchestrated from CLI tools or bespoke agents. +- Enable background maintenance workflows (expiry, consolidation, compression) once supporting services are provisioned. + +## Current Architecture Snapshot +- **Client façade**: `meshmind.client.MeshMind` composes the OpenAI client, configurable embedding model, registry bootstrap, + and a lazily created `MemgraphDriver`. Retrieval helpers expose hybrid, vector, regex, exact, and reranked flows. +- **Pipelines**: Extraction (LLM + function calling), preprocessing (deduplicate, score, compress), and storage utilities live in + `meshmind.pipeline`. Maintenance helpers consolidate duplicates and expire stale memories. +- **Graph layer**: `meshmind.db` defines `GraphDriver` and implements a Memgraph backend that handles node and relationship + upserts, queries, deletions, and basic vector search fallbacks. +- **Retrieval helpers**: `meshmind.retrieval` now covers BM25, fuzzy, hybrid, vector-only, regex, exact-match, and LLM rerank + workflows with shared filters and reranker utilities. +- **Task runners**: `meshmind.tasks` configures Celery beat to run expiry, consolidation, and compression. Drivers/managers + initialize lazily so import-time failures are avoided. +- **Support code**: `meshmind.core` provides configuration, data models, embeddings, similarity math, and optional dependency + guards around tokenization. +- **Tooling**: The CLI ingest command (`meshmind ingest`), updated example script, Makefile automation, and CI workflow illustrate + extraction → preprocessing → storage → retrieval without custom scripting. + +## Implemented Capabilities +- Serialize knowledge as `Memory` (nodes) and `Triplet` (relationships) Pydantic models with namespaces, metadata, embeddings, + timestamps, TTL, and importance fields. +- Extract structured memories from text via the OpenAI Responses API, with encoder registration handled during bootstrap. +- Deduplicate memories by name and cosine similarity, normalize importance, and compress metadata when `tiktoken` is installed. +- Persist memory nodes and triplet relationships through the storage pipeline and `MemoryManager` CRUD helpers. +- Perform hybrid, vector-only, regex, exact-match, fuzzy, and BM25 retrieval with optional metadata filters and LLM reranking. +- Provide CRUD surfaces on `MeshMind` for creating, updating, deleting, and listing memories and triplets. +- Run Celery maintenance tasks (expiry, consolidation, compression) that tolerate missing graph drivers until runtime. +- Demonstrate ingestion, relationship creation, and retrieval in `examples/extract_preprocess_store_example.py`. +- Automate linting, formatting, type checking, and testing through the Makefile and GitHub Actions. + +## Partially Implemented or Fragile Areas +- The OpenAI embedding wrapper still assumes dictionary-style responses; adjust once SDK models are fully adopted. +- Memgraph remains the only storage backend; introducing an in-memory or Neo4j implementation would improve portability. +- Maintenance tasks compute consolidation results in-process and do not yet write back optimized memories. +- Importance scoring remains heuristic; richer scoring logic or LLM-assisted ranking is still pending. +- Docker Compose does not provision Memgraph/Redis services; local setups must install them manually. +- The declared Python 3.13 requirement may exceed what optional dependencies officially support. + +## Missing or Broken Capabilities +- Retrieval still operates on caller-provided memory lists; direct graph queries for retrieval are not implemented. +- No public API exposes predicate registration beyond automatic registry bootstrap. +- Observability (structured logging, metrics) is minimal across pipelines and tasks. +- Service interfaces (REST/gRPC) for remote ingestion or retrieval have not been built. + +## External Services & Dependencies +- **Memgraph + pymgclient**: Required for persistence. Without `pymgclient`, the default driver cannot connect. +- **OpenAI SDK**: Required for extraction, embeddings, and LLM reranking; configure `OPENAI_API_KEY`. +- **tiktoken**: Optional but necessary for compression/token budgeting. +- **RapidFuzz, scikit-learn, numpy**: Support fuzzy and lexical retrieval. +- **Celery + Redis**: Optional but necessary for scheduled maintenance jobs. +- **sentence-transformers**: Optional embedding backend for offline models. +- **ruff, pyright, typeguard, toml-sort, yamllint**: Development tooling invoked by the Makefile and CI workflow. + +## Tooling and Operational State +- `Makefile` exposes `lint`, `fmt`, `fmt-check`, `typecheck`, `test`, `check`, `docker`, and `clean` targets. +- `.github/workflows/ci.yml` runs formatting/linting checks and pytest on push and pull requests. +- Tests now rely on fixtures (`memory_factory`, `dummy_encoder`) and avoid external services, though optional dependencies may + still need installation. +- Neo4j/Memgraph services are not provisioned automatically; see `NEEDED_FOR_TESTING.md` for setup requirements. + +## Roadmap Highlights +- Implement graph-backed retrieval queries (vector similarity, structured filters) instead of memory list inputs. +- Add observability (logging, metrics, tracing) across pipelines and Celery tasks. +- Enhance importance scoring with data-driven heuristics or LLM evaluation. +- Harden Celery tasks to persist consolidation output and surface failures clearly. +- Provide alternative graph drivers (Neo4j, in-memory) and docker-compose services for required infrastructure. +- Build service interfaces (REST/gRPC) for ingestion and retrieval, enabling integration with other systems. +- Continue refining documentation to reflect setup, troubleshooting, and architectural decisions. + +## Future Potential Extensions +- Plugin-based encoder and retriever registration for runtime extensibility. +- Streaming ingestion workers (queues, webhooks) beyond batch CLI workflows. +- UI or agent-facing dashboards for curation, monitoring, and analytics. +- Automated CI pipelines for release packaging, schema migrations, and integration tests. diff --git a/README_LATEST.md b/README_LATEST.md new file mode 100644 index 0000000..0e934f7 --- /dev/null +++ b/README_LATEST.md @@ -0,0 +1,136 @@ +# MeshMind + +MeshMind is an experimental memory orchestration service that pairs large language models with a property graph. It extracts +structured `Memory` records from unstructured text, enriches them with embeddings and metadata, and stores both nodes and +relationships via a Memgraph driver. Retrieval helpers operate on in-memory collections today, offering hybrid, vector-only, +regex, exact-match, fuzzy, and BM25 scoring with optional LLM reranking. + +## Status at a Glance +- ✅ `meshmind.client.MeshMind` orchestrates extraction, preprocessing, storage, CRUD helpers, and retrieval wrappers. +- ✅ Pipelines deduplicate memories, normalize importance, compress metadata, and persist nodes and triplets. +- ✅ Retrieval helpers expose hybrid, vector-only, regex, exact-match, BM25, fuzzy, and rerank workflows with namespace/entity + filters. +- ✅ Celery tasks for expiry, consolidation, and compression initialize lazily and run when Redis and Memgraph are configured. +- ✅ Makefile and GitHub Actions provide linting, formatting, type checking, and pytest automation. +- ⚠️ Retrieval still consumes caller-provided lists; graph-backed queries and observability remain on the roadmap. +- ⚠️ Docker Compose does not yet provision Memgraph/Redis services; manual setup is required. + +## Requirements +- Python 3.11 or 3.12 recommended (metadata targets 3.13; verify third-party support before upgrading). +- Memgraph instance reachable via Bolt and the `pymgclient` Python package. +- OpenAI API key for extraction, embeddings, and optional reranking. +- Optional: Redis and Celery for scheduled maintenance tasks. +- Install project dependencies with `pip install -e .`; see `pyproject.toml` for the full list. + +## Installation +1. Create and activate a virtual environment using Python 3.11/3.12 (e.g., `uv venv`, `python -m venv .venv`). +2. Install MeshMind: + ```bash + pip install -e . + ``` +3. Install optional dependencies as needed: + ```bash + pip install pymgclient tiktoken sentence-transformers celery[redis] ruff pyright typeguard toml-sort yamllint + ``` +4. Export required environment variables: + ```bash + export OPENAI_API_KEY=sk-... + export MEMGRAPH_URI=bolt://localhost:7687 + export MEMGRAPH_USERNAME=neo4j + export MEMGRAPH_PASSWORD=secret + export REDIS_URL=redis://localhost:6379/0 + export EMBEDDING_MODEL=text-embedding-3-small + ``` + +## Encoder Registration +`MeshMind` bootstraps encoders and entities during initialization, but custom scripts can register additional encoders: +```python +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder + +if not EncoderRegistry.is_registered("text-embedding-3-small"): + EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder("text-embedding-3-small")) +``` +You may register deterministic or local encoders (e.g., sentence-transformers) for offline testing. + +## Quick Start +```python +from meshmind.client import MeshMind +from meshmind.core.types import Memory, Triplet + +mm = MeshMind() +texts = ["Python is a programming language created by Guido van Rossum."] +memories = mm.extract_memories( + instructions="Extract key facts as Memory objects.", + namespace="demo", + entity_types=[Memory], + content=texts, +) +memories = mm.deduplicate(memories) +memories = mm.score_importance(memories) +memories = mm.compress(memories) +mm.store_memories(memories) + +if len(memories) >= 2: + relation = Triplet( + subject=str(memories[0].uuid), + predicate="RELATED_TO", + object=str(memories[1].uuid), + namespace="demo", + entity_label="Knowledge", + ) + mm.store_triplets([relation]) +``` + +## Retrieval +`MeshMind` exposes multiple retrieval helpers that operate on lists of `Memory` objects (e.g., fetched via +`mm.list_memories(namespace="demo")`). +```python +from meshmind.core.types import SearchConfig + +memories = mm.list_memories(namespace="demo") +config = SearchConfig(encoder=mm.embedding_model, top_k=5, rerank_model="gpt-4o-mini") + +hybrid = mm.search("Python", memories, namespace="demo", config=config, use_llm_rerank=True) +vector_only = mm.search_vector("programming", memories, namespace="demo") +regex_hits = mm.search_regex(r"Guido", memories, namespace="demo") +exact_hits = mm.search_exact("Python", memories, namespace="demo") +``` +The in-memory search helpers support namespace/entity filters and optional reranking via the OpenAI Responses API. Graph-backed +retrieval is planned for a future release. + +## Command-Line Ingestion +```bash +meshmind ingest \ + --namespace demo \ + --instructions "Extract key facts as Memory objects." \ + ./path/to/text/files +``` +The CLI bootstraps encoders/entities automatically. Ensure environment variables are set and Memgraph is reachable. + +## Maintenance Tasks +Celery tasks in `meshmind.tasks.scheduled` provide expiry, consolidation, and compression maintenance. +```bash +celery -A meshmind.tasks.celery_app.app worker -B +``` +Tasks instantiate the driver lazily; provide valid environment variables and ensure Memgraph/Redis are running. + +## Tooling +- **Makefile** – `make fmt`, `make lint`, `make typecheck`, `make test`, `make check`, `make docker`, `make clean`. +- **CI** – `.github/workflows/ci.yml` runs formatting checks (ruff, toml-sort, yamllint) and pytest on push/PR. +- **Examples** – `examples/extract_preprocess_store_example.py` demonstrates ingestion, triplet creation, and multiple retrieval + strategies. + +## Testing +- Run `pytest` to execute the suite; tests rely on fixtures and do not require external services. +- `make typecheck` invokes `pyright` and `typeguard`; install the tooling listed above beforehand. +- See `NEEDED_FOR_TESTING.md` for environment requirements and known blockers (Docker/Memgraph/Redis availability). + +## Known Limitations +- Retrieval operates on in-memory lists; direct Memgraph queries are not yet implemented. +- Consolidation/compression tasks do not persist results back into the graph. +- Observability (structured logging, metrics) is minimal. +- Docker Compose lacks service definitions; infrastructure setup is manual for now. + +## Roadmap Snapshot +Consult `PROJECT.md`, `PLAN.md`, and `RECOMMENDATIONS.md` for prioritized enhancements: graph-backed retrieval, maintenance +persistence, observability, alternative drivers, and service interfaces. diff --git a/RECOMMENDATIONS.md b/RECOMMENDATIONS.md new file mode 100644 index 0000000..375c2b5 --- /dev/null +++ b/RECOMMENDATIONS.md @@ -0,0 +1,31 @@ +# Recommendations + +## Stabilize the Foundation +- Maintain lazy initialization for optional dependencies and continue testing environments without Memgraph or OpenAI access. +- Align declared Python support with dependency compatibility (target 3.11/3.12 until third parties certify 3.13). +- Harden the OpenAI embedding adapter to consume SDK response objects directly and surface actionable errors for rate limits. + +## Restore and Extend Functionality +- Implement graph-backed retrieval queries so callers are not required to materialize memory lists in Python. +- Persist consolidation outputs back into the graph and close the loop on maintenance workflows. +- Design richer importance scoring heuristics (analytics-driven or LLM-evaluated) to replace the current constant fallback. +- Expand predicate/registry management APIs so custom relationship vocabularies can be registered explicitly. + +## Improve Developer Experience +- Provision docker-compose services (or scripts) for Memgraph and Redis to streamline local setup. +- Add Makefile targets for running Celery workers and seeding demo data once infrastructure is provisioned. +- Broaden pytest coverage to include Celery tasks, graph-backed retrieval, and error handling scenarios. +- Cache dependencies and split lint/test jobs in CI for faster feedback once the dependency stack stabilizes. + +## Documentation & Onboarding +- Keep `README_LATEST.md`, `SOT.md`, and onboarding guides synchronized with each release; document rerank, retrieval, and + registry flows with diagrams when possible. +- Publish troubleshooting sections for missing optional tooling (ruff, pyright, typeguard, toml-sort, yamllint) now referenced in + the Makefile. +- Provide walkthroughs for configuring LLM reranking, including sample prompts and response expectations. + +## Future Enhancements +- Introduce alternative storage drivers (Neo4j, in-memory) and plug-in registration for embeddings/retrievers. +- Expose REST/gRPC services for ingestion and retrieval so external agents can integrate without importing Python modules. +- Instrument ingestion and maintenance with structured logging and metrics to support observability and alerting. +- Explore streaming ingestion pipelines (queues, webhooks) for near-real-time updates. diff --git a/SOT.md b/SOT.md new file mode 100644 index 0000000..ef56455 --- /dev/null +++ b/SOT.md @@ -0,0 +1,113 @@ +# MeshMind Source of Truth + +This document summarizes the current architecture, supporting assets, and operational expectations for MeshMind. Update it +whenever workflows or modules change so new contributors can find accurate information in one place. + +## Repository Layout +``` +meshmind/ +├── api/ # MemoryManager CRUD adapter that wraps a GraphDriver +├── cli/ # CLI entry point and ingest command +├── client.py # High-level MeshMind façade and orchestration helpers +├── core/ # Configuration, embeddings, types, similarity, shared utilities +├── db/ # Graph driver abstractions plus Memgraph implementation +├── models/ # Entity and predicate registries shared across the pipeline +├── pipeline/ # Extract, preprocess, compression, storage, consolidation, expiry stages +├── retrieval/ # Search strategies (hybrid, lexical, fuzzy, vector, regex, rerank helpers) +├── tasks/ # Celery beat schedules and maintenance jobs +├── tests/ # Pytest suites with local fixtures (no external services required) +└── examples/ # Scripts and notebooks showing ingestion and retrieval flows +``` +Supporting assets: +- `Makefile`: Development automation (linting, formatting, type checks, tests, docker compose). +- `.github/workflows/ci.yml`: GitHub Actions workflow running linting/formatting checks and pytest. +- `pyproject.toml`: Project metadata and dependency list (targets Python 3.13, see blockers in `ISSUES.md`). +- Documentation (`PROJECT.md`, `PLAN.md`, `SOT.md`, `README_LATEST.md`, etc.) describing the system and roadmap. + +## Configuration (`meshmind/core/config.py`) +- Loads environment variables for Memgraph (`MEMGRAPH_URI`, `MEMGRAPH_USERNAME`, `MEMGRAPH_PASSWORD`), Redis + (`REDIS_URL`), OpenAI (`OPENAI_API_KEY`), and the default embedding model (`EMBEDDING_MODEL`). +- Uses `python-dotenv` when available to hydrate values from a `.env` file automatically. +- Provides a module-level `settings` instance used across the client, drivers, CLI, and Celery tasks. + +## Core Data Models (`meshmind/core/types.py`) +- `Memory`: Pydantic model that represents a knowledge record, including embeddings, metadata, and optional TTL. +- `Triplet`: Subject–predicate–object edge connecting two memory UUIDs with namespace and metadata. +- `SearchConfig`: Retrieval configuration (encoder name, `top_k`, `rerank_k`, optional rerank model, metadata filters, + hybrid weights). + +## Client (`meshmind/client.py`) +- `MeshMind` bootstraps: + - Default OpenAI client (Responses API) when the SDK is installed; custom clients can be injected for testing. + - Embedding model from configuration with encoder bootstrap that registers available adapters. + - Graph driver factory that creates a `MemgraphDriver` lazily only when persistence is required. +- Provides convenience helpers: + - Pipelines: `extract_memories`, `deduplicate`, `score_importance`, `compress`, `store_memories`, `store_triplets`. + - CRUD: `create_memory`, `update_memory`, `delete_memory`, `get_memory`, `list_memories`, `list_triplets`. + - Retrieval: `search` (hybrid + optional LLM rerank), `search_vector`, `search_regex`, `search_exact`. +- Exposes `graph_driver`/`driver` properties that surface the active graph driver instance on demand. + +## Embeddings & Utilities (`meshmind/core/embeddings.py`, `meshmind/core/utils.py`) +- `EncoderRegistry` manages encoder instances (OpenAI embeddings, sentence-transformers, custom fixtures). +- OpenAI and SentenceTransformer adapters provide encoding with retry logic and optional fallbacks. +- Utility functions provide UUIDs, timestamps, hashing, and token counting guarded behind optional `tiktoken` imports. + +## Database Layer (`meshmind/db`) +- `GraphDriver` defines the persistence contract (entity upserts, relationship upserts, querying, deletions, triplet listing). +- `MemgraphDriver` wraps `mgclient`, handles URI parsing, executes Cypher statements, and exposes a Python-side vector search + fallback when database-native similarity is unavailable. + +## Pipeline Modules (`meshmind/pipeline`) +1. **Extraction (`extract.py`)** – Orchestrates OpenAI function calling against the `Memory` schema, enforces entity label filters, + and populates embeddings via registered encoders. +2. **Preprocess (`preprocess.py`)** – Deduplicates by name/embedding similarity, ensures memories have importance scores, and + delegates to compression when available. +3. **Compress (`compress.py`)** – Truncates metadata payloads to configurable token budgets when `tiktoken` is installed. +4. **Store (`store.py`)** – Persists memories and triplets using the configured `GraphDriver`, registering predicates as needed. +5. **Consolidate & Expire (`consolidate.py`, `expire.py`)** – Maintenance utilities triggered by Celery tasks to group memories + and remove stale entries. + +## Retrieval (`meshmind/retrieval`) +- `filters.py`: Namespace, entity label, and metadata filtering helpers. +- `bm25.py`, `fuzzy.py`: Lexical and fuzzy scorers using scikit-learn TF-IDF + cosine and RapidFuzz WRatio respectively. +- `vector.py`: Vector-only search utilities with cosine similarity and optional precomputed query embeddings. +- `hybrid.py`: Combines vector and BM25 scores with configurable weights defined in `SearchConfig`. +- `search.py`: Dispatchers for hybrid, BM25, fuzzy, vector, regex, and exact-match search modes plus metadata filters. +- `rerank.py`: Generic reranker interface and LLM-based rerank helper compatible with the OpenAI Responses API. + +## CLI (`meshmind/cli`) +- `meshmind.cli.__main__`: Entry point exposing an `ingest` command for local pipelines. +- CLI bootstraps encoder and entity registries, validates configuration early, and surfaces actionable errors when optional + dependencies are missing. + +## Tasks (`meshmind/tasks`) +- `celery_app.py`: Creates the Celery application lazily, returning a shim when Celery is not installed. +- `scheduled.py`: Defines periodic consolidation, compression, and expiry jobs that now initialize drivers and managers lazily + to tolerate missing dependencies during import. + +## API Adapter (`meshmind/api/memory_manager.py`) +- Manages CRUD operations against the graph driver, including triplet persistence and deletion helpers. +- Returns Pydantic models for list/get operations and gracefully handles missing records. + +## Models (`meshmind/models/registry.py`) +- `EntityRegistry` and `PredicateRegistry` store class metadata and permitted predicates. +- Registries are populated during bootstrap and extended as new entity/predicate types are defined. + +## Examples & Tests +- `examples/extract_preprocess_store_example.py`: Demonstrates extraction, preprocessing, triplet creation, and multiple + retrieval strategies. +- `meshmind/tests`: Pytest suites rely on fixtures (`memory_factory`, `dummy_encoder`) and pure-Python stubs, allowing the + suite to run without Memgraph, OpenAI, or Redis dependencies. + +## External Dependencies +- Required: `openai`, `pydantic`, `pydantic-settings`, `numpy`, `scikit-learn`, `rapidfuzz`, `python-dotenv`, `pymgclient`. +- Optional but supported: `tiktoken`, `sentence-transformers`, `celery[redis]`, `typeguard`, `pyright`, `toml-sort`, `yamllint`. +- Development tooling introduced in the Makefile/CI expects `ruff`, `pyright`, `typeguard`, `toml-sort`, and `yamllint`. + +## Operational Notes +- Graph persistence requires a running Memgraph instance reachable via `settings.MEMGRAPH_URI` and `pymgclient` installed. +- Encoder registration occurs during bootstrap; ensure at least one embedding encoder is available before extraction/search. +- LLM reranking uses the OpenAI Responses API. Provide `OPENAI_API_KEY` and confirm the selected `SearchConfig.rerank_model` is + deployed to your account. +- Local development commands rely on external tooling (ruff, pyright, typeguard, toml-sort, yamllint); install them via the + Makefile or the CI workflow instructions in `README_LATEST.md`. diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..ab2e7ed --- /dev/null +++ b/TODO.md @@ -0,0 +1,19 @@ +# TODO + +- [x] Implement dependency guards and lazy imports for optional packages (`mgclient`, `tiktoken`, `celery`, `sentence-transformers`). +- [x] Add bootstrap helper for default encoder registration and call it from the CLI. +- [x] Update OpenAI encoder implementation to align with latest SDK responses and retry semantics. +- [x] Improve configuration guidance and automation for environment variables and service setup. +- [x] Wire `EntityRegistry` and `PredicateRegistry` into the storage pipeline and client. +- [x] Implement CRUD and triplet methods on `MeshMind`, including relationship persistence in `GraphDriver`. +- [x] Refresh examples to cover relationship-aware ingestion and retrieval flows. +- [x] Extend retrieval module with vector-only, regex, exact-match, and optional LLM rerank search helpers. +- [ ] Harden Celery maintenance tasks to initialize drivers lazily and persist consolidation results. +- [ ] Replace constant importance scoring with a data-driven or LLM-assisted heuristic. +- [x] Modernize pytest suites and add fixtures to run without external services. +- [x] Expand Makefile and add CI workflows for linting, testing, and type checks. +- [ ] Document or provision local Memgraph and Redis services (e.g., via docker-compose) for onboarding. +- [ ] Abstract `GraphDriver` to support alternative storage backends (Neo4j, in-memory, SQLite prototype). +- [ ] Add service interfaces (REST/gRPC) for ingestion and retrieval. +- [ ] Introduce observability (logging, metrics) for ingestion and maintenance pipelines. +- [ ] Promote NEW_README.md, archive legacy README, and maintain SOT diagrams and maps. diff --git a/examples/extract_preprocess_store_example.py b/examples/extract_preprocess_store_example.py index 74ac65d..a12511a 100644 --- a/examples/extract_preprocess_store_example.py +++ b/examples/extract_preprocess_store_example.py @@ -1,40 +1,62 @@ -""" -Example flow: extract → preprocess → store using Meshmind pipeline. -Requires a running Memgraph instance and a valid OPENAI_API_KEY. -""" -from meshmind.core.types import Memory +"""End-to-end MeshMind example covering extraction, storage, and retrieval.""" +from __future__ import annotations + from meshmind.client import MeshMind +from meshmind.core.types import Memory, Triplet + -def main(): - # Initialize MeshMind client (uses OpenAI and default MemgraphDriver) +def main() -> None: mm = MeshMind() - driver = mm.driver - # Sample content for extraction texts = [ "The Eiffel Tower is located in Paris and was built in 1889.", - "Python is a programming language created by Guido van Rossum." + "Python is a programming language created by Guido van Rossum.", ] - # Extract memories via LLM memories = mm.extract_memories( instructions="Extract key facts as Memory objects.", namespace="demo", - entity_types=[Memory], + entity_types=[Memory], content=texts, ) - print(f"Extracted {len(memories)} memories:") - for m in memories: - print(m.json()) - - # Preprocess: deduplicate, score importance, compress - memories = mm.deduplicate(memories, threshold=0.9) + memories = mm.deduplicate(memories) memories = mm.score_importance(memories) memories = mm.compress(memories) - - # Store into graph mm.store_memories(memories) - print("Memories stored to graph.") + print(f"Stored {len(memories)} memories.") + + if len(memories) >= 2: + relation = Triplet( + subject=str(memories[0].uuid), + predicate="RELATED_TO", + object=str(memories[1].uuid), + namespace="demo", + entity_label="Knowledge", + metadata={"confidence": 0.9}, + ) + mm.store_triplets([relation]) + print("Stored relationship between first two memories.") + + hits = mm.search("Eiffel Tower", memories, namespace="demo") + print("Hybrid search results:") + for mem in hits: + print(f"- {mem.name} (importance={mem.importance})") + + vector_hits = mm.search_vector("programming", memories, namespace="demo") + print("Vector-only search results:") + for mem in vector_hits: + print(f"- {mem.name}") + + regex_hits = mm.search_regex(r"Paris", memories, namespace="demo") + print("Regex search results:") + for mem in regex_hits: + print(f"- {mem.name}") + + exact_hits = mm.search_exact("Python", memories, fields=["name"], namespace="demo") + print("Exact match search results:") + for mem in exact_hits: + print(f"- {mem.name}") + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/meshmind/api/memory_manager.py b/meshmind/api/memory_manager.py index 935d270..bbd73a0 100644 --- a/meshmind/api/memory_manager.py +++ b/meshmind/api/memory_manager.py @@ -1,39 +1,50 @@ -from typing import Any, List, Optional +from __future__ import annotations + +from typing import Any, Dict, List, Optional from uuid import UUID +from pydantic import BaseModel + +from meshmind.core.types import Memory, Triplet + + class MemoryManager: - """ - Mid-level CRUD interface for Memory objects, delegating to an underlying graph driver. - """ + """Mid-level CRUD interface for ``Memory`` and ``Triplet`` objects.""" + def __init__(self, graph_driver: Any): # pragma: no cover self.driver = graph_driver - def add_memory(self, memory: Any) -> UUID: + @staticmethod + def _props(model: Any) -> Dict[str, Any]: + if isinstance(model, BaseModel): + return model.dict(exclude_none=True) + if hasattr(model, "dict"): + try: + return model.dict(exclude_none=True) # type: ignore[attr-defined] + except TypeError: + pass + if isinstance(model, dict): + return {k: v for k, v in model.items() if v is not None} + return {k: v for k, v in model.__dict__.items() if v is not None} + + def add_memory(self, memory: Memory) -> UUID: """ Add a new Memory object to the graph. :param memory: A Memory-like object to be stored. :return: The UUID of the newly added memory. """ - # Upsert the memory object into the graph - try: - props = memory.dict(exclude_none=True) - except Exception: - props = memory.__dict__ + props = self._props(memory) self.driver.upsert_entity(memory.entity_label, memory.name, props) return memory.uuid - def update_memory(self, memory: Any) -> None: + def update_memory(self, memory: Memory) -> None: """ Update an existing Memory object in the graph. :param memory: A Memory-like object with updated fields. """ - # Update an existing memory via upsert - try: - props = memory.dict(exclude_none=True) - except Exception: - props = memory.__dict__ + props = self._props(memory) self.driver.upsert_entity(memory.entity_label, memory.name, props) def delete_memory(self, memory_id: UUID) -> None: @@ -53,8 +64,6 @@ def get_memory(self, memory_id: UUID) -> Optional[Any]: :return: Memory-like object or None if not found. """ # Retrieve a memory by UUID - from meshmind.core.types import Memory - cypher = "MATCH (m) WHERE m.uuid = $uuid RETURN m" params = {"uuid": str(memory_id)} records = self.driver.find(cypher, params) @@ -68,7 +77,7 @@ def get_memory(self, memory_id: UUID) -> Optional[Any]: except Exception: return None - def list_memories(self, namespace: Optional[str] = None) -> List[Any]: + def list_memories(self, namespace: Optional[str] = None) -> List[Memory]: """ List Memory objects, optionally filtered by namespace. @@ -76,8 +85,6 @@ def list_memories(self, namespace: Optional[str] = None) -> List[Any]: :return: List of Memory-like objects. """ # List memories, optionally filtered by namespace - from meshmind.core.types import Memory - if namespace: cypher = "MATCH (m) WHERE m.namespace = $namespace RETURN m" params = {"namespace": namespace} @@ -85,11 +92,51 @@ def list_memories(self, namespace: Optional[str] = None) -> List[Any]: cypher = "MATCH (m) RETURN m" params = {} records = self.driver.find(cypher, params) - result: List[Any] = [] + result: List[Memory] = [] for record in records: data = record.get('m', record) try: result.append(Memory(**data)) except Exception: continue - return result \ No newline at end of file + return result + + def add_triplet(self, triplet: Triplet) -> None: + """Persist or update a ``Triplet`` relationship.""" + + props = self._props(triplet) + namespace = props.pop("namespace", None) + if namespace is not None: + props["namespace"] = namespace + self.driver.upsert_edge( + triplet.subject, + triplet.predicate, + triplet.object, + props, + ) + + def delete_triplet(self, subj: str, predicate: str, obj: str) -> None: + """Remove a relationship identified by subject/predicate/object.""" + + self.driver.delete_triplet(subj, predicate, obj) + + def list_triplets(self, namespace: Optional[str] = None) -> List[Triplet]: + """Return stored ``Triplet`` objects, optionally filtered by namespace.""" + + records = self.driver.list_triplets(namespace) + result: List[Triplet] = [] + for record in records: + data = { + "subject": record.get("subject"), + "predicate": record.get("predicate"), + "object": record.get("object"), + "namespace": record.get("namespace") or namespace, + "entity_label": record.get("predicate", "Relation"), + "metadata": record.get("metadata") or {}, + "reference_time": record.get("reference_time"), + } + try: + result.append(Triplet(**data)) + except Exception: + continue + return result diff --git a/meshmind/cli/__main__.py b/meshmind/cli/__main__.py index 24dab53..6a7a304 100644 --- a/meshmind/cli/__main__.py +++ b/meshmind/cli/__main__.py @@ -6,6 +6,8 @@ import sys from meshmind.cli.ingest import ingest_command +from meshmind.core.bootstrap import bootstrap_encoders, bootstrap_entities +from meshmind.core.config import settings def main(): @@ -35,6 +37,18 @@ def main(): args = parser.parse_args() + # Ensure default encoders and entities are registered before executing commands + bootstrap_entities() + bootstrap_encoders() + + missing = settings.missing() + if missing: + for group, keys in missing.items(): + print( + f"Warning: missing configuration for {group}: {', '.join(keys)}", + file=sys.stderr, + ) + if args.command == "ingest": ingest_command(args) else: @@ -43,4 +57,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/meshmind/client.py b/meshmind/client.py index 365eac0..8f1d070 100644 --- a/meshmind/client.py +++ b/meshmind/client.py @@ -1,80 +1,276 @@ -""" -MeshMind client combining LLM, embedding, and graph driver. -""" -from openai import OpenAI -from typing import Any, List, Type -from meshmind.db.memgraph_driver import MemgraphDriver +"""High-level MeshMind client orchestrating ingestion and storage flows.""" +from __future__ import annotations + +from typing import Any, Callable, Iterable, List, Optional, Sequence, Type +from uuid import UUID + +try: # pragma: no cover - optional dependency + from openai import OpenAI +except ImportError: # pragma: no cover - optional dependency + OpenAI = None # type: ignore + +from meshmind.api.memory_manager import MemoryManager +from meshmind.core.bootstrap import bootstrap_entities, bootstrap_encoders from meshmind.core.config import settings +from meshmind.core.types import Memory, Triplet, SearchConfig +from meshmind.db.base_driver import GraphDriver +from meshmind.db.memgraph_driver import MemgraphDriver +from meshmind.models.registry import EntityRegistry, PredicateRegistry class MeshMind: - """ - High-level client to manage extraction, preprocessing, and storage of memories. - """ + """High-level orchestration client for extraction, preprocessing, and persistence.""" + def __init__( self, llm_client: Any = None, embedding_model: str | None = None, - graph_driver: Any = None, + graph_driver: Optional[GraphDriver] = None, + graph_driver_factory: Callable[[], GraphDriver] | None = None, ): - # Initialize LLM client - self.llm_client = llm_client or OpenAI() - # Set embedding model name + if llm_client is None: + if OpenAI is None: + raise ImportError( + "openai package is required to construct a default MeshMind LLM client." + ) + client_kwargs: dict[str, Any] = {} + if settings.OPENAI_API_KEY: + client_kwargs["api_key"] = settings.OPENAI_API_KEY + llm_client = OpenAI(**client_kwargs) + + self.llm_client = llm_client self.embedding_model = embedding_model or settings.EMBEDDING_MODEL - # Initialize graph driver - self.driver = graph_driver or MemgraphDriver( - settings.MEMGRAPH_URI, - settings.MEMGRAPH_USERNAME, - settings.MEMGRAPH_PASSWORD, + + self._graph_driver: Optional[GraphDriver] = graph_driver + self._graph_driver_factory = graph_driver_factory + if self._graph_driver is None and self._graph_driver_factory is None: + self._graph_driver_factory = lambda: MemgraphDriver( # type: ignore[assignment] + settings.MEMGRAPH_URI, + settings.MEMGRAPH_USERNAME, + settings.MEMGRAPH_PASSWORD, + ) + + self._memory_manager: Optional[MemoryManager] = ( + MemoryManager(self._graph_driver) if self._graph_driver else None ) + self.entity_registry = EntityRegistry + self.predicate_registry = PredicateRegistry + bootstrap_entities([Memory]) + bootstrap_encoders() + + @property + def graph_driver(self) -> GraphDriver: + """Expose the active graph driver, creating it on demand.""" + return self._ensure_driver() + + @property + def driver(self) -> GraphDriver: + """Backward compatible alias for :attr:`graph_driver`.""" + return self.graph_driver + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + def _ensure_driver(self) -> GraphDriver: + if self._graph_driver is None: + if self._graph_driver_factory is None: + raise RuntimeError("No graph driver factory available for MeshMind") + self._graph_driver = self._graph_driver_factory() + return self._graph_driver + + def _ensure_manager(self) -> MemoryManager: + if self._memory_manager is None: + self._memory_manager = MemoryManager(self._ensure_driver()) + return self._memory_manager + # ------------------------------------------------------------------ + # Pipelines + # ------------------------------------------------------------------ def extract_memories( self, instructions: str, namespace: str, - entity_types: List[Type[Any]], - content: List[str], + entity_types: Sequence[Type[Any]], + content: Sequence[str], ) -> List[Any]: from meshmind.pipeline.extract import extract_memories return extract_memories( instructions=instructions, namespace=namespace, - entity_types=entity_types, + entity_types=list(entity_types), embedding_model=self.embedding_model, - content=content, + content=list(content), llm_client=self.llm_client, ) def deduplicate( self, - memories: List[Any], + memories: Sequence[Any], threshold: float = 0.95, ) -> List[Any]: from meshmind.pipeline.preprocess import deduplicate - return deduplicate(memories, threshold) + return deduplicate(list(memories), threshold) def score_importance( self, - memories: List[Any], + memories: Sequence[Any], ) -> List[Any]: from meshmind.pipeline.preprocess import score_importance - return score_importance(memories) + return score_importance(list(memories)) def compress( self, - memories: List[Any], + memories: Sequence[Any], ) -> List[Any]: from meshmind.pipeline.preprocess import compress - return compress(memories) + return compress(list(memories)) def store_memories( self, - memories: List[Any], + memories: Iterable[Any], ) -> None: from meshmind.pipeline.store import store_memories - store_memories(memories, self.driver) + store_memories( + memories, + self._ensure_driver(), + entity_registry=self.entity_registry, + ) + + def store_triplets( + self, + triplets: Iterable[Triplet], + ) -> None: + from meshmind.pipeline.store import store_triplets + + store_triplets( + triplets, + self._ensure_driver(), + predicate_registry=self.predicate_registry, + ) + + # ------------------------------------------------------------------ + # CRUD helpers + # ------------------------------------------------------------------ + def create_memory(self, memory: Memory) -> UUID: + return self._ensure_manager().add_memory(memory) + + def update_memory(self, memory: Memory) -> None: + self._ensure_manager().update_memory(memory) + + def delete_memory(self, memory_id: UUID) -> None: + self._ensure_manager().delete_memory(memory_id) + + def get_memory(self, memory_id: UUID) -> Optional[Memory]: + return self._ensure_manager().get_memory(memory_id) + + def list_memories(self, namespace: str | None = None) -> List[Memory]: + return self._ensure_manager().list_memories(namespace) + + def create_triplet(self, triplet: Triplet) -> None: + self.predicate_registry.add(triplet.predicate) + self._ensure_manager().add_triplet(triplet) + + def delete_triplet(self, triplet: Triplet) -> None: + self._ensure_manager().delete_triplet( + triplet.subject, triplet.predicate, triplet.object + ) + + def list_triplets(self, namespace: str | None = None) -> List[Triplet]: + return self._ensure_manager().list_triplets(namespace) + + # ------------------------------------------------------------------ + # Retrieval helpers + # ------------------------------------------------------------------ + def search( + self, + query: str, + memories: Sequence[Memory], + namespace: str | None = None, + entity_labels: Sequence[str] | None = None, + config: SearchConfig | None = None, + use_llm_rerank: bool = False, + reranker: Callable[[str, Sequence[Memory], int], Sequence[Memory]] | None = None, + ) -> List[Memory]: + from meshmind.retrieval import llm_rerank, search as hybrid_search + + cfg = config or SearchConfig(encoder=self.embedding_model) + active_reranker = reranker + if use_llm_rerank: + active_reranker = lambda q, c, k: llm_rerank( + q, c, self.llm_client, k, model=cfg.rerank_model + ) + return hybrid_search( + query, + list(memories), + namespace=namespace, + entity_labels=list(entity_labels) if entity_labels else None, + config=cfg, + reranker=active_reranker, + ) + + def search_vector( + self, + query: str, + memories: Sequence[Memory], + namespace: str | None = None, + entity_labels: Sequence[str] | None = None, + config: SearchConfig | None = None, + ) -> List[Memory]: + from meshmind.retrieval import search_vector + + cfg = config or SearchConfig(encoder=self.embedding_model) + return search_vector( + query, + list(memories), + namespace=namespace, + entity_labels=list(entity_labels) if entity_labels else None, + config=cfg, + ) + + def search_regex( + self, + pattern: str, + memories: Sequence[Memory], + namespace: str | None = None, + entity_labels: Sequence[str] | None = None, + flags: int | None = None, + top_k: int = 10, + ) -> List[Memory]: + from meshmind.retrieval import search_regex + + return search_regex( + pattern, + list(memories), + namespace=namespace, + entity_labels=list(entity_labels) if entity_labels else None, + flags=flags, + top_k=top_k, + ) + + def search_exact( + self, + query: str, + memories: Sequence[Memory], + namespace: str | None = None, + entity_labels: Sequence[str] | None = None, + fields: Sequence[str] | None = None, + case_sensitive: bool = False, + top_k: int = 10, + ) -> List[Memory]: + from meshmind.retrieval import search_exact + + return search_exact( + query, + list(memories), + namespace=namespace, + entity_labels=list(entity_labels) if entity_labels else None, + fields=list(fields) if fields else None, + case_sensitive=case_sensitive, + top_k=top_k, + ) + diff --git a/meshmind/core/bootstrap.py b/meshmind/core/bootstrap.py new file mode 100644 index 0000000..0db5610 --- /dev/null +++ b/meshmind/core/bootstrap.py @@ -0,0 +1,45 @@ +"""Bootstrap helpers for wiring encoders and registries.""" +from __future__ import annotations + +import warnings +from typing import Iterable, Sequence, Type + +from pydantic import BaseModel + +from meshmind.core.config import settings +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder +from meshmind.core.types import Memory +from meshmind.models.registry import EntityRegistry, PredicateRegistry + + +def bootstrap_encoders(default_models: Sequence[str] | None = None) -> None: + """Ensure a default set of embedding encoders are registered.""" + + models = list(default_models) if default_models else [settings.EMBEDDING_MODEL] + for model_name in models: + if EncoderRegistry.is_registered(model_name): + continue + try: + EncoderRegistry.register(model_name, OpenAIEmbeddingEncoder(model_name)) + except ImportError as exc: + warnings.warn( + f"Skipping registration of OpenAI encoder '{model_name}': {exc}", + RuntimeWarning, + stacklevel=2, + ) + + +def bootstrap_entities(entity_models: Iterable[Type[BaseModel]] | None = None) -> None: + """Register default entity models used throughout the application.""" + + models = list(entity_models) if entity_models else [Memory] + for model in models: + EntityRegistry.register(model) + + +def register_predicates(predicates: Iterable[str]) -> None: + """Register predicate labels in the global predicate registry.""" + + for predicate in predicates: + PredicateRegistry.add(predicate) + diff --git a/meshmind/core/config.py b/meshmind/core/config.py index 9b14ffd..2d81bc3 100644 --- a/meshmind/core/config.py +++ b/meshmind/core/config.py @@ -20,6 +20,12 @@ class Settings: OPENAI_API_KEY: str = os.getenv("OPENAI_API_KEY", "") EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small") + REQUIRED_GROUPS = { + "graph": ("MEMGRAPH_URI",), + "openai": ("OPENAI_API_KEY",), + "redis": ("REDIS_URL",), + } + def __repr__(self) -> str: return ( f"Settings(MEMGRAPH_URI={self.MEMGRAPH_URI}, " @@ -28,5 +34,35 @@ def __repr__(self) -> str: f"EMBEDDING_MODEL={self.EMBEDDING_MODEL})" ) + @staticmethod + def _mask(value: str) -> str: + if not value: + return "" + if len(value) <= 4: + return "*" * len(value) + return f"{value[:2]}***{value[-2:]}" + + def missing(self) -> dict[str, list[str]]: + """Return missing environment variables grouped by capability.""" + + missing: dict[str, list[str]] = {} + for group, keys in self.REQUIRED_GROUPS.items(): + absent = [key for key in keys if not getattr(self, key)] + if absent: + missing[group] = absent + return missing + + def summary(self) -> dict[str, str]: + """Return a sanitized summary of active configuration values.""" + + return { + "MEMGRAPH_URI": self.MEMGRAPH_URI, + "MEMGRAPH_USERNAME": self.MEMGRAPH_USERNAME, + "MEMGRAPH_PASSWORD": self._mask(self.MEMGRAPH_PASSWORD), + "REDIS_URL": self.REDIS_URL, + "OPENAI_API_KEY": self._mask(self.OPENAI_API_KEY), + "EMBEDDING_MODEL": self.EMBEDDING_MODEL, + } + -settings = Settings() \ No newline at end of file +settings = Settings() diff --git a/meshmind/core/embeddings.py b/meshmind/core/embeddings.py index 84a67e8..ed2eb56 100644 --- a/meshmind/core/embeddings.py +++ b/meshmind/core/embeddings.py @@ -1,18 +1,16 @@ -""" -Embedding encoders and registry for MeshMind. -""" -from typing import List, Dict, Any +"""Embedding encoder implementations and registry utilities.""" +from __future__ import annotations + import time +from typing import Any, Dict, List _OPENAI_AVAILABLE = True -try: +try: # pragma: no cover - environment dependent from openai import OpenAI - from openai.error import RateLimitError -except ImportError: + from openai import RateLimitError +except ImportError: # pragma: no cover - environment dependent _OPENAI_AVAILABLE = False - openai = None # type: ignore - class RateLimitError(Exception): # type: ignore - pass + OpenAI = None # type: ignore from .config import settings @@ -31,12 +29,11 @@ def __init__( raise ImportError( "openai package is required for OpenAIEmbeddingEncoder" ) - try: - openai.api_key = settings.OPENAI_API_KEY - except Exception: - pass + client_kwargs: Dict[str, Any] = {} + if settings.OPENAI_API_KEY: + client_kwargs["api_key"] = settings.OPENAI_API_KEY - self.llm_client = OpenAI() + self.llm_client = OpenAI(**client_kwargs) self.RateLimitError = RateLimitError self.model_name = model_name self.max_retries = max_retries @@ -49,14 +46,23 @@ def encode(self, texts: List[str] | str) -> List[List[float]]: """ if isinstance(texts, str): texts = [texts] - + for attempt in range(self.max_retries): try: response = self.llm_client.embeddings.create( model=self.model_name, input=texts, ) - return [item['embedding'] for item in response['data']] + data = getattr(response, "data", None) + if data is None: + data = response.get("data", []) # type: ignore[assignment] + embeddings: List[List[float]] = [] + for item in data: + if hasattr(item, "embedding"): + embeddings.append(list(getattr(item, "embedding"))) + else: + embeddings.append(list(item["embedding"])) + return embeddings except self.RateLimitError: time.sleep(self.backoff_factor * (2 ** attempt)) except Exception: @@ -71,7 +77,13 @@ class SentenceTransformerEncoder: Encoder that uses a local SentenceTransformer model. """ def __init__(self, model_name: str): - from sentence_transformers import SentenceTransformer + try: # pragma: no cover - optional dependency + from sentence_transformers import SentenceTransformer + except ImportError as exc: + raise ImportError( + "sentence-transformers is required for SentenceTransformerEncoder." + " Install the optional 'sentence-transformers' extra to enable this encoder." + ) from exc self.model = SentenceTransformer(model_name) @@ -106,4 +118,22 @@ def get(cls, name: str) -> Any: encoder = cls._encoders.get(name) if encoder is None: raise KeyError(f"Encoder '{name}' not found in registry") - return encoder \ No newline at end of file + return encoder + + @classmethod + def is_registered(cls, name: str) -> bool: + """Return True if an encoder ``name`` has been registered.""" + + return name in cls._encoders + + @classmethod + def available(cls) -> List[str]: + """Return the list of registered encoder identifiers.""" + + return list(cls._encoders.keys()) + + @classmethod + def clear(cls) -> None: + """Remove all registered encoders. Intended for testing.""" + + cls._encoders.clear() diff --git a/meshmind/core/types.py b/meshmind/core/types.py index 112d31e..1ab71e9 100644 --- a/meshmind/core/types.py +++ b/meshmind/core/types.py @@ -43,5 +43,6 @@ class SearchConfig(BaseModel): encoder: str = "text-embedding-3-small" top_k: int = 20 rerank_k: int = 10 + rerank_model: Optional[str] = None filters: Optional[dict[str, Any]] = None hybrid_weights: Tuple[float, float] = (0.5, 0.5) \ No newline at end of file diff --git a/meshmind/core/utils.py b/meshmind/core/utils.py index 74e90dc..1884212 100644 --- a/meshmind/core/utils.py +++ b/meshmind/core/utils.py @@ -1,9 +1,44 @@ -"""Utility functions for MeshMind.""" +"""Utility helpers for MeshMind with optional dependency guards.""" +from __future__ import annotations + +import hashlib import uuid from datetime import datetime -import hashlib -from typing import Any -import tiktoken +from functools import lru_cache +from typing import Any, Optional + +_TIKTOKEN = None + + +def _ensure_tiktoken() -> Any: + """Return the ``tiktoken`` module if it is installed.""" + + global _TIKTOKEN + if _TIKTOKEN is None: + try: + import tiktoken # type: ignore + except ImportError as exc: # pragma: no cover - exercised in minimal envs + raise RuntimeError( + "tiktoken is required for token counting but is not installed." + " Install the optional 'tiktoken' extra to enable compression features." + ) from exc + _TIKTOKEN = tiktoken + return _TIKTOKEN + + +@lru_cache(maxsize=8) +def get_token_encoder(encoding_name: str = "o200k_base", optional: bool = False) -> Optional[Any]: + """Return a cached tiktoken encoder or ``None`` when optional.""" + + try: + module = _ensure_tiktoken() + except RuntimeError: + if optional: + return None + raise + return module.get_encoding(encoding_name) + + def generate_uuid() -> str: """Generate a UUID4 string.""" return str(uuid.uuid4()) @@ -21,13 +56,7 @@ def hash_dict(data: Any) -> str: return hash_string(str(data)) def num_tokens_from_string(string: str, encoding_name: str = "o200k_base") -> int: - """Returns the number of tokens in a text string. - Args: - string: The text string to count tokens for. - encoding_name: The name of the encoding to use. Defaults to "o200k_base". - Returns: - The number of tokens in the text string. - """ - encoding = tiktoken.get_encoding(encoding_name) - num_tokens = len(encoding.encode(string)) - return num_tokens \ No newline at end of file + """Return the number of tokens in ``string`` for ``encoding_name``.""" + + encoder = get_token_encoder(encoding_name, optional=False) + return len(encoder.encode(string)) diff --git a/meshmind/db/base_driver.py b/meshmind/db/base_driver.py index 1e1366c..4c6a9a6 100644 --- a/meshmind/db/base_driver.py +++ b/meshmind/db/base_driver.py @@ -1,6 +1,8 @@ """Abstract base class for graph database drivers.""" +from __future__ import annotations + from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional import uuid @@ -25,4 +27,16 @@ def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: @abstractmethod def delete(self, uuid: uuid.UUID) -> None: """Delete a node or relationship by UUID.""" - raise NotImplementedError \ No newline at end of file + raise NotImplementedError + + @abstractmethod + def delete_triplet(self, subj: str, pred: str, obj: str) -> None: + """Delete a relationship identified by subject/predicate/object.""" + + raise NotImplementedError + + @abstractmethod + def list_triplets(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + """Return stored triplets, optionally filtered by namespace.""" + + raise NotImplementedError diff --git a/meshmind/db/memgraph_driver.py b/meshmind/db/memgraph_driver.py index f4e8c4e..8d6b59f 100644 --- a/meshmind/db/memgraph_driver.py +++ b/meshmind/db/memgraph_driver.py @@ -1,110 +1,136 @@ -"""Memgraph implementation of GraphDriver.""" -from typing import Any, Dict, List -from .base_driver import GraphDriver +"""Memgraph implementation of :class:`GraphDriver` using ``mgclient``.""" +from __future__ import annotations - -"""Memgraph implementation of GraphDriver using mgclient.""" from typing import Any, Dict, List, Optional from urllib.parse import urlparse -try: +from meshmind.db.base_driver import GraphDriver + +try: # pragma: no cover - optional dependency import mgclient -except ImportError: +except ImportError: # pragma: no cover - optional dependency mgclient = None # type: ignore -from .base_driver import GraphDriver - class MemgraphDriver(GraphDriver): - """Memgraph driver implementation of GraphDriver using mgclient.""" + """Memgraph driver implementation backed by ``mgclient``.""" - def __init__(self, uri: str, username: str = None, password: str = None) -> None: - """Initialize Memgraph driver with Bolt URI and credentials.""" + def __init__(self, uri: str, username: str = "", password: str = "") -> None: if mgclient is None: raise ImportError("mgclient is required for MemgraphDriver") + self.uri = uri self.username = username self.password = password - # Parse URI: bolt://host:port + parsed = urlparse(uri) - host = parsed.hostname or 'localhost' + host = parsed.hostname or "localhost" port = parsed.port or 7687 - # Establish connection - self._conn = mgclient.connect( + + self._conn = mgclient.connect( # type: ignore[union-attr] host=host, port=port, - username=username, - password=password, + username=username or None, + password=password or None, ) self._cursor = self._conn.cursor() - def _execute(self, cypher: str, params: Optional[Dict[str, Any]] = None): - if params is None: - params = {} + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + def _execute(self, cypher: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + params = params or {} self._cursor.execute(cypher, params) try: rows = self._cursor.fetchall() cols = [col[0] for col in self._cursor.description] - results: List[Dict[str, Any]] = [] - for row in rows: - rec: Dict[str, Any] = {} - for idx, val in enumerate(row): - rec[cols[idx]] = val - results.append(rec) - return results except Exception: return [] + results: List[Dict[str, Any]] = [] + for row in rows: + record: Dict[str, Any] = {} + for idx, value in enumerate(row): + record[cols[idx]] = value + results.append(record) + return results + + @staticmethod + def _sanitize_predicate(predicate: str) -> str: + return predicate.replace("`", "") + + # ------------------------------------------------------------------ + # GraphDriver API + # ------------------------------------------------------------------ def upsert_entity(self, label: str, name: str, props: Dict[str, Any]) -> None: - """Insert or update an entity node by uuid.""" - uid = props.get('uuid') + uid = props.get("uuid") cypher = ( f"MERGE (n:{label} {{uuid: $uuid}})\n" f"SET n += $props" ) - params = {'uuid': str(uid), 'props': props} + params = {"uuid": str(uid), "props": props} self._execute(cypher, params) self._conn.commit() def upsert_edge(self, subj: str, pred: str, obj: str, props: Dict[str, Any]) -> None: - """Insert or update an edge between two entities identified by uuid.""" + predicate = self._sanitize_predicate(pred) cypher = ( - f"MATCH (a {{uuid: $subj}}), (b {{uuid: $obj}})\n" - f"MERGE (a)-[r:`{pred}`]->(b)\n" - f"SET r += $props" + "MATCH (a {uuid: $subj}), (b {uuid: $obj})\n" + f"MERGE (a)-[r:`{predicate}`]->(b)\n" + "SET r += $props" ) - params = {'subj': str(subj), 'obj': str(obj), 'props': props} + params = {"subj": str(subj), "obj": str(obj), "props": props} self._execute(cypher, params) self._conn.commit() def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: - """Execute a Cypher query and return results as list of dicts.""" return self._execute(cypher, params) def delete(self, uuid: Any) -> None: - """Delete a node (and detach relationships) by uuid.""" cypher = "MATCH (n {uuid: $uuid}) DETACH DELETE n" - params = {'uuid': str(uuid)} + params = {"uuid": str(uuid)} self._execute(cypher, params) self._conn.commit() + def delete_triplet(self, subj: str, pred: str, obj: str) -> None: + predicate = self._sanitize_predicate(pred) + cypher = ( + "MATCH (a {uuid: $subj})-[r:`{predicate}`]->(b {uuid: $obj}) " + "DELETE r" + ) + params = {"subj": str(subj), "obj": str(obj)} + self._execute(cypher, params) + self._conn.commit() + + def list_triplets(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + cypher = ( + "MATCH (a)-[r]->(b)\n" + "WHERE $namespace IS NULL OR r.namespace = $namespace\n" + "RETURN a.uuid AS subject, type(r) AS predicate, b.uuid AS object, " + "r.namespace AS namespace, r.metadata AS metadata, r.reference_time AS reference_time" + ) + params = {"namespace": namespace} + return self._execute(cypher, params) + + # ------------------------------------------------------------------ + # Convenience helpers + # ------------------------------------------------------------------ def vector_search(self, embedding: List[float], top_k: int = 10) -> List[Dict[str, Any]]: - """ - Fallback vector search: loads all embeddings and ranks by cosine similarity. - """ from meshmind.core.similarity import cosine_similarity - # Load all entities with embeddings - records = self.find("MATCH (n) WHERE exists(n.embedding) RETURN n.embedding AS emb, n AS node", {}) + + records = self.find( + "MATCH (n) WHERE exists(n.embedding) RETURN n.embedding AS emb, n AS node", + {}, + ) scored = [] for rec in records: - emb = rec.get('emb') + emb = rec.get("emb") if not isinstance(emb, list): continue try: score = cosine_similarity(embedding, emb) except Exception: score = 0.0 - scored.append({'node': rec.get('node'), 'score': float(score)}) - # Sort and take top_k - scored.sort(key=lambda x: x['score'], reverse=True) - return scored[:top_k] \ No newline at end of file + scored.append({"node": rec.get("node"), "score": float(score)}) + scored.sort(key=lambda item: item["score"], reverse=True) + return scored[:top_k] diff --git a/meshmind/models/registry.py b/meshmind/models/registry.py index bd98f87..498a25c 100644 --- a/meshmind/models/registry.py +++ b/meshmind/models/registry.py @@ -30,4 +30,16 @@ def add(cls, label: str) -> None: @classmethod def allowed(cls, label: str) -> bool: """Check if a predicate label is allowed.""" - return label in cls._predicates \ No newline at end of file + return label in cls._predicates + + @classmethod + def all(cls) -> Set[str]: + """Return all registered predicate labels.""" + + return set(cls._predicates) + + @classmethod + def clear(cls) -> None: + """Remove all registered predicates (testing helper).""" + + cls._predicates.clear() diff --git a/meshmind/pipeline/compress.py b/meshmind/pipeline/compress.py index 9ded746..b1f65ac 100644 --- a/meshmind/pipeline/compress.py +++ b/meshmind/pipeline/compress.py @@ -1,13 +1,10 @@ -""" -Pipeline for token-aware compression/summarization of memories. -""" +"""Token-aware compression helpers for memory metadata.""" +from __future__ import annotations + from typing import List -from meshmind.core.types import Memory -try: - import tiktoken -except ImportError: - tiktoken = None # type: ignore +from meshmind.core.types import Memory +from meshmind.core.utils import get_token_encoder def compress_memories( @@ -20,7 +17,9 @@ def compress_memories( :param max_tokens: Maximum number of tokens allowed per memory. :return: List of Memory objects with content possibly shortened. """ - encoder = tiktoken.get_encoding('o200k_base') + encoder = get_token_encoder("o200k_base", optional=True) + if encoder is None: + return memories compressed = [] for mem in memories: content = mem.metadata.get('content') @@ -35,4 +34,4 @@ def compress_memories( truncated = encoder.decode(tokens[:max_tokens]) mem.metadata['content'] = truncated compressed.append(mem) - return compressed \ No newline at end of file + return compressed diff --git a/meshmind/pipeline/extract.py b/meshmind/pipeline/extract.py index 613073b..4897527 100644 --- a/meshmind/pipeline/extract.py +++ b/meshmind/pipeline/extract.py @@ -1,11 +1,11 @@ -from typing import Any, List, Type +from typing import Any, List, Sequence, Type def extract_memories( instructions: str, namespace: str, - entity_types: List[Type[Any]], + entity_types: Sequence[Type[Any]], embedding_model: str, - content: List[str], + content: Sequence[str], llm_client: Any = None, ) -> List[Any]: """ @@ -26,6 +26,7 @@ def extract_memories( raise RuntimeError("openai package is required for extraction pipeline") from meshmind.core.types import Memory from meshmind.core.embeddings import EncoderRegistry + from meshmind.models.registry import EntityRegistry # Initialize default LLM client if not provided if llm_client is None: @@ -46,6 +47,9 @@ def extract_memories( } # Build system prompt using a default template and user instructions + entity_types = list(entity_types) or [Memory] + for model in entity_types: + EntityRegistry.register(model) allowed_labels = [cls.__name__ for cls in entity_types] default_prompt = ( "You are an agent that extracts structured memories from text segments. " @@ -58,7 +62,7 @@ def extract_memories( prompt += f"\nAllowed entity labels: {', '.join(allowed_labels)}." messages = [{"role": "system", "content": prompt}] # Add each text segment as a user message - messages += [{"role": "user", "content": text} for text in content] + messages += [{"role": "user", "content": text} for text in list(content)] # Call chat completion with function-calling response = llm_client.responses.create( @@ -99,4 +103,4 @@ def extract_memories( emb = encoder.encode([mem.name])[0] mem.embedding = emb memories.append(mem) - return memories \ No newline at end of file + return memories diff --git a/meshmind/pipeline/store.py b/meshmind/pipeline/store.py index 3d0b07c..b8199bb 100644 --- a/meshmind/pipeline/store.py +++ b/meshmind/pipeline/store.py @@ -1,23 +1,60 @@ +"""Persistence helpers for storing memories and triplets.""" +from __future__ import annotations + from typing import Any, Iterable + +from pydantic import BaseModel + +from meshmind.core.types import Triplet from meshmind.db.base_driver import GraphDriver +from meshmind.models.registry import EntityRegistry, PredicateRegistry + + +def _props(obj: Any) -> dict[str, Any]: + if isinstance(obj, BaseModel): + return obj.dict(exclude_none=True) + if hasattr(obj, "dict"): + try: + return obj.dict(exclude_none=True) # type: ignore[attr-defined] + except TypeError: + pass + if isinstance(obj, dict): + return {k: v for k, v in obj.items() if v is not None} + return {k: v for k, v in obj.__dict__.items() if v is not None} + def store_memories( memories: Iterable[Any], graph_driver: GraphDriver, + *, + entity_registry: type[EntityRegistry] | None = None, ) -> None: - """ - Persist a sequence of Memory objects into the graph database. + """Persist a sequence of Memory objects into the graph database.""" - :param memories: An iterable of Memory-like objects with attributes for upsert. - :param graph_driver: An instance of GraphDriver to perform database operations. - """ - # Iterate over Memory-like objects and upsert into graph + registry = entity_registry or EntityRegistry for mem in memories: - # Use Pydantic-like dict to extract properties - try: - props = mem.dict(exclude_none=True) - except Exception: - # Fallback for non-Pydantic objects - props = mem.__dict__ - # Upsert entity node with label and name - graph_driver.upsert_entity(mem.entity_label, mem.name, props) \ No newline at end of file + props = _props(mem) + label = getattr(mem, "entity_label", None) + if label and registry.model_for_label(label) is None and isinstance(mem, BaseModel): + registry.register(type(mem)) + graph_driver.upsert_entity(label or "Memory", getattr(mem, "name", ""), props) + + +def store_triplets( + triplets: Iterable[Triplet], + graph_driver: GraphDriver, + *, + predicate_registry: type[PredicateRegistry] | None = None, +) -> None: + """Persist a collection of ``Triplet`` relationships.""" + + registry = predicate_registry or PredicateRegistry + for triplet in triplets: + registry.add(triplet.predicate) + props = _props(triplet) + graph_driver.upsert_edge( + triplet.subject, + triplet.predicate, + triplet.object, + props, + ) diff --git a/meshmind/retrieval/__init__.py b/meshmind/retrieval/__init__.py index e69de29..7864fbb 100644 --- a/meshmind/retrieval/__init__.py +++ b/meshmind/retrieval/__init__.py @@ -0,0 +1,25 @@ +"""Retrieval helpers exposed for external consumers.""" + +from .search import ( + search, + search_bm25, + search_exact, + search_fuzzy, + search_regex, + search_vector, +) +from .vector import vector_search, vector_search_from_embeddings +from .rerank import llm_rerank, apply_reranker + +__all__ = [ + "search", + "search_bm25", + "search_exact", + "search_fuzzy", + "search_regex", + "search_vector", + "vector_search", + "vector_search_from_embeddings", + "llm_rerank", + "apply_reranker", +] diff --git a/meshmind/retrieval/rerank.py b/meshmind/retrieval/rerank.py new file mode 100644 index 0000000..71085ec --- /dev/null +++ b/meshmind/retrieval/rerank.py @@ -0,0 +1,80 @@ +"""Helpers for reranking retrieval results.""" +from __future__ import annotations + +from typing import Callable, List, Sequence + +from meshmind.core.types import Memory + +Reranker = Callable[[str, Sequence[Memory], int], Sequence[Memory]] + + +def llm_rerank( + query: str, + memories: Sequence[Memory], + llm_client: object | None, + top_k: int, + model: str | None = None, +) -> List[Memory]: + """Rerank results using an LLM client that supports the Responses API.""" + if llm_client is None or not memories: + return list(memories)[:top_k] + + model_name = model or "gpt-4o-mini" + prompt = "\n".join( + [ + "You are a ranking assistant.", + "Given the query and numbered memory summaries, return a JSON array of memory indexes", + "sorted from best to worst match.", + f"Query: {query}", + "Memories:", + ] + ) + for idx, memory in enumerate(memories): + prompt += f"\n{idx}: {memory.name}" + + try: # pragma: no cover - network interaction mocked in tests + response = llm_client.responses.create( # type: ignore[attr-defined] + model=model_name, + input=[{"role": "user", "content": prompt}], + response_format={"type": "json_schema", "json_schema": { + "name": "rankings", + "schema": { + "type": "object", + "properties": { + "order": { + "type": "array", + "items": {"type": "integer"}, + } + }, + "required": ["order"], + }, + }}, + ) + content = response.output[0].content[0].text # type: ignore[index] + except Exception: + return list(memories)[:top_k] + + try: + import json + + data = json.loads(content) + indexes = [idx for idx in data.get("order", []) if 0 <= idx < len(memories)] + except Exception: + return list(memories)[:top_k] + + ranked = [memories[idx] for idx in indexes] + remaining = [mem for mem in memories if mem not in ranked] + ranked.extend(remaining) + return ranked[:top_k] + + +def apply_reranker( + query: str, + candidates: Sequence[Memory], + top_k: int, + reranker: Reranker | None = None, +) -> List[Memory]: + if reranker is None: + return list(candidates)[:top_k] + ranked = reranker(query, candidates, top_k) + return list(ranked)[:top_k] diff --git a/meshmind/retrieval/search.py b/meshmind/retrieval/search.py index 666d7b1..19724ce 100644 --- a/meshmind/retrieval/search.py +++ b/meshmind/retrieval/search.py @@ -1,7 +1,8 @@ -""" -Unified dispatcher for various retrieval strategies. -""" -from typing import List, Optional +"""Unified dispatcher for various retrieval strategies.""" +from __future__ import annotations + +import re +from typing import Callable, List, Optional, Sequence from meshmind.core.types import Memory, SearchConfig from meshmind.retrieval.bm25 import bm25_search @@ -12,6 +13,23 @@ filter_by_entity_labels, filter_by_metadata, ) +from meshmind.retrieval.rerank import apply_reranker +from meshmind.retrieval.vector import vector_search + +Reranker = Callable[[str, Sequence[Memory], int], Sequence[Memory]] + + +def _apply_filters( + memories: Sequence[Memory], + namespace: Optional[str], + entity_labels: Optional[List[str]], + config: Optional[SearchConfig], +) -> List[Memory]: + mems = filter_by_namespace(list(memories), namespace) + mems = filter_by_entity_labels(mems, entity_labels) + if config and config.filters: + mems = filter_by_metadata(mems, config.filters) + return mems def search( @@ -20,28 +38,16 @@ def search( namespace: Optional[str] = None, entity_labels: Optional[List[str]] = None, config: Optional[SearchConfig] = None, + reranker: Reranker | None = None, ) -> List[Memory]: - """ - Perform hybrid search over memories with optional filters. - - :param query: Query string. - :param memories: List of Memory objects. - :param namespace: Filter by namespace. - :param entity_labels: Filter by entity labels. - :param config: SearchConfig overriding defaults. - :return: Ranked list of Memory objects. - """ - # Apply filters - mems = filter_by_namespace(memories, namespace) - mems = filter_by_entity_labels(mems, entity_labels) - if config and config.filters: - mems = filter_by_metadata(mems, config.filters) - - # Use hybrid search by default + """Perform hybrid search with optional reranking.""" cfg = config or SearchConfig() + mems = _apply_filters(memories, namespace, entity_labels, cfg) ranked = hybrid_search(query, mems, cfg) - # Return only Memory objects - return [m for m, _ in ranked] + ordered = [m for m, _ in ranked] + if not ordered: + return [] + return apply_reranker(query, ordered, cfg.rerank_k, reranker) def search_bm25( @@ -51,8 +57,7 @@ def search_bm25( entity_labels: Optional[List[str]] = None, top_k: int = 10, ) -> List[Memory]: - mems = filter_by_namespace(memories, namespace) - mems = filter_by_entity_labels(mems, entity_labels) + mems = _apply_filters(memories, namespace, entity_labels, None) results = bm25_search(query, mems, top_k=top_k) return [m for m, _ in results] @@ -64,7 +69,73 @@ def search_fuzzy( entity_labels: Optional[List[str]] = None, top_k: int = 10, ) -> List[Memory]: - mems = filter_by_namespace(memories, namespace) - mems = filter_by_entity_labels(mems, entity_labels) + mems = _apply_filters(memories, namespace, entity_labels, None) results = fuzzy_search(query, mems, top_k=top_k) - return [m for m, _ in results] \ No newline at end of file + return [m for m, _ in results] + + +def search_vector( + query: str, + memories: List[Memory], + namespace: Optional[str] = None, + entity_labels: Optional[List[str]] = None, + config: Optional[SearchConfig] = None, +) -> List[Memory]: + cfg = config or SearchConfig() + mems = _apply_filters(memories, namespace, entity_labels, cfg) + results = vector_search(query, mems, cfg) + return [m for m, _ in results] + + +def search_regex( + pattern: str, + memories: List[Memory], + namespace: Optional[str] = None, + entity_labels: Optional[List[str]] = None, + flags: int | None = None, + top_k: int = 10, +) -> List[Memory]: + mems = _apply_filters(memories, namespace, entity_labels, None) + regex = re.compile(pattern, flags or re.IGNORECASE) + scored: List[tuple[Memory, int]] = [] + for mem in mems: + haystacks = [mem.name] + [str(value) for value in mem.metadata.values()] + matches = [len(regex.findall(h)) for h in haystacks] + score = max(matches, default=0) + if score > 0: + scored.append((mem, score)) + scored.sort(key=lambda item: item[1], reverse=True) + return [mem for mem, _ in scored[:top_k]] + + +def search_exact( + query: str, + memories: List[Memory], + namespace: Optional[str] = None, + entity_labels: Optional[List[str]] = None, + fields: Optional[List[str]] = None, + case_sensitive: bool = False, + top_k: int = 10, +) -> List[Memory]: + mems = _apply_filters(memories, namespace, entity_labels, None) + needle = query if case_sensitive else query.lower() + fields = fields or ["name"] + + def normalize(value: object) -> str: + text = "" if value is None else str(value) + return text if case_sensitive else text.lower() + + matched: List[Memory] = [] + for mem in mems: + for field in fields: + value = getattr(mem, field, None) + if value is None and field == "metadata": + for meta_val in mem.metadata.values(): + if normalize(meta_val) == needle: + matched.append(mem) + break + continue + if normalize(value) == needle: + matched.append(mem) + break + return matched[:top_k] diff --git a/meshmind/retrieval/vector.py b/meshmind/retrieval/vector.py new file mode 100644 index 0000000..2ca83a6 --- /dev/null +++ b/meshmind/retrieval/vector.py @@ -0,0 +1,57 @@ +"""Vector-only retrieval helpers.""" +from __future__ import annotations + +from typing import Iterable, List, Sequence, Tuple + +from meshmind.core.embeddings import EncoderRegistry +from meshmind.core.similarity import cosine_similarity +from meshmind.core.types import Memory, SearchConfig + + +def vector_search( + query: str, + memories: Sequence[Memory], + config: SearchConfig | None = None, +) -> List[Tuple[Memory, float]]: + """Rank memories using cosine similarity against the query embedding.""" + if not memories: + return [] + + cfg = config or SearchConfig() + encoder = EncoderRegistry.get(cfg.encoder) + query_embedding = encoder.encode([query])[0] + + scored: List[Tuple[Memory, float]] = [] + for memory in memories: + embedding = getattr(memory, "embedding", None) + if embedding is None: + continue + try: + score = cosine_similarity(query_embedding, embedding) + except Exception: + score = 0.0 + scored.append((memory, float(score))) + + scored.sort(key=lambda item: item[1], reverse=True) + return scored[: cfg.top_k] + + +def vector_search_from_embeddings( + query_embedding: Sequence[float], + memories: Iterable[Memory], + top_k: int = 10, +) -> List[Tuple[Memory, float]]: + """Rank memories when the query embedding is precomputed.""" + scored: List[Tuple[Memory, float]] = [] + for memory in memories: + embedding = getattr(memory, "embedding", None) + if embedding is None: + continue + try: + score = cosine_similarity(query_embedding, embedding) + except Exception: + score = 0.0 + scored.append((memory, float(score))) + + scored.sort(key=lambda item: item[1], reverse=True) + return scored[:top_k] diff --git a/meshmind/tasks/scheduled.py b/meshmind/tasks/scheduled.py index eaa4ce2..59428d8 100644 --- a/meshmind/tasks/scheduled.py +++ b/meshmind/tasks/scheduled.py @@ -1,6 +1,8 @@ """ Scheduled Celery tasks for expiry, consolidation, and compression. """ +from __future__ import annotations + try: from celery.schedules import crontab _CELERY_BEAT = True @@ -17,17 +19,25 @@ def crontab(*args, **kwargs): # type: ignore from meshmind.db.memgraph_driver import MemgraphDriver from meshmind.core.config import settings -# Initialize database driver and memory manager (fallback if mgclient missing) -try: - driver = MemgraphDriver( - settings.MEMGRAPH_URI, - settings.MEMGRAPH_USERNAME, - settings.MEMGRAPH_PASSWORD, - ) - manager = MemoryManager(driver) -except Exception: - driver = None # type: ignore - manager = None # type: ignore +_MANAGER: MemoryManager | None = None + + +def _get_manager() -> MemoryManager | None: + global _MANAGER + if _MANAGER is not None: + return _MANAGER + + try: + driver = MemgraphDriver( + settings.MEMGRAPH_URI, + settings.MEMGRAPH_USERNAME, + settings.MEMGRAPH_PASSWORD, + ) + except Exception: + return None + + _MANAGER = MemoryManager(driver) + return _MANAGER # Define periodic task schedule if Celery is available if _CELERY_BEAT and hasattr(app, 'conf'): @@ -50,6 +60,7 @@ def crontab(*args, **kwargs): # type: ignore @app.task(name='meshmind.tasks.scheduled.expire_task') def expire_task(): """Delete expired memories based on TTL.""" + manager = _get_manager() if manager is None: return [] return expire_memories(manager) @@ -58,6 +69,7 @@ def expire_task(): @app.task(name='meshmind.tasks.scheduled.consolidate_task') def consolidate_task(): """Merge duplicate memories and summarise.""" + manager = _get_manager() if manager is None: return 0 memories = manager.list_memories() @@ -70,10 +82,11 @@ def consolidate_task(): @app.task(name='meshmind.tasks.scheduled.compress_task') def compress_task(): """Compress long memories to respect token limits.""" + manager = _get_manager() if manager is None: return 0 memories = manager.list_memories() compressed = compress_memories(memories) for mem in compressed: manager.update_memory(mem) - return len(compressed) \ No newline at end of file + return len(compressed) diff --git a/meshmind/tests/conftest.py b/meshmind/tests/conftest.py new file mode 100644 index 0000000..fa92a05 --- /dev/null +++ b/meshmind/tests/conftest.py @@ -0,0 +1,28 @@ +import pytest + +from meshmind.core.embeddings import EncoderRegistry +from meshmind.core.types import Memory + + +@pytest.fixture +def memory_factory(): + def _factory(name: str, **overrides): + payload = {"namespace": "ns", "name": name, "entity_label": "Test"} + payload.update(overrides) + return Memory(**payload) + + return _factory + + +@pytest.fixture +def dummy_encoder(): + EncoderRegistry.clear() + + class DummyEncoder: + def encode(self, texts): + return [[1.0 if "apple" in text else 0.0] for text in texts] + + name = "dummy-encoder" + EncoderRegistry.register(name, DummyEncoder()) + yield name + EncoderRegistry.clear() diff --git a/meshmind/tests/test_memgraph_driver.py b/meshmind/tests/test_memgraph_driver.py index 03f7e79..66de8cf 100644 --- a/meshmind/tests/test_memgraph_driver.py +++ b/meshmind/tests/test_memgraph_driver.py @@ -55,8 +55,23 @@ def commit(self): # Test upsert_edge does not raise edge_props = {'rel': 'value'} driver.upsert_edge('id1', 'REL', 'id2', edge_props) + # Test delete_triplet uses predicate sanitisation + driver.delete_triplet('id1', 'REL', 'id2') + assert 'DELETE r' in driver._cursor._last_query + # Test list_triplets returns parsed dicts + driver._cursor.description = [ + ('subject',), + ('predicate',), + ('object',), + ('namespace',), + ('metadata',), + ('reference_time',), + ] + driver._cursor._rows = [(('s',), ('p',), ('o',), ('ns',), ({'k': 'v'},), (None,))] + triplets = driver.list_triplets() + assert triplets and triplets[0]['subject'] == ('s',) # Test vector_search returns list # Use dummy record driver._cursor._rows = [([1.0], {'uuid': 'id1'})] out = driver.vector_search([1.0], top_k=1) - assert isinstance(out, list) \ No newline at end of file + assert isinstance(out, list) diff --git a/meshmind/tests/test_pipeline_extract.py b/meshmind/tests/test_pipeline_extract.py index 6b72966..30ac3ff 100644 --- a/meshmind/tests/test_pipeline_extract.py +++ b/meshmind/tests/test_pipeline_extract.py @@ -1,11 +1,10 @@ import json + import pytest -import openai from meshmind.client import MeshMind from meshmind.core.types import Memory from meshmind.core.embeddings import EncoderRegistry -from meshmind.db.memgraph_driver import MemgraphDriver class DummyEncoder: @@ -14,39 +13,26 @@ def encode(self, texts): return [[len(text)] for text in texts] -class DummyChoice: - def __init__(self, message): - self.message = message - class DummyResponse: - def __init__(self, arg_json): - func_call = {'arguments': arg_json} - self.choices = [DummyChoice({'function_call': func_call})] + def __init__(self, payload): + self.choices = [type('Choice', (), {'message': payload})] -@pytest.fixture(autouse=True) -def patch_openai(monkeypatch): - """Patch OpenAI client to use DummyChat for responses.""" - class DummyChat: +class DummyLLMClient: + class responses: # type: ignore[assignment] @staticmethod def create(model, messages, functions, function_call): names = [m['content'] for m in messages if m['role'] == 'user'] items = [{'name': n, 'entity_label': 'Memory'} for n in names] arg_json = json.dumps({'memories': items}) - return DummyResponse(arg_json) - class DummyModelClient: - def __init__(self): - self.responses = DummyChat - monkeypatch.setattr(openai, 'OpenAI', lambda *args, **kwargs: DummyModelClient()) - return None + return DummyResponse({'function_call': {'arguments': arg_json}}) def test_extract_memories_basic(tmp_path): # Register dummy encoder + EncoderRegistry.clear() EncoderRegistry.register('text-embedding-3-small', DummyEncoder()) - mm = MeshMind() - # override default llm_client to use dummy - mm.llm_client = openai.OpenAI() + mm = MeshMind(llm_client=DummyLLMClient()) # Run extraction texts = ['alpha', 'beta'] results = mm.extract_memories( @@ -64,16 +50,17 @@ def test_extract_memories_basic(tmp_path): assert mem.embedding == [len(text)] -def test_extract_invalid_label(monkeypatch): - # Monkeypatch to return an entry with invalid label - def bad_create(*args, **kwargs): - arg_json = json.dumps({'memories': [{'name': 'x', 'entity_label': 'Bad'}]}) - return DummyResponse(arg_json) - from openai import OpenAI - llm_client = OpenAI() - monkeypatch.setattr(llm_client.responses, 'create', bad_create) +def test_extract_invalid_label(): + class BadLLMClient: + class responses: # type: ignore[assignment] + @staticmethod + def create(*args, **kwargs): + arg_json = json.dumps({'memories': [{'name': 'x', 'entity_label': 'Bad'}]}) + return DummyResponse({'function_call': {'arguments': arg_json}}) + + EncoderRegistry.clear() EncoderRegistry.register('text-embedding-3-small', DummyEncoder()) - mm = MeshMind(llm_client=llm_client) + mm = MeshMind(llm_client=BadLLMClient()) with pytest.raises(ValueError) as e: mm.extract_memories( instructions='Extract:', @@ -81,4 +68,4 @@ def bad_create(*args, **kwargs): entity_types=[Memory], content=['x'], ) - assert 'Invalid entity_label' in str(e.value) \ No newline at end of file + assert 'Invalid entity_label' in str(e.value) diff --git a/meshmind/tests/test_pipeline_preprocess_store.py b/meshmind/tests/test_pipeline_preprocess_store.py index 2d10abd..b02b414 100644 --- a/meshmind/tests/test_pipeline_preprocess_store.py +++ b/meshmind/tests/test_pipeline_preprocess_store.py @@ -1,26 +1,38 @@ import pytest from meshmind.pipeline.preprocess import deduplicate, score_importance, compress -from meshmind.pipeline.store import store_memories +from meshmind.pipeline.store import store_memories, store_triplets from meshmind.api.memory_manager import MemoryManager -from meshmind.core.types import Memory +from meshmind.core.types import Memory, Triplet +from meshmind.models.registry import PredicateRegistry class DummyDriver: def __init__(self): self.entities = [] self.deleted = [] + self.edges = [] + self.deleted_edges = [] def upsert_entity(self, label, name, props): self.entities.append((label, name, props)) + def upsert_edge(self, subj, pred, obj, props): + self.edges.append((subj, pred, obj, props)) + def delete(self, uuid): self.deleted.append(uuid) + def delete_triplet(self, subj, pred, obj): + self.deleted_edges.append((subj, pred, obj)) + def find(self, cypher, params): # Return empty for simplicity return [] + def list_triplets(self, namespace=None): + return [] + def make_memory(name: str) -> Memory: return Memory(namespace="ns", name=name, entity_label="Test") @@ -57,6 +69,21 @@ def test_store_memories_calls_driver(): assert d.entities[0][1] == "node1" +def test_store_triplets_registers_predicate(): + PredicateRegistry.clear() + d = DummyDriver() + triplet = Triplet( + subject="s", + predicate="RELATES", + object="o", + namespace="ns", + entity_label="Relation", + ) + store_triplets([triplet], d) + assert d.edges and d.edges[0][1] == "RELATES" + assert "RELATES" in PredicateRegistry.all() + + def test_memory_manager_add_update_delete(): d = DummyDriver() mgr = MemoryManager(d) @@ -76,6 +103,23 @@ def test_memory_manager_add_update_delete(): # list returns empty or list lst = mgr.list_memories() assert isinstance(lst, list) + + +def test_memory_manager_triplet_roundtrip(): + d = DummyDriver() + mgr = MemoryManager(d) + triplet = Triplet( + subject="s", + predicate="RELATES", + object="o", + namespace="ns", + entity_label="Relation", + ) + mgr.add_triplet(triplet) + assert d.edges + mgr.delete_triplet(triplet.subject, triplet.predicate, triplet.object) + assert d.deleted_edges + assert mgr.list_triplets() == [] def test_deduplicate_by_embedding_similarity(): # Two memories with similar embeddings should be deduplicated @@ -89,4 +133,4 @@ def test_deduplicate_by_embedding_similarity(): assert len(result_high) == 1 # With low threshold, keep both result_low = deduplicate([m1, m2], threshold=0.1) - assert len(result_low) == 2 \ No newline at end of file + assert len(result_low) == 2 diff --git a/meshmind/tests/test_retrieval.py b/meshmind/tests/test_retrieval.py index 66abd33..f4b90d7 100644 --- a/meshmind/tests/test_retrieval.py +++ b/meshmind/tests/test_retrieval.py @@ -1,76 +1,113 @@ import pytest -from meshmind.core.types import Memory, SearchConfig -from meshmind.retrieval.bm25 import bm25_search -from meshmind.retrieval.fuzzy import fuzzy_search +from meshmind.core.types import SearchConfig +from meshmind.retrieval import ( + apply_reranker, + llm_rerank, + search, + search_bm25, + search_exact, + search_fuzzy, + search_regex, + search_vector, +) from meshmind.retrieval.hybrid import hybrid_search -from meshmind.retrieval.search import search, search_bm25, search_fuzzy - - -def make_memory(name: str) -> Memory: - return Memory(namespace="ns", name=name, entity_label="Test") - - -@pytest.fixture(autouse=True) -def add_embeddings(): - # Assign dummy embeddings equal to length of name - def _hook(mem: Memory): - mem.embedding = [len(mem.name)] - return mem - Memory.pre_init = _hook - yield - delattr(Memory, 'pre_init') - - -def test_bm25_search(): - docs = [make_memory("apple pie"), make_memory("banana split"), make_memory("cherry tart")] - results = bm25_search("apple", docs, top_k=2) - # Expect 'apple pie' first - assert results and results[0][0].name == "apple pie" - assert results[0][1] > 0 - - -def test_fuzzy_search(): - docs = [make_memory("apple pie"), make_memory("banana split")] - results = fuzzy_search("apple pie", docs, top_k=2) - assert results and results[0][0].name == "apple pie" - assert 0 < results[0][1] <= 1.0 - - -def test_hybrid_search(): - # Setup memories - m1 = make_memory("apple") - m2 = make_memory("banana") - m1.embedding = [1.0] - m2.embedding = [0.0] - docs = [m1, m2] - config = SearchConfig(encoder="dummy", top_k=2, hybrid_weights=(0.5, 0.5)) - # Register dummy encoder that returns [1] for 'apple' and [0] for 'banana' - class DummyEncoder: - def encode(self, texts): - return [[1.0] if "apple" in t else [0.0] for t in texts] - from meshmind.core.embeddings import EncoderRegistry - EncoderRegistry.register("dummy", DummyEncoder()) - results = hybrid_search("apple", docs, config) - # apple should have highest hybrid score - assert results[0][0].name == "apple" - - -def test_search_dispatcher(): - m1 = make_memory("apple") - m2 = make_memory("banana") - m1.embedding = [1.0] - m2.embedding = [0.0] - docs = [m1, m2] - from meshmind.core.embeddings import EncoderRegistry - class DummyEncoder: - def encode(self, texts): return [[1.0] if "apple" in t else [0.0] for t in texts] - EncoderRegistry.register("dummy", DummyEncoder()) - config = SearchConfig(encoder="dummy", top_k=1, hybrid_weights=(0.5,0.5)) - res = search("apple", docs, namespace="ns", entity_labels=["Test"], config=config) - assert len(res) == 1 and res[0].name == "apple" - # BM25 and fuzzy via dispatcher - res2 = search_bm25("banana", docs) - assert res2 and res2[0].name == "banana" - res3 = search_fuzzy("banana", docs) - assert res3 and res3[0].name == "banana" \ No newline at end of file + + +def test_bm25_search(memory_factory): + docs = [ + memory_factory("apple pie"), + memory_factory("banana split"), + memory_factory("cherry tart"), + ] + results = search_bm25("apple", docs, top_k=2) + assert results and results[0].name == "apple pie" + + +def test_fuzzy_search(memory_factory): + docs = [memory_factory("apple pie"), memory_factory("banana split")] + results = search_fuzzy("apple pie", docs, top_k=2) + assert results and results[0].name == "apple pie" + + +def test_hybrid_search(memory_factory, dummy_encoder): + m1 = memory_factory("apple", embedding=[1.0]) + m2 = memory_factory("banana", embedding=[0.0]) + config = SearchConfig(encoder=dummy_encoder, top_k=2, hybrid_weights=(0.5, 0.5)) + ranked = hybrid_search("apple", [m1, m2], config) + assert ranked[0][0].name == "apple" + + +def test_vector_search(memory_factory, dummy_encoder): + m1 = memory_factory("apple", embedding=[1.0]) + m2 = memory_factory("banana", embedding=[0.0]) + config = SearchConfig(encoder=dummy_encoder, top_k=1) + results = search_vector("apple", [m1, m2], config=config) + assert results == [m1] + + +def test_regex_search(memory_factory): + docs = [ + memory_factory("Visit Paris", metadata={"city": "Paris"}), + memory_factory("Visit Berlin", metadata={"city": "Berlin"}), + ] + results = search_regex("paris", docs, top_k=5) + assert len(results) == 1 and results[0].name == "Visit Paris" + + +def test_exact_search(memory_factory): + docs = [ + memory_factory("Python"), + memory_factory("Rust", metadata={"language": "Rust"}), + ] + results = search_exact("rust", docs, fields=["metadata"], case_sensitive=False) + assert results and results[0].name == "Rust" + + +def test_search_dispatcher_with_rerank(memory_factory, dummy_encoder): + m1 = memory_factory("apple", embedding=[1.0]) + m2 = memory_factory("banana", embedding=[0.1]) + docs = [m2, m1] + config = SearchConfig(encoder=dummy_encoder, top_k=2, rerank_k=2) + + class DummyLLM: + class responses: + @staticmethod + def create(**kwargs): + return type( + "Resp", + (), + { + "output": [ + type( + "Out", + (), + { + "content": [ + type("Text", (), {"text": '{"order": [1, 0]}'}) + ] + }, + ) + ] + }, + ) + + reranked = search( + "apple", + docs, + config=config, + reranker=lambda q, c, k: llm_rerank(q, c, DummyLLM(), k, model="dummy"), + ) + assert reranked[0].name == "apple" + + +def test_apply_reranker_default(memory_factory): + docs = [memory_factory("alpha"), memory_factory("beta")] + ranked = apply_reranker("alpha", docs, top_k=1) + assert ranked == [docs[0]] + + +def test_llm_rerank_failure(memory_factory): + docs = [memory_factory("alpha"), memory_factory("beta")] + result = llm_rerank("alpha", docs, llm_client=None, top_k=2) + assert len(result) == 2