diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..e842439 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +# Agent Instructions + +- 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/DISCREPANCIES.md b/DISCREPANCIES.md new file mode 100644 index 0000000..5e391f9 --- /dev/null +++ b/DISCREPANCIES.md @@ -0,0 +1,43 @@ +# README vs Implementation Discrepancies + +## Overview +- The legacy README promises a fully featured memory graph service with multi-level APIs, relationship storage, and diverse retrieval methods. +- The current codebase delivers a narrower pipeline that focuses on extracting `Memory` nodes, preprocessing them, and writing them to Memgraph. +- Many examples in the README are not executable because the described methods, configuration defaults, and dependency behaviors do not exist. + +## API Surface +- README shows `MeshMind.register_entity`, `register_allowed_predicates`, `add_predicate`, `store_memory`, `add_memory`, `add_triplet`, `search`, `search_facts`, `search_procedures`, `update_memory`, and `delete_memory` methods. Only `extract_memories`, `deduplicate`, `score_importance`, `compress`, and `store_memories` exist on `meshmind.client.MeshMind`. +- Entity registration is depicted as a first-class feature that validates custom Pydantic models (e.g., `Person`). The implementation merely checks that the `entity_label` returned by the LLM matches the class name; there is no registry enforcing schemas or instantiating the models. +- Triplet storage is central to the README narrative, yet the pipeline never creates edges. `MesMind` exposes no method that calls `GraphDriver.upsert_edge`, and the tests never cover triplet scenarios. +- CRUD operations (`add_memory`, `update_memory`, `delete_memory`) are discussed as mid-level helpers. Only the lower-level `MemoryManager` class (not surfaced through `MeshMind`) contains these methods. + +## Retrieval Capabilities +- README advertises embedding vector search, BM25, LLM reranking, fuzzy search, exact comparison, regex search, filter support, and hybrid methods. The codebase exposes BM25, fuzzy, hybrid, and metadata filters. There are no endpoints for exact comparison, regex search, or LLM reranking, and vector search exists only as a helper inside `MemgraphDriver.vector_search` with no integration. +- The README implies that search operates against the graph database. Actual retrieval utilities work on in-memory lists provided by the caller and do not query Memgraph. +- Usage examples show three layered search calls (`search`, `search_facts`, `search_procedures`). Only the single `search` dispatcher exists, alongside `search_bm25` and `search_fuzzy` helper functions. + +## Data & Relationship Modeling +- README claims memories encompass nodes and edges, including relationship predicates registered ahead of time. The code lacks predicate management beyond an unused `PredicateRegistry` and never writes relationships to the database. +- Low-level `add_triplet` examples assume subject/object lookups by name. `MemgraphDriver.upsert_edge` expects UUIDs and assumes the nodes already exist, so the documented behavior cannot work. +- Memory importance, consolidation, and expiry are presented as rich features. Implementations are minimal: importance defaults to `1.0`, consolidation simply keeps the highest-importance duplicate in-memory, and expiry only runs inside a Celery task that depends on optional infrastructure. + +## Configuration & Dependencies +- README omits instructions for registering embedding encoders. In practice, `meshmind.pipeline.extract` fails with `KeyError` unless `EncoderRegistry.register` is called before extraction. +- README suggests broad Python support and effortless setup. `pyproject.toml` requires Python 3.13, yet many dependencies (Celery, mgclient) do not publish wheels for that version. Missing `mgclient` triggers an import-time failure inside `MeshMind.__init__`. +- Required environment variables (OpenAI API key, Memgraph URI/credentials, Redis URL) are not documented. The README examples instantiate `MeshMind()` with no mention of configuration, but the code depends on these settings. +- README instructs storing and searching without mentioning external services. The current implementation requires a running Memgraph instance and optional Redis for Celery beat. + +## Example Code Path +- README’s extraction example instantiates `MeshMind`, registers entity classes, and expects memories to be generated with corresponding attributes. Actual extraction enforces `entity_label` only by string name and returns `Memory` objects rather than instances of the provided Pydantic models. +- The `mesh_mind.store_memory(memory)` loop in the README references a nonexistent method; the equivalent in code is `store_memories([memory])` or direct use of `MemoryManager`. +- Search examples call `mesh_mind.search`, `search_facts`, and `search_procedures`. Only `search` exists, and it requires preloaded `Memory` objects. +- Update/delete examples rely on `mesh_mind.update_memory` and `mesh_mind.delete_memory`, which are absent. + +## Tooling & Operations +- README does not mention the need to register encoders before running the CLI; the default ingest command fails unless `EncoderRegistry` has an entry matching the configured embedding model. +- README implies functioning Celery maintenance processes. The Celery tasks are importable but disabled when dependencies are missing, and they do not persist consolidated results. +- README lacks troubleshooting guidance for OpenAI SDK changes. The shipped code uses response access patterns (`response['data']`) incompatible with the current SDK, leading to runtime errors. + +## Documentation State +- README positions the document as the authoritative source of truth, yet large sections (triplet storage, relationship management, retrieval coverage) describe unimplemented functionality. +- The original README does not point readers to supporting documents such as configuration references, dependency requirements, or operational runbooks, leaving gaps for anyone onboarding today. diff --git a/FINDINGS.md b/FINDINGS.md new file mode 100644 index 0000000..4fb2bf4 --- /dev/null +++ b/FINDINGS.md @@ -0,0 +1,32 @@ +# Findings + +## General Observations +- The codebase compiles a wide range of functionality, but most modules are loosely integrated; many components are present without being wired into the main client or CLI flows. +- Optional dependencies are imported eagerly (OpenAI, tiktoken, mgclient). In minimal environments the package cannot be imported without installing every dependency. +- Tests capture the intended behaviors more accurately than the runtime code, yet they rely on deprecated OpenAI interfaces and attributes that no longer exist on the models. + +## Dependency & Environment Issues +- `MeshMind` always instantiates `MemgraphDriver`, so installing `mgclient` and having a reachable Memgraph instance are prerequisites even for local experimentation. +- `meshmind.pipeline.extract` calls the OpenAI Responses API but never registers an encoder. Unless callers register a matching encoder manually, extraction fails before returning any memories. +- The OpenAI SDK objects expose attributes rather than dictionary keys, so `response['data']` in `OpenAIEmbeddingEncoder` raises at runtime. +- `meshmind.core.utils` imports `tiktoken` globally; importing `meshmind.core.utils` without the package installed raises immediately. +- Celery and Redis are referenced in tasks, yet `docker-compose.yml` does not provision those services. There is no documented way to launch a working stack. + +## Data Flow & Persistence +- Pipelines only ever upsert nodes. `GraphDriver.upsert_edge` and the `Triplet` model are unused, so relationship data is currently lost. +- Compression, consolidation, and expiry utilities operate on lists of `Memory` objects but do not persist the results back into the graph within standard workflows. +- `Memory.importance` defaults to `1.0` and is not recalculated; there is no ranking algorithm or heuristics as described in the README. + +## CLI & Tooling +- The CLI hardcodes `entity_types=[Memory]` for extraction, which undermines the intent of user-defined entity models. +- CLI ingestion will fail without an encoder registered under the configured embedding model name, yet the CLI does not perform or document this registration step. +- `docker-compose.yml` appears to be a placeholder. It lacks service definitions for Memgraph or Redis and cannot launch the environment described in the README. + +## Testing & Quality +- Tests import `pytest` but there is no automated workflow or Makefile target for running them; `pytest` will still fail because mgclient, OpenAI, and tiktoken are missing. +- Some tests reference `Memory.pre_init` hooks that are absent from the production model, indicating drift between tests and implementation. +- There is no linting, formatting, or type-checking configuration, despite the project aiming for production-level reliability. + +## Documentation +- The new README must explain encoder registration, dependency setup, and realistic capabilities. The existing README significantly overpromises. +- Additional developer onboarding materials (environment setup, service provisioning, troubleshooting) are required to make the project approachable. diff --git a/ISSUES.md b/ISSUES.md new file mode 100644 index 0000000..45e8572 --- /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 +- [ ] Implement relationship persistence (`GraphDriver.upsert_edge`) within the storage pipeline and expose triplet APIs. +- [ ] Restore high-level API methods promised in README (`register_entity`, predicate management, `add_memory`, `update_memory`, `delete_memory`). +- [ ] Ensure CLI ingestion registers entity models and embedding encoders or fails fast with actionable messaging. +- [ ] Provide configuration documentation and examples for Memgraph, Redis, and OpenAI environment variables. +- [ ] 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. +- [ ] 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. +- [ ] Reconcile tests that depend on `Memory.pre_init` and outdated OpenAI interfaces with the current implementation. +- [ ] 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/NEEDED_FOR_TESTING.md b/NEEDED_FOR_TESTING.md new file mode 100644 index 0000000..ad6f3cf --- /dev/null +++ b/NEEDED_FOR_TESTING.md @@ -0,0 +1,51 @@ +# Needed for Testing MeshMind + +## Python Runtime +- Python 3.11 or 3.12 is recommended; project metadata claims 3.13+, but several dependencies (e.g., `pymgclient`, `sentence-transformers`) + have not been validated there. +- A virtual environment (via `venv`, `uv`, or `conda`) to isolate dependencies. + +## Python Dependencies +- Install the project in editable mode: `pip install -e .` from the repository root. +- Ensure optional extras that ship as hard dependencies in `pyproject.toml` are present: + - `pymgclient` for Memgraph connectivity. + - `celery[redis]` for scheduled maintenance tasks. + - `tiktoken`, `sentence-transformers`, `rapidfuzz`, `scikit-learn`, and `numpy` for embedding, retrieval, and compression. +- Additional development tooling that may be required when running the current test suite: + - `pytest` (already listed but verify it installs successfully under the chosen Python version). + - `python-dotenv` if you plan to load environment variables from a `.env` file. + +## External Services and Infrastructure +- **Memgraph** (preferred) or a Neo4j-compatible Bolt graph database reachable at the URI exported via `MEMGRAPH_URI`. + - Requires the Bolt port (default `7687`) to be exposed. + - Ensure user credentials provided in `MEMGRAPH_USERNAME` and `MEMGRAPH_PASSWORD` have write access. +- **Redis** instance when exercising Celery tasks (expiry, consolidation, compression). Set its location via `REDIS_URL`. +- **OpenAI API access** for extraction and embedding encoders used throughout the pipeline. +- Optional but useful: an orchestration layer (Docker Compose or Kubernetes) to manage Memgraph and Redis in tandem if you plan to + mimic production workflows. + +## Environment Variables +- `OPENAI_API_KEY` — required for any extraction or embedding calls via the OpenAI SDK. +- `MEMGRAPH_URI` — Bolt connection string, e.g., `bolt://localhost:7687`. +- `MEMGRAPH_USERNAME` — username for the Memgraph (or Neo4j) instance. +- `MEMGRAPH_PASSWORD` — password for the database user. +- `REDIS_URL` — Redis connection URI (defaults to `redis://localhost:6379/0`). +- `EMBEDDING_MODEL` — key used by `EncoderRegistry` (defaults to `text-embedding-3-small`). Ensure a matching encoder is + registered at runtime before running ingestion or retrieval steps. + +## Local Configuration Steps +- Register an embedding encoder before tests that rely on embeddings: + ```python + from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder + EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder("text-embedding-3-small")) + ``` +- Provide seed data or fixtures for the graph database if end-to-end tests assume pre-existing memories. +- Optionally create a `.env` file mirroring the environment variables above for convenient local setup. + +## Current Blockers in This Environment +- Neo4j and Memgraph binaries are not available, and container tooling (Docker, `mgconsole`, `neo4j-admin`) cannot be installed, + preventing local graph database provisioning inside this workspace. +- Outbound network restrictions may block installation of proprietary dependencies or remote database provisioning without + pre-baked artifacts. +- Redis is likewise unavailable without Docker or host-level package managers; Celery tasks cannot be validated locally until a + remote instance is supplied. diff --git a/NEW_README.md b/NEW_README.md new file mode 100644 index 0000000..78c2b55 --- /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. The current code turns unstructured text into `Memory` records, applies light preprocessing, and stores them via a Memgraph driver. Retrieval helpers run in-memory using lexical, fuzzy, and hybrid scoring strategies. The project is a work in progress; many features described in the legacy README are not yet implemented. + +## Status at a Glance +- ✅ High-level client (`meshmind.client.MeshMind`) with helpers for extraction, preprocessing, and storage. +- ✅ Pipelines for deduplication, default importance scoring, compression, and persistence of memory nodes. +- ✅ Retrieval helpers for TF-IDF (BM25-style), fuzzy string matching, and hybrid vector + lexical ranking (requires registered encoder). +- ✅ Celery task stubs for expiry, consolidation, and compression (require Celery + Redis + Memgraph to function). +- ⚠️ Relationship handling (triplets, predicate registration) is scaffolded but not wired into the storage pipeline. +- ⚠️ CLI ingestion requires manual encoder registration and a running Memgraph instance with `mgclient` installed. +- ❌ Mid-level CRUD APIs (`add_memory`, `update_memory`, `delete_memory`), triplet storage, regex search, and LLM reranking are not implemented. + +## Requirements +- Python 3.11 or 3.12 recommended (project metadata claims 3.13 but dependency support is unverified). +- Memgraph instance reachable via Bolt and the `mgclient` Python package. +- OpenAI API key for extraction and embeddings. +- Optional but recommended: Redis and Celery for scheduled maintenance tasks. +- Additional Python packages installed via `pip install -e .` (see `pyproject.toml`). Some optional modules (`tiktoken`, `sentence-transformers`) are required for specific features. + +## Installation +1. Create and activate a virtual environment using Python 3.11 or 3.12. +2. Install the package in editable mode: + ```bash + pip install -e . + ``` +3. Install optional dependencies as needed: + ```bash + pip install mgclient tiktoken sentence-transformers redis celery + ``` +4. Export required environment variables (adapt values to your setup): + ```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.pipeline.extract` expects an encoder registered in `meshmind.core.embeddings.EncoderRegistry` that matches `settings.EMBEDDING_MODEL`. Before calling extraction or the CLI, register an encoder: +```python +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder +EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder("text-embedding-3-small")) +``` +For offline experimentation you may register a custom encoder that returns deterministic embeddings. + +## Quick Start (Python) +```python +from meshmind.client import MeshMind +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder +from meshmind.core.types import Memory + +# Register an embedding encoder once at startup +EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder()) + +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) +``` +This workflow persists `Memory` nodes in Memgraph. Relationships are not yet created automatically. + +## Command-Line Ingestion +The CLI performs the same pipeline for files and folders of text documents. +```bash +meshmind ingest \ + --namespace demo \ + --instructions "Extract key facts as Memory objects." \ + ./path/to/text/files +``` +Before running: +- Ensure `mgclient` can connect to Memgraph using credentials in environment variables. +- Register an embedding encoder in a startup script (e.g., run a small Python snippet prior to invocation) or extend the CLI to perform registration. +- Provide a `.env` file or export environment variables for configuration. + +## Retrieval Helpers +Retrieval utilities operate on in-memory lists of `Memory` objects. Load the records from your graph (e.g., via `MemoryManager.list_memories`) before calling these helpers. +```python +from meshmind.api.memory_manager import MemoryManager +from meshmind.core.types import SearchConfig +from meshmind.db.memgraph_driver import MemgraphDriver +from meshmind.retrieval.search import search + +# Load memories from Memgraph +manager = MemoryManager(MemgraphDriver("bolt://localhost:7687")) +memories = manager.list_memories(namespace="demo") + +# Register the same encoder used during ingestion +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder +EncoderRegistry.register("text-embedding-3-small", OpenAIEmbeddingEncoder()) + +config = SearchConfig(encoder="text-embedding-3-small", top_k=10) +results = search("Python", memories, namespace="demo", config=config) +for memory in results: + print(memory.name, memory.metadata) +``` +`search_bm25` and `search_fuzzy` are also available for lexical and fuzzy-only scoring. Vector search against Memgraph is not implemented; hybrid search uses embeddings stored on each memory object. + +## Maintenance Tasks +`meshmind.tasks.scheduled` defines Celery beat jobs for: +- Expiring memories once `created_at + ttl_seconds` has elapsed. +- Consolidating duplicate names by keeping the highest-importance record. +- Compressing long metadata content with tiktoken. + +To enable these tasks: +1. Ensure Celery and Redis are installed and running. +2. Start a Celery worker with the meshmind app: + ```bash + celery -A meshmind.tasks.celery_app.app worker -B + ``` + The module attempts to instantiate `MemgraphDriver` at import time; provide valid credentials and ensure `mgclient` is available. + +## Testing +- Pytests live under `meshmind/tests`. They rely on heavy monkeypatching and may need updates to align with the current OpenAI SDK. +- Some tests assume fixtures or hooks (`Memory.pre_init`) that are absent; expect failures until the suite is modernized. +- No continuous integration pipeline is currently provided. + +## Known Limitations +- No triplet/relationship persistence; only nodes are stored. +- Mid-level CRUD helpers and predicate registration are missing from the public client API. +- CLI lacks ergonomics for registering encoders or custom entity models. +- Optional dependencies (`tiktoken`, `sentence-transformers`) are required at import time, leading to crashes when absent. +- The project requires significant configuration (Memgraph, Redis, OpenAI) before any end-to-end scenario succeeds. + +## Roadmap Snapshot +See `PROJECT.md` and `PLAN.md` for prioritized workstreams, including restoring the documented API, improving dependency handling, and expanding retrieval coverage. diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..1dc1ca8 --- /dev/null +++ b/PLAN.md @@ -0,0 +1,52 @@ +# Plan of Action + +## Phase 1 – Stabilize Runtime Basics +1. **Dependency Guards** + - Add lazy imports and feature flags for `mgclient`, `tiktoken`, `celery`, and `sentence-transformers` so missing packages do not crash the import path. + - Provide clear error messages and documentation when dependencies are absent. +2. **Default Encoder Registration** + - Introduce a bootstrap helper that registers an OpenAI encoder (or configurable fallback) on package import or via CLI option. + - Update CLI to call the bootstrap helper before extraction. +3. **OpenAI SDK Compatibility** + - Refactor `OpenAIEmbeddingEncoder` and pipeline extraction to use the latest `openai` SDK response objects with proper retry and error handling. +4. **Configuration Clarity** + - Publish a setup guide covering environment variables, Memgraph/Redis provisioning, and CLI usage. + +## Phase 2 – Restore Promised API Surface +1. **Entity & Predicate Registry Wiring** + - Connect `EntityRegistry` and `PredicateRegistry` to `MeshMind`, ensuring registered models and predicates persist to the database. +2. **CRUD & Triplet Support** + - Add `add_memory`, `update_memory`, `delete_memory`, `register_entity`, `register_allowed_predicates`, and `add_triplet` methods on `MeshMind` that wrap pipeline + driver operations. + - Extend storage pipeline to create relationships via `GraphDriver.upsert_edge` with sensible defaults for subject/object resolution. +3. **Relationship-Aware Examples** + - Update example scripts and documentation to demonstrate triplet storage and retrieval once implemented. + +## Phase 3 – Retrieval & Maintenance Enhancements +1. **Search Coverage** + - Implement vector-only, regex, and exact-match search helpers and expose them through `meshmind.retrieval.search`. + - Optionally integrate LLM reranking for high-quality results. +2. **Maintenance Tasks** + - Ensure consolidation and compression tasks persist results back to Memgraph. + - Move Celery driver initialization into task functions to avoid import-time failures and add logging for missing dependencies. +3. **Importance Scoring Improvements** + - Replace the constant importance score with a heuristic or LLM-based evaluator aligned with README claims. + +## Phase 4 – Developer Experience & Tooling +1. **Testing Overhaul** + - Modernize pytest suites to align with new SDKs, provide fixtures for stub drivers, and ensure tests run without external services. + - Add coverage for new API methods and relationship handling. +2. **Automation & CI** + - Expand the Makefile with lint, format, test, and type-check targets. + - Configure CI (GitHub Actions or similar) to run the suite and static checks on push/PR. +3. **Environment Provisioning** + - Replace `docker-compose.yml` with services for Memgraph and Redis or document alternative local development setups. + +## Phase 5 – Strategic Enhancements +1. **Pluggable Storage Backends** + - Abstract `GraphDriver` further to support alternative backends (Neo4j, in-memory driver, SQLite prototype). +2. **Service Interfaces** + - Expose REST/gRPC endpoints for ingestion and retrieval to enable external integrations. +3. **Operational Observability** + - Add logging, metrics, and dashboards for maintenance jobs, ingestion throughput, and retrieval latency. +4. **Onboarding & Documentation** + - Promote `NEW_README.md` to `README.md`, archive the legacy document, and maintain `SOT.md` with diagrams and workflow maps. diff --git a/PROJECT.md b/PROJECT.md new file mode 100644 index 0000000..ba9d9f6 --- /dev/null +++ b/PROJECT.md @@ -0,0 +1,69 @@ +# MeshMind Project Overview + +## Vision and Scope +- Build a practical memory service that turns unstructured text into graph-backed `Memory` records. +- Provide pipelines for extraction, preprocessing, storage, and retrieval without tightly coupling to a specific UI. +- Support background maintenance (expiry, consolidation, compression) once storage and scheduling dependencies are available. + +## Current Architecture Snapshot +- **Client façade**: `meshmind.client.MeshMind` wires together an OpenAI client, a configured embedding model name, and a `MemgraphDriver` instance. Every workflow starts here. +- **Pipelines**: Extraction (LLM + function-calling), preprocessing (deduplicate, score, compress), and storage utilities live in `meshmind.pipeline`. +- **Graph layer**: `meshmind.db` exposes an abstract `GraphDriver` and a Memgraph implementation that relies on `mgclient`. +- **Retrieval helpers**: `meshmind.retrieval` operates on in-memory `Memory` lists with TF-IDF, fuzzy, and hybrid (vector + lexical) scoring. +- **Task runners**: `meshmind.tasks` defines Celery wiring for expiry, consolidation, and compression jobs when Redis and Celery are present. +- **Support code**: `meshmind.core` contains models, configuration, similarity utilities, and embedding encoder helpers. Optional dependencies (OpenAI, sentence-transformers, tiktoken) are required at runtime for many modules. +- **Tooling**: A CLI ingest command (`meshmind ingest`) demonstrates the extract → preprocess → store loop. Tests exist but depend on heavy monkeypatching and outdated SDK assumptions. + +## Implemented Capabilities +- Serialize knowledge as `Memory` Pydantic models with namespace, entity label, metadata, optional embeddings, timestamps, TTL, and importance fields. +- Extract structured memories from text via the OpenAI Responses API using function-calling against the `Memory` schema (requires manual encoder registration). +- Deduplicate memories by name and (optionally) cosine similarity, assign a default importance score, and truncate metadata content using a tiktoken-based compressor. +- Persist memory nodes to Memgraph by calling `GraphDriver.upsert_entity` for each record. +- Run lexical, fuzzy, and hybrid retrieval against caller-provided in-memory lists of `Memory` objects, including optional metadata and namespace filters. +- Schedule expiry, consolidation, and compression maintenance tasks through Celery beat when both Celery and Redis are configured and the Memgraph driver initializes successfully. +- Provide an example script and CLI entry point that ingest plaintext files into a configured Memgraph instance. + +## Partially Implemented or Fragile Areas +- Hybrid search depends on an encoder registered in `EncoderRegistry`; nothing is auto-registered, so out-of-the-box calls fail. +- The OpenAI embedding wrapper assumes dictionary-style responses and does not match the latest SDK payload objects. +- Celery tasks instantiate the Memgraph driver at import time; without `mgclient` they silently degrade to no-ops. +- The compressor and utility helpers import `tiktoken` eagerly and fail if the package is absent. +- Tests reference hooks (`Memory.pre_init`, OpenAI chat completions) that are no longer present, so the suite does not execute cleanly. +- Python 3.13 is declared in `pyproject.toml`, yet third-party dependencies (Celery, mgclient) have not been validated for that interpreter. + +## Missing or Broken Capabilities +- No public API for registering entities, predicates, or storing triplets as promised in the legacy README. +- Graph relationships are never created because the storage pipeline only upserts nodes. +- There is no mid-level `add_memory`/`update_memory`/`delete_memory` surface on `MeshMind`; the CLI relies solely on extraction and store helpers. +- Vector search, regex search, exact match search, and LLM re-ranking endpoints described in the README are absent. +- Memory consolidation and expiry are not integrated into the ingestion workflow, and consolidation never writes results back to the database. +- Configuration guidance is minimal; missing environment variables lead to runtime failures when constructing the client. + +## External Services & Dependencies +- **Memgraph + mgclient**: Required for any persistence. Without `mgclient`, constructing `MeshMind` raises immediately. +- **OpenAI SDK**: Needed for both extraction and embeddings. Newer SDK versions return typed objects, not dicts, which breaks current assumptions. +- **tiktoken**: Used by compression and token counting utilities. Imported at module load time without fallbacks. +- **RapidFuzz, scikit-learn, numpy**: Support fuzzy and lexical retrieval. +- **Celery + Redis**: Optional but necessary for scheduled maintenance tasks. +- **sentence-transformers**: Optional embedding backend for offline models. + +## Tooling and Operational State +- `docker-compose.yml` lists services but does not provision Memgraph or Redis containers. +- No encoder instances are registered automatically; setup scripts are missing. +- Pytests rely on manual monkeypatches to simulate OpenAI and mgclient. Running `pytest` out of the box fails due to missing optional dependencies and incompatible SDK interfaces. +- Continuous integration or linting workflows are not defined. + +## Roadmap Highlights +- Restore the high-level API surface promised in the README (entity registration, predicate management, CRUD helpers, triplet storage). +- Introduce a safe, dependency-light initialization path (lazy imports, graceful fallbacks, injectable storage backends). +- Expand retrieval to include driver-backed vector search, regex/exact match helpers, and optional LLM-based re-ranking. +- Implement relationship persistence and richer metadata handling in the graph layer. +- Harden maintenance jobs to run independently of import-time side effects and to write results back to Memgraph. +- Rewrite the test suite around modern OpenAI SDK semantics and provide fixtures for running without external services. +- Document setup thoroughly (encoders, environment variables, dependency installation, service provisioning) and provide automation scripts. + +## Future Potential Extensions +- Plug-in architecture for alternative vector databases or document stores. +- Streaming ingestion workers that watch queues or webhooks instead of filesystem batches. +- UI or API gateway to expose memory search and curation to downstream agents or humans. +- Analytics dashboards that summarize namespace health, expiry cadence, and consolidation outcomes. diff --git a/RECOMMENDATIONS.md b/RECOMMENDATIONS.md new file mode 100644 index 0000000..fce74de --- /dev/null +++ b/RECOMMENDATIONS.md @@ -0,0 +1,33 @@ +# Recommendations + +## Stabilize the Foundation +- Refactor `MeshMind` initialization so that graph and OpenAI dependencies are optional or injectable, enabling local development without Memgraph. +- Provide a bootstrap module (or CLI option) that registers default embedding encoders and entity models before extraction. +- Update the OpenAI integration to use the current SDK response objects and add robust error handling for rate limits and API failures. +- Introduce optional dependency guards across the package; defer importing `tiktoken`, `mgclient`, and `celery` until the functionality is invoked. +- Align Python version support with dependency availability (target 3.11/3.12 until 3.13 is validated). + +## Restore Promised Functionality +- Implement entity and predicate registries with persistence hooks so that README workflows (`register_entity`, `add_predicate`, `add_triplet`) become real. +- Add mid-level CRUD methods (`add_memory`, `update_memory`, `delete_memory`) on `MeshMind` that delegate to `MemoryManager` and ensure stored records round-trip correctly. +- Extend the storage pipeline to create relationships via `GraphDriver.upsert_edge`, including handling for subject/object lookups by name or UUID. +- Build out retrieval helpers for vector-only, regex, and exact-match queries, and optionally integrate reranking via the LLM client. + +## Improve Developer Experience +- Replace the placeholder `docker-compose.yml` with services for Memgraph and Redis (or document how to run them separately). +- Ship sample scripts that register encoders, seed demo data, and demonstrate retrieval end-to-end. +- Add Makefile tasks for running tests, linting, type checking, and starting Celery workers. +- Modernize the pytest suite to rely on fixtures that do not require live services and that mirror the new OpenAI SDK APIs. +- Set up continuous integration to run unit tests and static checks on every change. + +## Documentation & Onboarding +- Promote `NEW_README.md` to `README.md` after validation and archive the legacy document for historical reference. +- Document configuration and dependency expectations in a dedicated setup guide linked from the README. +- Expand `SOT.md` with diagrams or tables that map modules to workflows once the architecture stabilizes. +- Provide troubleshooting steps for common failures (missing encoder registration, mgclient import errors, OpenAI authentication). + +## Future Enhancements +- Explore alternative storage backends (e.g., Neo4j driver, SQLite) for environments without Memgraph. +- Offer a lightweight REST or gRPC API to interact with memories programmatically. +- Instrument maintenance jobs with metrics and logging so operators can observe expiry/consolidation outcomes. +- Investigate incremental ingestion pipelines (message queues, streaming connectors) for real-time memory updates. diff --git a/SOT.md b/SOT.md new file mode 100644 index 0000000..4b21d05 --- /dev/null +++ b/SOT.md @@ -0,0 +1,127 @@ +# MeshMind Source of Truth + +This document captures the current architecture, modules, and operational expectations for MeshMind. Update it whenever code structure or workflows change so new contributors can ramp up quickly. + +## Repository Layout +``` +meshmind/ +├── api/ # MemoryManager CRUD adapter +├── cli/ # CLI entry point and ingest command +├── client.py # High-level MeshMind façade +├── core/ # Config, types, embeddings, similarity, utilities +├── db/ # Graph driver abstractions (Memgraph) +├── models/ # Entity/predicate registries (not yet integrated) +├── pipeline/ # Extraction, preprocessing, storage, maintenance steps +├── retrieval/ # In-memory search strategies and filters +├── tasks/ # Celery wiring and scheduled jobs +├── tests/ # Pytest suites (require extensive monkeypatching) +└── examples/ # Example extraction/preprocess/store script +``` +Supporting files include: +- `pyproject.toml`: project metadata (declares Python >=3.13, which is aspirational). +- `docker-compose.yml`: placeholder with no services defined. +- `Makefile`: minimal development targets (currently none for testing). +- Documentation artifacts (`PROJECT.md`, `PLAN.md`, etc.). + +## Configuration (`meshmind/core/config.py`) +- Loads environment variables for Memgraph (`MEMGRAPH_URI`, `MEMGRAPH_USERNAME`, `MEMGRAPH_PASSWORD`), Redis (`REDIS_URL`), OpenAI (`OPENAI_API_KEY`), and default embedding model (`EMBEDDING_MODEL`). +- Uses `python-dotenv` if available to load `.env` files at import time. +- Exposes a module-level `settings` object consumed by clients, drivers, and tasks. + +## Core Data Models (`meshmind/core/types.py`) +- `Memory`: Pydantic model representing a knowledge record. Fields include `uuid`, `namespace`, `name`, `entity_label`, optional `embedding`, `metadata`, `reference_time`, `created_at`, `updated_at`, `importance`, and `ttl_seconds`. +- `Triplet`: Subject–predicate–object relationship with namespace, entity label, metadata, and optional reference time. Not used in current flows. +- `SearchConfig`: Retrieval configuration (encoder name, `top_k`, `rerank_k`, metadata filters, hybrid weights). + +## Client (`meshmind/client.py`) +- `MeshMind` constructor wires: + - `OpenAI()` as the default LLM client (fails if OpenAI SDK is missing or API key absent). + - Embedding model name from `settings.EMBEDDING_MODEL`. + - `MemgraphDriver` instantiated with configured URI/credentials (raises `ImportError` if `mgclient` is unavailable). +- Provides high-level helpers: + - `extract_memories` delegates to `pipeline.extract.extract_memories` (requires an encoder in `EncoderRegistry`). + - `deduplicate`, `score_importance`, `compress` delegate to `pipeline.preprocess`. + - `store_memories` delegates to `pipeline.store.store_memories`. +- Does **not** expose registration, triplet storage, or CRUD methods promised in the legacy README. + +## Embeddings & Utilities (`meshmind/core/embeddings.py`, `meshmind/core/utils.py`) +- `EncoderRegistry`: Class-level map from string key to encoder instance. Call `register(name, encoder)` before extraction or hybrid search. +- `OpenAIEmbeddingEncoder`: Wraps the OpenAI Embeddings API with retry logic but assumes dictionary-style responses (`response['data']`), which is incompatible with current SDK objects. +- `SentenceTransformerEncoder`: Provides local embedding support via `sentence-transformers`. +- `meshmind.core.utils`: Supplies UUID generation, timestamp helpers, hashing, and token counting. Imports `tiktoken` at module load, so missing the package raises immediately. + +## Database Layer (`meshmind/db`) +- `base_driver.py`: Abstract `GraphDriver` defining `upsert_entity`, `upsert_edge`, `find`, and `delete` signatures. +- `memgraph_driver.py`: + - Imports `mgclient` and opens a Bolt connection on instantiation. + - Implements node upserts (MERGE by `uuid`), edge upserts (MERGE by `uuid`), arbitrary Cypher `find`, deletion, and a naive Python-based vector search over stored embeddings. + - Requires `mgclient`; otherwise constructing the driver raises `ImportError`. + +## Pipeline Modules (`meshmind/pipeline`) +1. **Extraction (`extract.py`)** + - Builds an OpenAI Responses API call with function-calling against the `Memory` JSON schema. + - Validates `entity_label` against supplied `entity_types` (string comparison only) and populates embeddings via `EncoderRegistry`. +2. **Preprocess (`preprocess.py`)** + - `deduplicate`: Removes duplicates by name and optionally by cosine similarity when embeddings exist. + - `score_importance`: Assigns a default importance of `1.0` when missing. + - `compress`: Delegates to `pipeline.compress.compress_memories` and falls back on errors. +3. **Compress (`compress.py`)** + - Uses `tiktoken` to truncate `metadata['content']` to a token budget (requires the package). +4. **Consolidate (`consolidate.py`)** + - Groups memories by name and selects the highest-importance entry (no persistence built in). +5. **Expire (`expire.py`)** + - Deletes memories whose `created_at + ttl_seconds` is in the past using `MemoryManager`. +6. **Store (`store.py`)** + - Iterates memories and calls `GraphDriver.upsert_entity`. Relationships are not touched. + +## Retrieval (`meshmind/retrieval`) +- `filters.py`: Filter helpers by namespace, entity labels, and metadata equality. +- `bm25.py`: TF-IDF vectorizer + cosine similarity (scikit-learn) used as a lexical scorer. +- `fuzzy.py`: RapidFuzz WRatio scoring for fuzzy name matching. +- `hybrid.py`: Combines query embeddings (from registered encoder) with BM25 scores using configurable weights. +- `search.py`: Dispatchers for hybrid (`search`), lexical (`search_bm25`), and fuzzy (`search_fuzzy`) retrieval. Operate on caller-provided lists of `Memory` objects; no direct graph querying. + +## CLI (`meshmind/cli`) +- `__main__.py`: Defines the `meshmind` CLI with an `ingest` subcommand. +- `ingest.py`: Walks files/directories, reads text contents, runs extraction + preprocessing + storage. Hardcodes `entity_types=[Memory]` and assumes an encoder is already registered. + +## Tasks (`meshmind/tasks`) +- `celery_app.py`: Creates a Celery app when the library is installed; otherwise exposes a no-op shim. +- `scheduled.py`: + - Attempts to instantiate `MemgraphDriver` and `MemoryManager` at import time, falling back to `None` when dependencies fail. + - Configures Celery beat schedules for expiry (daily), consolidation (every 6 hours), and compression (every 12 hours). + - Defines tasks that operate on the global `manager` instance; if initialization failed they return empty results. + +## API Adapter (`meshmind/api/memory_manager.py`) +- Wraps a graph driver to provide CRUD helpers (`add_memory`, `update_memory`, `delete_memory`, `get_memory`, `list_memories`). +- Converts Pydantic objects to dicts via `memory.dict(exclude_none=True)` with fallback to `__dict__`. +- Currently the primary way to list memories for retrieval; not exposed through the CLI or `MeshMind` convenience methods. + +## Models (`meshmind/models/registry.py`) +- `EntityRegistry` and `PredicateRegistry` store registered models and allowed relationship labels. +- No production code writes to these registries yet; integrating them is part of the future roadmap. + +## Examples & Tests +- `examples/extract_preprocess_store_example.py`: Demonstrates extraction and storage using `MeshMind`. Requires valid OpenAI credentials and Memgraph. +- Tests under `meshmind/tests` cover extraction, preprocessing, driver behavior, retrieval, similarity, and maintenance tasks. They rely on monkeypatching dummy encoders, OpenAI clients, and mgclient modules. Some tests assume attributes (`Memory.pre_init`) that are not defined in production code, so the suite will fail until updated. + +## External Dependencies +- **OpenAI SDK**: Required for extraction and embeddings. Update to latest `openai` package and adjust code accordingly. +- **mgclient**: Required for Memgraph persistence; missing package prevents `MeshMind` construction. +- **tiktoken**: Required for compression and utility token counting. Currently imported eagerly. +- **scikit-learn**, **rapidfuzz**, **numpy**: Support retrieval algorithms. +- **sentence-transformers** (optional): Alternative embedding encoder. +- **celery** and **redis** (optional): Required for scheduled maintenance tasks. + +## Operational Caveats +- No encoder is registered by default; failing to register one causes extraction and hybrid search to raise `KeyError`. +- `MeshMind` cannot be instantiated in environments lacking Memgraph or mgclient, limiting portability. +- Relationship data is not persisted, so graph analyses beyond isolated nodes are impossible. +- Tests and CLI commands assume manual setup of encoders and environment variables. +- `docker-compose.yml` does not start required services; developers must provision Memgraph and Redis separately. + +## Related Documentation +- `PROJECT.md`: Architectural summary, capability matrix, and roadmap themes. +- `PLAN.md`: Actionable next steps to close gaps. +- `DISCREPANCIES.md`: Detailed comparison between the legacy README and actual implementation. +- `RECOMMENDATIONS.md`: Suggested improvements ranked by impact. diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..ee792c5 --- /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`. +- [ ] Refresh examples to cover relationship-aware ingestion and retrieval flows. +- [ ] 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. +- [ ] Modernize pytest suites and add fixtures to run without external services. +- [ ] 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/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..d2a184a 100644 --- a/meshmind/client.py +++ b/meshmind/client.py @@ -1,80 +1,175 @@ -""" -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 +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() + + # ------------------------------------------------------------------ + # 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) + 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/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/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/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