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..aa2cbc3 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,28 @@ +# Changelog + +## [Unreleased] - 2025-02-14 +### Added +- Configurable graph driver factory with in-memory, SQLite, Memgraph, and optional Neo4j implementations plus supporting tests. +- REST and gRPC service layers (with FastAPI stub fallback) for ingestion and retrieval, including coverage in the test suite. +- Observability utilities that collect metrics and structured logs across pipelines and scheduled Celery tasks. +- Docker Compose definition provisioning Memgraph, Redis, and a Celery worker for local development. +- 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 +- Settings now surface `GRAPH_BACKEND`, Neo4j, and SQLite options while README/NEEDED_FOR_TESTING document the expanded setup. +- README, README_LATEST, and NEW_README were consolidated so the promoted README reflects current behaviour. +- PROJECT, PLAN, SOT, FINDINGS, DISCREPANCIES, ISSUES, RECOMMENDATIONS, and TODO were refreshed to capture new capabilities and + re-homed backlog items under a "Later" section. +- 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..d49c4ba --- /dev/null +++ b/DISCREPANCIES.md @@ -0,0 +1,57 @@ +# README vs Implementation Discrepancies + +## Overview +- The legacy README has been superseded by `README.md`/`README_LATEST.md`, which now reflect the implemented feature set. +- 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. +- ⚠️ README still omits optional tooling now required by the Makefile/CI (ruff, pyright, typeguard, toml-sort, yamllint); + highlight these prerequisites more prominently. +- ⚠️ 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 now provisions Memgraph, Redis, and a Celery worker; README sections should highlight the workflow and + caveats for environments lacking container tooling. +- ⚠️ Celery tasks remain best-effort shims; README should clarify that maintenance requires the optional infrastructure and that + consolidation outputs are not persisted yet. +- ⚠️ Celery tasks remain best-effort shims; README should clarify that maintenance requires the optional infrastructure. + +## Documentation State +- Continue promoting `README_LATEST.md`/`README.md` as the authoritative guides and propagate updates 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..52a3e18 --- /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 has now been promoted/archived. + +## Dependency & Environment Notes +- `MeshMind` defers driver creation until persistence is required, enabling workflows without `pymgclient` and selecting between in-memory, SQLite, Memgraph, or Neo4j backends. +- 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 now provisions both + services alongside a worker for local testing. + +## 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 graph credentials and OpenAI keys are configured when using external backends. +- The Makefile introduces lint, format, type-check, and test targets, plus Docker helpers. 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.md`/`README_LATEST.md` document setup, pipelines, retrieval, tooling, and now highlight service interfaces and observability. +- 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..ea39667 --- /dev/null +++ b/ISSUES.md @@ -0,0 +1,32 @@ +# Issues Checklist + +## Blockers +- [x] MeshMind client fails without `mgclient`; introduce lazy driver initialization or documented in-memory fallback. +- [x] Register a default embedding encoder (OpenAI or sentence-transformers) during startup so extraction and hybrid search can run. +- [x] Update OpenAI integration to match the current SDK (Responses API payload, embeddings API response structure). +- [x] 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). +- [x] Create real docker-compose services for Memgraph and Redis or remove the placeholder file. +- [ ] Document Neo4j driver requirements and verify connectivity against a live cluster. + +## 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) +- [ ] Implement graph-backed retrieval queries leveraging Memgraph/Neo4j search capabilities once available. +- [ ] 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 +- [x] 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..075b3fa --- /dev/null +++ b/NEEDED_FOR_TESTING.md @@ -0,0 +1,52 @@ +# 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` (for Memgraph backends). +- Optional drivers: install `neo4j` if exercising the Neo4j backend; SQLite support ships with the standard library. +- 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`/`grpcio-tools` (for service interface experimentation). + +## External Services and Infrastructure +- **Graph backend** options: + - In-memory / SQLite require no external services (set `GRAPH_BACKEND=memory` or `sqlite`). + - **Memgraph** reachable via `MEMGRAPH_URI` with credentials exported in `MEMGRAPH_USERNAME`/`MEMGRAPH_PASSWORD`. + - **Neo4j** reachable via `NEO4J_URI` with credentials `NEO4J_USERNAME`/`NEO4J_PASSWORD` when the optional driver is installed. +- **Redis** for Celery task queues, referenced through `REDIS_URL`. +- **OpenAI API access** for extraction, embeddings, and LLM reranking (`OPENAI_API_KEY`). +- Recommended: Docker Compose (shipped in repo) to run Memgraph, Redis, and the Celery worker together when developing locally. + +## Environment Variables +- `GRAPH_BACKEND` — `memory`, `sqlite`, `memgraph`, or `neo4j` (defaults to `memory`). +- `OPENAI_API_KEY` — required for extraction, embeddings, and reranking. +- `MEMGRAPH_URI` — e.g., `bolt://localhost:7687` (when using Memgraph). +- `MEMGRAPH_USERNAME` and `MEMGRAPH_PASSWORD` — credentials for the Memgraph database. +- `NEO4J_URI`, `NEO4J_USERNAME`, `NEO4J_PASSWORD` — optional Neo4j connectivity details. +- `SQLITE_PATH` — filesystem path for the SQLite graph backend (defaults to in-memory). +- `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()`. +- For REST/gRPC testing, instantiate the `RestAPIStub`/`GrpcServiceStub` with the in-memory driver to avoid external services. +- 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; use the in-memory or SQLite drivers instead. +- 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..527b169 --- /dev/null +++ b/NEW_README.md @@ -0,0 +1,4 @@ +# README Consolidated + +The up-to-date project overview now lives in [`README.md`](README.md) and [`README_LATEST.md`](README_LATEST.md). +This file remains as a pointer for historical context. diff --git a/PLAN.md b/PLAN.md new file mode 100644 index 0000000..3fce9fa --- /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/Neo4j search endpoints are available. +2. **Maintenance Tasks** – Tasks now emit telemetry but still return in-memory results. Persist consolidation outputs and improve + failure handling. +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** – Docker Compose now provisions Memgraph, Redis, and a Celery worker. Track multi-backend + examples and ensure documentation stays current. + +## Phase 5 – Strategic Enhancements (Planned) +1. **Graph-Backed Retrieval** – Push search workloads into the configured graph driver once backend capabilities land. +2. **Operational Observability** – Export telemetry to Prometheus/OpenTelemetry and surface dashboards/alerts. +3. **Celery Hardening** – Persist consolidation/compression outputs back into the graph and add retry policies. +4. **Importance Scoring** – Replace constant heuristics with data-driven scoring or LLM evaluation workflows. diff --git a/PROJECT.md b/PROJECT.md new file mode 100644 index 0000000..77de438 --- /dev/null +++ b/PROJECT.md @@ -0,0 +1,87 @@ +# 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 graph driver selected via `GRAPH_BACKEND` (memory, SQLite, Memgraph, Neo4j). 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 in-memory, SQLite, Memgraph, and optional Neo4j drivers + with a shared factory for local testing and production use. +- **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. +- **Service adapters**: `meshmind.api.rest` and `.grpc` expose REST/gRPC entry points (with lightweight stubs for tests) so + ingestion and retrieval can run as services. +- **Observability**: `meshmind.core.observability` collects metrics, gauges, and structured log events across pipelines and + scheduled tasks. +- **Tooling**: The CLI ingest command (`meshmind ingest`), updated example script, Makefile automation, CI workflow, and Docker + Compose file illustrate extraction → preprocessing → storage → retrieval locally. + +## 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. +- Neo4j driver support is import-guarded; automated verification against a live cluster is still pending. +- 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. +- SQLite driver currently stores JSON blobs; future work may normalize columns for structured querying. +- 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. +- 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. +- Metrics remain in-memory; external exporters (Prometheus/OpenTelemetry) are not wired up. +- gRPC wiring currently relies on stubs; production-ready servers are still future work. + +## External Services & Dependencies +- **Graph backend**: Choose via `GRAPH_BACKEND`. In-memory and SQLite require no external services. Memgraph needs `pymgclient`; + Neo4j requires the official driver and a live instance. +- **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 rely on fixtures (`memory_factory`, `dummy_encoder`, in-memory drivers) so they pass without external services, though + optional dependencies may still need installation. +- Docker Compose now provisions Memgraph, Redis, and a Celery worker; see `NEEDED_FOR_TESTING.md` for enabling optional + services locally. + +## Roadmap Highlights +- Implement graph-backed retrieval queries (vector similarity, structured filters) instead of memory list inputs. +- Export observability metrics to external sinks (Prometheus/OpenTelemetry) and surface dashboards. +- Enhance importance scoring with data-driven heuristics or LLM evaluation. +- Harden Celery tasks to persist consolidation output and surface failures clearly. +- Validate Neo4j driver behaviour against a live cluster and ship official test doubles for Memgraph/Redis/encoders. +- 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.md b/README.md index 04882e7..1cfe900 100644 --- a/README.md +++ b/README.md @@ -1,244 +1,155 @@ # MeshMind -MeshMind is a knowledge management system that uses LLMs and graph databases to store and retrieve information. -Adding, Searching, Updating, and Deleting memories is supported. -Retrieval of memories using LLMs and graph databases by comparing embeddings and graph database metadata. - -## Features - -- Adding memories -- Searching memories -- Updating memories -- Deleting memories -- Extracting memories from content -- Expiring memories -- Memory Importance Ranking -- Memory Deduplication -- Memory Consolidation -- Memory Compression - -## Retrieval Methods - -- Embedding Vector Search -- BM25 Retrieval -- ReRanking with LLM -- Fuzzy Search -- Exact Comparison Search -- Regex Search -- Search Filters -- Hybrid Search Methods - -## Components -- OpenAI API -- MemGraphDB (Alternative to Neo4j) -- Embedding Model - -# Types of Memory - -- Long-Term Memory - Persistent Memory - - Explicit Memory - Conscious Memory - Active Hotpath Memory (Triggered in Response to Input) - - Declarative Memory - Conscious Memory - - Semantic Memory - What is known - - Eposodic Memory - What has been experienced - - Procedural Memory - How to perform tasks - - - Implicit Memory - Subconscious Memory - Background Process Memory (Triggered in Intervals) - - Non-Declarative Memory - Subconscious Memory - - Semantic Memory - Implicitly acquired - - Eposodic Memory - Implicitly acquired - - Procedural Memory - Implicitly acquired - -- Short-Term Memory - Transient Memory - - Working Memory (Processing) [reasoning messages, scratchpad, etc] - - Sensory Memory (Input) [user input, system input, etc] - - Log Storage (Output) [assistant responses, tool logs, etc] - -# Methods -- Extract Memory -- Add Memory -- Add Triplet -- Search Memory -- Update Memory -- Delete Memory - -## Usage - +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. +- ✅ Docker Compose provisions Memgraph, Redis, and a Celery worker for local orchestration. +- ✅ Built-in observability surfaces structured events and in-memory metrics for pipelines and scheduled tasks. +- ⚠️ Retrieval still consumes caller-provided lists; graph-backed querying remains future work. + +## Requirements +- Python 3.11 or 3.12 recommended (metadata targets 3.13; verify third-party support before upgrading). +- Configurable graph backend via `GRAPH_BACKEND` (`memory`, `sqlite`, `memgraph`, `neo4j`). +- Memgraph instance reachable via Bolt and the `pymgclient` Python package (when `GRAPH_BACKEND=memgraph`). +- Optional Neo4j instance with the official Python driver (when `GRAPH_BACKEND=neo4j`). +- 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 GRAPH_BACKEND=memory # or memgraph/sqlite/neo4j + export MEMGRAPH_URI=bolt://localhost:7687 + export MEMGRAPH_USERNAME=neo4j + export MEMGRAPH_PASSWORD=secret + export SQLITE_PATH=/tmp/meshmind.db + export NEO4J_URI=bolt://localhost:7687 + export NEO4J_USERNAME=neo4j + export NEO4J_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 import MeshMind -from pydantic import BaseModel, Field - -# Pydantic model's name will be used as the node label in the graph database. -# This defines the attributes of the entity. -# Attributes of the entity are stored in entity metadata['attributes'] -class Person(BaseModel): - first_name: str | None = Field(..., description="First name of the person") - last_name: str | None = Field(None, description="Last name of the person") - description: str | None = Field(None, description="Description of the person") - job_title: str | None = Field(None, description="Job title of the person") - -# Initialize MeshMind -mesh_mind = MeshMind() - -# Register pydantic model as entity -mesh_mind.register_entity(Person) - -# Register allowed relationship predicates labels. -mesh_mind.register_allowed_predicates([ - "employee_of", - "on_team", - "on_project", -]) - -# Add edge - This will create an edge in the graph database. (Alternative to `register_allowed_types`) -mesh_mind.add_predicate("has_skill") - -# Extract memories - [High Level API] - This will create nodes and edges based on the content provided. Extracts memories from the content. -extracted_memories = mesh_mind.extract_memories( - instructions="Extract all memories from the database.", - namespace="Company Employees", - entity_types=[Person], - content=["John Doe, Software Engineer 10 years of experience.", "Jane Doe, Software Engineer 5 years of experience."] - ) +from meshmind.core.embeddings import EncoderRegistry, OpenAIEmbeddingEncoder -for memory in extracted_memories: - # Store memory - This will perform deduplication, add uuid, add timestamps, format memory and store the memory in the graph database using `add_triplet`. - mesh_mind.store_memory(memory) - -# Add memory - [Mid Level API] - This will perform all preprocessing steps i.e. deduplication, add uuid, add timestamps, format memory, etc. and store the memory in the graph database using `add_triplet`. ( Skips extraction and automatically adds memory to the graph database ) **Useful for adding custom memories** -mesh_mind.add_memory( - namespace="Company Employees", - name="John Doe", - entity_label="Person", - entity=Person( - first_name="John", - last_name="Doe", - description="John Doe", - job_title="Software Engineer", - ), - metadata={ - "source": "John Doe Employee Record", - "source_type": "text", - } -) +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. -# `add_triplet` - [Low Level API] -# -# Add a [node, edge, node] triplet - This will create a pair of nodes and a connecting edge in the graph database using db driver. -# This bypasses deduplication and other preprocessing steps and directly adds the data to the graph database. -# This is useful for adding data that is not in the format of a memory. -# -# subject: The subject of the triplet. ( Source Entity Node Label Name ) -# predicate: The predicate of the triplet. ( Relationship between the subject and object, registered using `register_allowed_predicates` ) -# object: The object of the triplet. ( Target Entity Node Label Name ) -# namespace: The namespace of the triplet. ( Group of related nodes and edges ) -# entity_label: The label of the entity. ( Type of node, registered using `register_entity` ) -# metadata: The metadata of the triplet. ( Additional information about the triplet ) -# reference_time: The time at which the triplet was created. ( Optional ) -# -# If the subject, predicate, object, namespace, or entity_label does not exist, it will be created. -# -# Example: -mesh_mind.add_triplet( - subject="John Doe", - predicate="on_project", - object="Project X", - namespace="Company Employees", - entity_label="Person", - metadata={ - "source": "John Doe Project Record", - "source_type": "text", - "summary": "John Doe is on project X.", - "attributes": { - "first_name": "John", - "last_name": "Doe", - "description": "John Doe", - "job_title": "Software Engineer", - } - }, - reference_time="2025-05-09T23:31:51-04:00" +## 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, ) - -# Search - [High Level API] - This will search the graph database for nodes and edges based on the query. -search_results = mesh_mind.search( - query="John Doe", - namespace="Company Employees", - entity_types=[Person], +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", ) - -for search_result in search_results: - print(search_result) - -# Search - [Mid Level API] - This will search the graph database for nodes and edges based on the query. -search_results = mesh_mind.search_facts( - query="John Doe", - namespace="Company Employees", - entity_types=[Person], - config=SearchConfig( - encoder="text-embedding-3-small", - ) - ) - -for search_result in search_results: - print(search_result) - -# Search - [Low Level API] - This will search the graph database for nodes and edges based on the query. -search_results = mesh_mind.search_procedures( - query="John Doe", - namespace="Company Employees", - entity_types=[Person], - config=SearchConfig( - encoder="text-embedding-3-small", - ) - ) - -# Update Memory - Same as `add_memory` but updates an existing memory. -mesh_mind.update_memory( - uuid="12345678-1234-1234-1234-123456789012", - namespace="Company Employees", - name="John Doe", - entity_label="Person", - entity=Person( - first_name="John", - last_name="Doe", - description="John Doe", - job_title="Software Engineer", - ), - metadata={ - "source": "John Doe Employee Record", - "source_type": "text", - } -) - -# Delete Memory -mesh_mind.delete_memory( - uuid="12345678-1234-1234-1234-123456789012" -) - -for search_result in search_results: - print(search_result) - + mm.store_triplets([relation]) ``` -## Command-Line Interface (CLI) +## Retrieval +`MeshMind` exposes multiple retrieval helpers that operate on lists of `Memory` objects (e.g., fetched via +`mm.list_memories(namespace="demo")`). The active graph backend is selected via `GRAPH_BACKEND` or by supplying a driver +instance to `MeshMind`. +```python +from meshmind.core.types import SearchConfig -MeshMind includes a `meshmind` CLI tool for ingesting content via the extract → preprocess → store pipeline. +memories = mm.list_memories(namespace="demo") +config = SearchConfig(encoder=mm.embedding_model, top_k=5, rerank_model="gpt-4o-mini") -Usage: -```bash -meshmind --help +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") ``` -Primary command: +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 \ - -n \ - [-e ] \ - [-i ""] \ - [ ...] + --namespace demo \ + --instructions "Extract key facts as Memory objects." \ + ./path/to/text/files ``` -Example: +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 -meshmind ingest -n demo --embedding-model text-embedding-3-small ./data/articles +celery -A meshmind.tasks.celery_app.app worker -B ``` -This reads text files under `./data/articles`, extracts memories, deduplicates, scores, compresses, -and stores them in your Memgraph database under the `demo` namespace. - +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. + +## Service Interfaces +- **REST** – `meshmind.api.rest.create_app` returns a FastAPI app (or lightweight stub) that exposes `/memories`, `/triplets`, + and `/search` endpoints. +- **gRPC** – `meshmind.api.grpc.GrpcServiceStub` mirrors the ingestion and retrieval RPC surface for integration tests and + future server wiring. + +## Observability +- `meshmind.core.observability.telemetry` collects counters, gauges, and durations for pipelines and Celery tasks. +- `meshmind.core.observability.log_event` emits structured log messages that annotate pipeline progress. +- Metrics remain in-memory today; export hooks (Prometheus, OpenTelemetry) are future enhancements. + +## 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. +- Metrics remain in-memory; no external exporter is wired up yet. + +## Roadmap Snapshot +Consult `PROJECT.md`, `PLAN.md`, and `RECOMMENDATIONS.md` for prioritized enhancements: graph-backed retrieval, maintenance +persistence, external metrics exporters, and richer driver/test doubles. diff --git a/README_LATEST.md b/README_LATEST.md new file mode 100644 index 0000000..1cfe900 --- /dev/null +++ b/README_LATEST.md @@ -0,0 +1,155 @@ +# 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. +- ✅ Docker Compose provisions Memgraph, Redis, and a Celery worker for local orchestration. +- ✅ Built-in observability surfaces structured events and in-memory metrics for pipelines and scheduled tasks. +- ⚠️ Retrieval still consumes caller-provided lists; graph-backed querying remains future work. + +## Requirements +- Python 3.11 or 3.12 recommended (metadata targets 3.13; verify third-party support before upgrading). +- Configurable graph backend via `GRAPH_BACKEND` (`memory`, `sqlite`, `memgraph`, `neo4j`). +- Memgraph instance reachable via Bolt and the `pymgclient` Python package (when `GRAPH_BACKEND=memgraph`). +- Optional Neo4j instance with the official Python driver (when `GRAPH_BACKEND=neo4j`). +- 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 GRAPH_BACKEND=memory # or memgraph/sqlite/neo4j + export MEMGRAPH_URI=bolt://localhost:7687 + export MEMGRAPH_USERNAME=neo4j + export MEMGRAPH_PASSWORD=secret + export SQLITE_PATH=/tmp/meshmind.db + export NEO4J_URI=bolt://localhost:7687 + export NEO4J_USERNAME=neo4j + export NEO4J_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")`). The active graph backend is selected via `GRAPH_BACKEND` or by supplying a driver +instance to `MeshMind`. +```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. + +## Service Interfaces +- **REST** – `meshmind.api.rest.create_app` returns a FastAPI app (or lightweight stub) that exposes `/memories`, `/triplets`, + and `/search` endpoints. +- **gRPC** – `meshmind.api.grpc.GrpcServiceStub` mirrors the ingestion and retrieval RPC surface for integration tests and + future server wiring. + +## Observability +- `meshmind.core.observability.telemetry` collects counters, gauges, and durations for pipelines and Celery tasks. +- `meshmind.core.observability.log_event` emits structured log messages that annotate pipeline progress. +- Metrics remain in-memory today; export hooks (Prometheus, OpenTelemetry) are future enhancements. + +## 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. +- Metrics remain in-memory; no external exporter is wired up yet. + +## Roadmap Snapshot +Consult `PROJECT.md`, `PLAN.md`, and `RECOMMENDATIONS.md` for prioritized enhancements: graph-backed retrieval, maintenance +persistence, external metrics exporters, and richer driver/test doubles. diff --git a/RECOMMENDATIONS.md b/RECOMMENDATIONS.md new file mode 100644 index 0000000..aec308a --- /dev/null +++ b/RECOMMENDATIONS.md @@ -0,0 +1,32 @@ +# 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. +- Add automated smoke tests for the new SQLite/Neo4j drivers to ensure regressions are caught early. + +## 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 +- Document usage patterns for each graph backend (memory/sqlite/memgraph/neo4j) and provide Makefile shortcuts for switching. +- 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. +- Add onboarding notes for the REST/gRPC service layers with sample payloads and curl/grpcurl snippets. + +## Future Enhancements +- Export telemetry to Prometheus/OpenTelemetry and wire alerts/dashboards around ingestion and maintenance. +- Explore streaming ingestion pipelines (queues, webhooks) for near-real-time updates. +- Investigate lightweight web UI tooling for inspecting memories, triplets, and telemetry snapshots. diff --git a/SOT.md b/SOT.md new file mode 100644 index 0000000..3f3ab1b --- /dev/null +++ b/SOT.md @@ -0,0 +1,125 @@ +# 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 plus REST/gRPC service layers +├── 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 in-memory, SQLite, Memgraph, and Neo4j implementations +├── models/ # Entity and predicate registries shared across the pipeline +├── pipeline/ # Extract, preprocess, compression, storage, consolidation, expiry stages with telemetry hooks +├── 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). +- `docker-compose.yml`: Provisions Memgraph, Redis, and a Celery worker for local orchestration. +- `.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 the active graph backend (`GRAPH_BACKEND`), Memgraph (`MEMGRAPH_URI`, `MEMGRAPH_USERNAME`, + `MEMGRAPH_PASSWORD`), SQLite (`SQLITE_PATH`), optional Neo4j (`NEO4J_URI`, `NEO4J_USERNAME`, `NEO4J_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 the configured backend (memory, SQLite, Memgraph, Neo4j) lazily 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). +- `InMemoryGraphDriver` and `SQLiteGraphDriver` power local development/testing without external services. +- `MemgraphDriver` wraps `mgclient`, handles URI parsing, executes Cypher statements, and exposes a Python-side vector search + fallback when database-native similarity is unavailable. +- `Neo4jGraphDriver` mirrors the Memgraph contract using the official driver (optional dependency). +- `factory.py` exposes helpers (`create_graph_driver`, `graph_driver_factory`) to instantiate backends based on configuration. + +## 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 and records telemetry counters/durations. +4. **Store (`store.py`)** – Persists memories and triplets using the configured `GraphDriver`, registering predicates as needed and emitting observability events. +5. **Consolidate & Expire (`consolidate.py`, `expire.py`)** – Maintenance utilities triggered by Celery tasks to group memories + and remove stale entries. + +## Service Layers (`meshmind/api`) +- `memory_manager.py`: CRUD façade over the active graph driver. +- `service.py`: Pydantic payloads and orchestration helpers shared by REST/gRPC surfaces. +- `rest.py`: `create_app` returns a FastAPI application when available or a `RestAPIStub` for tests. +- `grpc.py`: `GrpcServiceStub` plus simple request/response dataclasses mirroring planned RPCs. + +## 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, + emit observability events, and 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`, in-memory drivers, service stubs) and + pure-Python doubles, 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` (when using Memgraph). +- Optional but supported: `neo4j`, `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 configured backend: in-memory/SQLite need no services; Memgraph requires a running instance reachable via `settings.MEMGRAPH_URI` and `pymgclient`; Neo4j requires the official driver and credentials. +- 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`. +- Docker Compose can be used to run Memgraph/Redis/Celery locally; ensure container tooling is available or provision services manually. diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..fdc8c3a --- /dev/null +++ b/TODO.md @@ -0,0 +1,24 @@ +# 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. +- [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. + - [x] Document or provision local Memgraph and Redis services (e.g., via docker-compose) for onboarding. + - [x] Abstract `GraphDriver` to support alternative storage backends (Neo4j, in-memory, SQLite prototype). + - [x] Add service interfaces (REST/gRPC) for ingestion and retrieval. + - [x] Introduce observability (logging, metrics) for ingestion and maintenance pipelines. + - [x] Promote NEW_README.md, archive legacy README, and maintain SOT diagrams and maps. + +## Later + +- [ ] Harden Celery maintenance tasks to initialize drivers lazily and persist consolidation results. +- [ ] Replace constant importance scoring with a data-driven or LLM-assisted heuristic. +- [ ] Create a fake memgraph driver for testing purposes. +- [ ] Create fake REDIS and Embedding model drivers for testing purposes. diff --git a/docker-compose.yml b/docker-compose.yml index 96be0b4..1c273da 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,15 +1,43 @@ version: "3.8" + services: memgraph: image: memgraph/memgraph-platform:latest + container_name: meshmind-memgraph ports: - "7687:7687" - "3000:3000" + healthcheck: + test: ["CMD", "bash", "-c", "cypher-shell --version || exit 1"] + interval: 10s + timeout: 5s + retries: 5 + volumes: + - memgraph-data:/var/lib/memgraph + redis: image: redis:7-alpine + container_name: meshmind-redis + ports: + - "6379:6379" + command: redis-server --save 60 1000 --loglevel warning + volumes: + - redis-data:/data + worker: build: . command: celery -A meshmind.tasks.celery_app worker -B -l info depends_on: - - memgraph - - redis \ No newline at end of file + memgraph: + condition: service_started + redis: + condition: service_started + environment: + MEMGRAPH_URI: bolt://memgraph:7687 + REDIS_URL: redis://redis:6379/0 + volumes: + - .:/app + +volumes: + memgraph-data: + redis-data: 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/grpc.py b/meshmind/api/grpc.py new file mode 100644 index 0000000..ba34a22 --- /dev/null +++ b/meshmind/api/grpc.py @@ -0,0 +1,80 @@ +"""gRPC-style adapters for MeshMind.""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Iterable, List + +from meshmind.api.service import MemoryPayload, MemoryService, SearchPayload, TripletPayload + + +@dataclass +class IngestMemoriesRequest: + memories: List[dict] = field(default_factory=list) + + +@dataclass +class IngestMemoriesResponse: + uuids: List[str] = field(default_factory=list) + + +@dataclass +class IngestTripletsRequest: + triplets: List[dict] = field(default_factory=list) + + +@dataclass +class IngestTripletsResponse: + stored: int = 0 + + +@dataclass +class SearchRequest: + query: str + namespace: str | None = None + top_k: int = 10 + encoder: str | None = None + rerank_k: int | None = None + + +@dataclass +class SearchResponse: + results: List[dict] = field(default_factory=list) + + +class GrpcServiceStub: + """Simple callable object that mirrors gRPC service behaviour for tests.""" + + def __init__(self, service: MemoryService) -> None: + self.service = service + + def IngestMemories(self, request: IngestMemoriesRequest) -> IngestMemoriesResponse: # noqa: N802 + payloads = [MemoryPayload(**item) for item in request.memories] + uuids = self.service.ingest_memories(payloads) + return IngestMemoriesResponse(uuids=uuids) + + def IngestTriplets(self, request: IngestTripletsRequest) -> IngestTripletsResponse: # noqa: N802 + payloads = [TripletPayload(**item) for item in request.triplets] + stored = self.service.ingest_triplets(payloads) + return IngestTripletsResponse(stored=stored) + + def Search(self, request: SearchRequest) -> SearchResponse: # noqa: N802 + payload = SearchPayload( + query=request.query, + namespace=request.namespace, + top_k=request.top_k, + encoder=request.encoder, + rerank_k=request.rerank_k, + ) + results = self.service.search(payload) + return SearchResponse(results=[mem.dict() for mem in results]) + + +__all__ = [ + "GrpcServiceStub", + "IngestMemoriesRequest", + "IngestMemoriesResponse", + "IngestTripletsRequest", + "IngestTripletsResponse", + "SearchRequest", + "SearchResponse", +] diff --git a/meshmind/api/memory_manager.py b/meshmind/api/memory_manager.py index 935d270..0ff023b 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: @@ -52,44 +63,66 @@ def get_memory(self, memory_id: UUID) -> Optional[Any]: :param memory_id: UUID of the memory to retrieve. :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) - if not records: + payload = self.driver.get_entity(str(memory_id)) + if not payload: return None - # Extract node properties - record = records[0] - data = record.get('m', record) try: - return Memory(**data) + return Memory(**payload) 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. :param namespace: If provided, only return memories in this namespace. :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} - else: - cypher = "MATCH (m) RETURN m" - params = {} - records = self.driver.find(cypher, params) - result: List[Any] = [] - for record in records: - data = record.get('m', record) + records = self.driver.list_entities(namespace) + result: List[Memory] = [] + for data in records: 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/api/rest.py b/meshmind/api/rest.py new file mode 100644 index 0000000..e22cfc7 --- /dev/null +++ b/meshmind/api/rest.py @@ -0,0 +1,88 @@ +"""REST adapters for the :mod:`meshmind` service layer.""" +from __future__ import annotations + +from typing import Any, Dict, Iterable + +from meshmind.api.service import MemoryPayload, MemoryService, SearchPayload, TripletPayload + + +class RestAPIStub: + """Fallback handler that emulates REST routes without FastAPI.""" + + def __init__(self, service: MemoryService) -> None: + self.service = service + + def dispatch(self, method: str, path: str, payload: Dict[str, Any] | None = None) -> Dict[str, Any]: + method = method.upper() + payload = payload or {} + if method == "POST" and path == "/memories": + memories = [MemoryPayload(**item) for item in payload.get("memories", [])] + uuids = self.service.ingest_memories(memories) + return {"uuids": uuids} + if method == "POST" and path == "/triplets": + triplets = [TripletPayload(**item) for item in payload.get("triplets", [])] + count = self.service.ingest_triplets(triplets) + return {"stored": count} + if method == "POST" and path == "/search": + request = SearchPayload(**payload) + results = self.service.search(request) + return {"results": [mem.dict() for mem in results]} + if method == "GET" and path == "/memories": + namespace = payload.get("namespace") + memories = self.service.list_memories(namespace) + return {"memories": [mem.dict() for mem in memories]} + if method == "GET" and path == "/triplets": + namespace = payload.get("namespace") + triplets = self.service.list_triplets(namespace) + return {"triplets": [triplet.dict() for triplet in triplets]} + raise ValueError(f"Unsupported route {method} {path}") + + +def create_app(service: MemoryService) -> Any: + """Create a FastAPI application if FastAPI is installed, otherwise return a stub.""" + + try: # pragma: no cover - optional dependency path + from fastapi import FastAPI, HTTPException + except ImportError: # pragma: no cover - executed in tests without fastapi + return RestAPIStub(service) + + app = FastAPI(title="MeshMind API") + + @app.post("/memories") + def create_memories(payload: Dict[str, Iterable[Dict[str, Any]]]): + try: + items = [MemoryPayload(**item) for item in payload.get("memories", [])] + except Exception as exc: # pragma: no cover - FastAPI handles validation + raise HTTPException(status_code=400, detail=str(exc)) + uuids = service.ingest_memories(items) + return {"uuids": uuids} + + @app.post("/triplets") + def create_triplets(payload: Dict[str, Iterable[Dict[str, Any]]]): + try: + items = [TripletPayload(**item) for item in payload.get("triplets", [])] + except Exception as exc: # pragma: no cover + raise HTTPException(status_code=400, detail=str(exc)) + stored = service.ingest_triplets(items) + return {"stored": stored} + + @app.post("/search") + def search(payload: Dict[str, Any]): + try: + request = SearchPayload(**payload) + except Exception as exc: # pragma: no cover + raise HTTPException(status_code=400, detail=str(exc)) + results = service.search(request) + return {"results": [mem.dict() for mem in results]} + + @app.get("/memories") + def list_memories(namespace: str | None = None): + memories = service.list_memories(namespace) + return {"memories": [mem.dict() for mem in memories]} + + @app.get("/triplets") + def list_triplets(namespace: str | None = None): + triplets = service.list_triplets(namespace) + return {"triplets": [triplet.dict() for triplet in triplets]} + + return app diff --git a/meshmind/api/service.py b/meshmind/api/service.py new file mode 100644 index 0000000..89d783c --- /dev/null +++ b/meshmind/api/service.py @@ -0,0 +1,101 @@ +"""Service layer abstractions for REST and gRPC adapters.""" +from __future__ import annotations + +from typing import Iterable, List, Sequence + +from pydantic import BaseModel, Field + +from meshmind.api.memory_manager import MemoryManager +from meshmind.core.types import Memory, SearchConfig, Triplet +from meshmind.retrieval import search as retrieval_search + + +class MemoryPayload(BaseModel): + """Serializable payload for creating or updating a memory.""" + + uuid: str | None = None + namespace: str + name: str + entity_label: str = "Memory" + embedding: List[float] | None = None + metadata: dict[str, object] = Field(default_factory=dict) + reference_time: str | None = None + importance: float | None = None + ttl_seconds: int | None = None + + def to_memory(self) -> Memory: + payload = self.model_dump(exclude_none=True) + if self.uuid: + payload["uuid"] = self.uuid + return Memory(**payload) + + +class TripletPayload(BaseModel): + subject: str + predicate: str + object: str + namespace: str + entity_label: str = "Relation" + metadata: dict[str, object] = Field(default_factory=dict) + reference_time: str | None = None + + def to_triplet(self) -> Triplet: + return Triplet(**self.model_dump(exclude_none=True)) + + +class SearchPayload(BaseModel): + query: str + namespace: str | None = None + top_k: int = 10 + encoder: str | None = None + rerank_k: int | None = None + + def to_config(self) -> SearchConfig: + config = SearchConfig(top_k=self.top_k) + if self.encoder: + config.encoder = self.encoder + if self.rerank_k is not None: + config.rerank_k = self.rerank_k + return config + + +class MemoryService: + """Business logic for ingestion, retrieval, and triplet persistence.""" + + def __init__(self, manager: MemoryManager) -> None: + self.manager = manager + + # ------------------------------------------------------------------ + # Ingestion + # ------------------------------------------------------------------ + def ingest_memories(self, payloads: Sequence[MemoryPayload]) -> List[str]: + uuids: List[str] = [] + for payload in payloads: + memory = payload.to_memory() + self.manager.add_memory(memory) + uuids.append(str(memory.uuid)) + return uuids + + def ingest_triplets(self, payloads: Iterable[TripletPayload]) -> int: + count = 0 + for payload in payloads: + self.manager.add_triplet(payload.to_triplet()) + count += 1 + return count + + # ------------------------------------------------------------------ + # Retrieval + # ------------------------------------------------------------------ + def search(self, request: SearchPayload) -> List[Memory]: + memories = self.manager.list_memories(namespace=request.namespace) + config = request.to_config() + return retrieval_search(request.query, memories, config=config) + + # ------------------------------------------------------------------ + # CRUD proxies + # ------------------------------------------------------------------ + def list_memories(self, namespace: str | None = None) -> List[Memory]: + return self.manager.list_memories(namespace) + + def list_triplets(self, namespace: str | None = None) -> List[Triplet]: + return self.manager.list_triplets(namespace) 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..ca70a9b 100644 --- a/meshmind/client.py +++ b/meshmind/client.py @@ -1,80 +1,272 @@ -""" -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.factory import graph_driver_factory as make_graph_driver_factory +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 = make_graph_driver_factory() + + 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..da78b2e 100644 --- a/meshmind/core/config.py +++ b/meshmind/core/config.py @@ -13,20 +13,73 @@ class Settings: """Application settings loaded from environment variables.""" + GRAPH_BACKEND: str = os.getenv("GRAPH_BACKEND", "memory") MEMGRAPH_URI: str = os.getenv("MEMGRAPH_URI", "bolt://localhost:7687") MEMGRAPH_USERNAME: str = os.getenv("MEMGRAPH_USERNAME", "") MEMGRAPH_PASSWORD: str = os.getenv("MEMGRAPH_PASSWORD", "") + NEO4J_URI: str = os.getenv("NEO4J_URI", "bolt://localhost:7687") + NEO4J_USERNAME: str = os.getenv("NEO4J_USERNAME", "neo4j") + NEO4J_PASSWORD: str = os.getenv("NEO4J_PASSWORD", "") + SQLITE_PATH: str = os.getenv("SQLITE_PATH", ":memory:") REDIS_URL: str = os.getenv("REDIS_URL", "redis://localhost:6379/0") OPENAI_API_KEY: str = os.getenv("OPENAI_API_KEY", "") EMBEDDING_MODEL: str = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small") + REQUIRED_GROUPS = { + "memgraph": ("MEMGRAPH_URI",), + "neo4j": ("NEO4J_URI",), + "openai": ("OPENAI_API_KEY",), + "redis": ("REDIS_URL",), + } + def __repr__(self) -> str: return ( - f"Settings(MEMGRAPH_URI={self.MEMGRAPH_URI}, " + f"Settings(GRAPH_BACKEND={self.GRAPH_BACKEND}, " + f"MEMGRAPH_URI={self.MEMGRAPH_URI}, " f"MEMGRAPH_USERNAME={self.MEMGRAPH_USERNAME}, " + f"NEO4J_URI={self.NEO4J_URI}, " f"REDIS_URL={self.REDIS_URL}, " 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(): + if group == "memgraph" and self.GRAPH_BACKEND != "memgraph": + continue + if group == "neo4j" and self.GRAPH_BACKEND != "neo4j": + continue + 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), + "NEO4J_URI": self.NEO4J_URI, + "NEO4J_USERNAME": self.NEO4J_USERNAME, + "NEO4J_PASSWORD": self._mask(self.NEO4J_PASSWORD), + "GRAPH_BACKEND": self.GRAPH_BACKEND, + "SQLITE_PATH": self.SQLITE_PATH, + "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/observability.py b/meshmind/core/observability.py new file mode 100644 index 0000000..6c36fe1 --- /dev/null +++ b/meshmind/core/observability.py @@ -0,0 +1,80 @@ +"""Shared logging and metrics utilities for MeshMind pipelines.""" +from __future__ import annotations + +import logging +from collections import defaultdict +from contextlib import contextmanager +from time import perf_counter +from typing import Any, Dict, Iterable + +_LOGGER_NAME = "meshmind" +logger = logging.getLogger(_LOGGER_NAME) +if not logger.handlers: # pragma: no cover - avoid duplicate handlers in tests + handler = logging.StreamHandler() + formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s") + handler.setFormatter(formatter) + logger.addHandler(handler) +logger.setLevel(logging.INFO) + + +class Telemetry: + """Lightweight in-memory metrics collector.""" + + def __init__(self) -> None: + self._counters: Dict[str, int] = defaultdict(int) + self._durations: Dict[str, list[float]] = defaultdict(list) + self._gauges: Dict[str, float] = {} + + # ------------------------------------------------------------------ + # Counter helpers + # ------------------------------------------------------------------ + def increment(self, metric: str, value: int = 1) -> None: + self._counters[metric] += value + + def gauge(self, metric: str, value: float) -> None: + self._gauges[metric] = value + + def observe(self, metric: str, value: float) -> None: + self._durations[metric].append(value) + + @contextmanager + def track_duration(self, metric: str): + start = perf_counter() + try: + yield + finally: + elapsed = perf_counter() - start + self.observe(metric, elapsed) + + def extend_counter(self, metric: str, values: Iterable[Any]) -> None: + count = sum(1 for _ in values) + self.increment(metric, count) + + # ------------------------------------------------------------------ + # Snapshot helpers + # ------------------------------------------------------------------ + def snapshot(self) -> Dict[str, Any]: + return { + "counters": dict(self._counters), + "durations": {k: list(v) for k, v in self._durations.items()}, + "gauges": dict(self._gauges), + } + + def reset(self) -> None: + self._counters.clear() + self._durations.clear() + self._gauges.clear() + + +telemetry = Telemetry() + + +def log_event(event: str, **fields: Any) -> None: + """Emit a structured log entry and update a counter for the event.""" + + telemetry.increment(f"events.{event}") + if fields: + formatted = " ".join(f"{key}={value}" for key, value in fields.items()) + logger.info("event=%s %s", event, formatted) + else: + logger.info("event=%s", event) 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/__init__.py b/meshmind/db/__init__.py index e69de29..e5c249e 100644 --- a/meshmind/db/__init__.py +++ b/meshmind/db/__init__.py @@ -0,0 +1,14 @@ +"""Graph driver implementations exposed for convenience.""" +from .base_driver import GraphDriver +from .in_memory_driver import InMemoryGraphDriver +from .memgraph_driver import MemgraphDriver +from .neo4j_driver import Neo4jGraphDriver +from .sqlite_driver import SQLiteGraphDriver + +__all__ = [ + "GraphDriver", + "InMemoryGraphDriver", + "MemgraphDriver", + "Neo4jGraphDriver", + "SQLiteGraphDriver", +] diff --git a/meshmind/db/base_driver.py b/meshmind/db/base_driver.py index 1e1366c..67414eb 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 @@ -22,7 +24,29 @@ def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: """Execute a Cypher query and return results.""" raise NotImplementedError + @abstractmethod + def get_entity(self, uid: str) -> Optional[Dict[str, Any]]: + """Return a single entity by UUID, if it exists.""" + raise NotImplementedError + + @abstractmethod + def list_entities(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + """Return entities, optionally filtered by namespace.""" + raise NotImplementedError + @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/factory.py b/meshmind/db/factory.py new file mode 100644 index 0000000..e7d18cd --- /dev/null +++ b/meshmind/db/factory.py @@ -0,0 +1,66 @@ +"""Factory helpers for constructing :class:`GraphDriver` instances.""" +from __future__ import annotations + +from typing import Callable, Dict, Type + +from meshmind.core.config import settings +from meshmind.db.base_driver import GraphDriver +from meshmind.db.in_memory_driver import InMemoryGraphDriver +from meshmind.db.memgraph_driver import MemgraphDriver +from meshmind.db.neo4j_driver import Neo4jGraphDriver +from meshmind.db.sqlite_driver import SQLiteGraphDriver + + +def _normalize_backend(name: str) -> str: + return name.replace("-", "_").lower() + + +def available_backends() -> Dict[str, Type[GraphDriver]]: + """Return the mapping of backend names to driver classes.""" + + return { + "memory": InMemoryGraphDriver, + "in_memory": InMemoryGraphDriver, + "inmemory": InMemoryGraphDriver, + "sqlite": SQLiteGraphDriver, + "memgraph": MemgraphDriver, + "neo4j": Neo4jGraphDriver, + } + + +def create_graph_driver(backend: str | None = None, **kwargs) -> GraphDriver: + """Instantiate a :class:`GraphDriver` for the requested backend.""" + + backend_name = _normalize_backend(backend or settings.GRAPH_BACKEND) + drivers = available_backends() + if backend_name not in drivers: + raise ValueError(f"Unsupported graph backend '{backend_name}'") + + driver_cls: Type[GraphDriver] = drivers[backend_name] + if driver_cls is InMemoryGraphDriver: + return driver_cls() + if driver_cls is SQLiteGraphDriver: + path = kwargs.get("path") or settings.SQLITE_PATH + return driver_cls(path) + if driver_cls is MemgraphDriver: + return driver_cls( + settings.MEMGRAPH_URI, + settings.MEMGRAPH_USERNAME, + settings.MEMGRAPH_PASSWORD, + ) + if driver_cls is Neo4jGraphDriver: + return driver_cls( + settings.NEO4J_URI, + settings.NEO4J_USERNAME, + settings.NEO4J_PASSWORD, + ) + return driver_cls(**kwargs) + + +def graph_driver_factory(backend: str | None = None, **kwargs) -> Callable[[], GraphDriver]: + """Return a callable that lazily constructs the configured driver.""" + + def _factory() -> GraphDriver: + return create_graph_driver(backend=backend, **kwargs) + + return _factory diff --git a/meshmind/db/in_memory_driver.py b/meshmind/db/in_memory_driver.py new file mode 100644 index 0000000..89bfe96 --- /dev/null +++ b/meshmind/db/in_memory_driver.py @@ -0,0 +1,126 @@ +"""In-memory implementation of :class:`GraphDriver` for tests and local development.""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple +from uuid import UUID, uuid4 + +from meshmind.db.base_driver import GraphDriver + + +class InMemoryGraphDriver(GraphDriver): + """A lightweight graph driver that stores entities and triplets in dictionaries.""" + + def __init__(self) -> None: + self._nodes: Dict[str, Dict[str, Any]] = {} + self._labels: Dict[str, str] = {} + self._triplets: Dict[Tuple[str, str, str], Dict[str, Any]] = {} + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + def _ensure_uuid(self, props: Dict[str, Any]) -> str: + uid = props.get("uuid") + if isinstance(uid, UUID): + props["uuid"] = str(uid) + return str(uid) + if isinstance(uid, str): + return uid + new_uid = str(uuid4()) + props["uuid"] = new_uid + return new_uid + + # ------------------------------------------------------------------ + # GraphDriver API + # ------------------------------------------------------------------ + def upsert_entity(self, label: str, name: str, props: Dict[str, Any]) -> None: + props = dict(props) + uid = self._ensure_uuid(props) + props.setdefault("name", name) + props.setdefault("entity_label", label) + self._nodes[uid] = props + self._labels[uid] = label + + def upsert_edge(self, subj: str, pred: str, obj: str, props: Dict[str, Any]) -> None: + key = (str(subj), pred, str(obj)) + payload = dict(props) + payload.setdefault("subject", str(subj)) + payload.setdefault("predicate", pred) + payload.setdefault("object", str(obj)) + self._triplets[key] = payload + + def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: + cypher = cypher.lower().strip() + if "where m.uuid" in cypher: + uid = str(params.get("uuid", "")) + node = self._nodes.get(uid) + if node is None: + return [] + return [{"m": dict(node)}] + if "where m.namespace" in cypher: + namespace = params.get("namespace") + results = [ + {"m": dict(node)} + for node in self._nodes.values() + if node.get("namespace") == namespace + ] + return results + if cypher.startswith("match (m) return m"): + return [{"m": dict(node)} for node in self._nodes.values()] + if cypher.startswith("match (a)-[r"): + namespace = params.get("namespace") + results: List[Dict[str, Any]] = [] + for payload in self._triplets.values(): + if namespace and payload.get("namespace") != namespace: + continue + record = { + "subject": payload.get("subject"), + "predicate": payload.get("predicate"), + "object": payload.get("object"), + "namespace": payload.get("namespace"), + "metadata": payload.get("metadata"), + "reference_time": payload.get("reference_time"), + } + results.append(record) + return results + return [] + + def get_entity(self, uid: str) -> Optional[Dict[str, Any]]: + node = self._nodes.get(str(uid)) + return dict(node) if node else None + + def list_entities(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + if namespace is None: + return [dict(node) for node in self._nodes.values()] + return [ + dict(node) + for node in self._nodes.values() + if node.get("namespace") == namespace + ] + + def delete(self, uuid: UUID) -> None: + uid = str(uuid) + self._nodes.pop(uid, None) + self._labels.pop(uid, None) + to_delete = [key for key in self._triplets if key[0] == uid or key[2] == uid] + for key in to_delete: + self._triplets.pop(key, None) + + def delete_triplet(self, subj: str, pred: str, obj: str) -> None: + self._triplets.pop((str(subj), pred, str(obj)), None) + + def list_triplets(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + results: List[Dict[str, Any]] = [] + for payload in self._triplets.values(): + if namespace and payload.get("namespace") != namespace: + continue + results.append( + { + "subject": payload.get("subject"), + "predicate": payload.get("predicate"), + "object": payload.get("object"), + "namespace": payload.get("namespace"), + "metadata": payload.get("metadata"), + "reference_time": payload.get("reference_time"), + } + ) + return results diff --git a/meshmind/db/memgraph_driver.py b/meshmind/db/memgraph_driver.py index f4e8c4e..1fd6961 100644 --- a/meshmind/db/memgraph_driver.py +++ b/meshmind/db/memgraph_driver.py @@ -1,110 +1,173 @@ -"""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("`", "") + + @staticmethod + def _normalize_node(node: Any) -> Dict[str, Any]: + if hasattr(node, "properties"): + try: + return dict(node.properties) # type: ignore[attr-defined] + except Exception: + pass + if isinstance(node, dict): + return dict(node) + if hasattr(node, "_properties"): + return dict(getattr(node, "_properties")) + return {k: v for k, v in getattr(node, "__dict__", {}).items() if not k.startswith("_")} + + # ------------------------------------------------------------------ + # 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 get_entity(self, uid: str) -> Optional[Dict[str, Any]]: + records = self.find( + "MATCH (m) WHERE m.uuid = $uuid RETURN m", + {"uuid": str(uid)}, + ) + if not records: + return None + node = records[0].get("m", records[0]) + return self._normalize_node(node) + + def list_entities(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + if namespace: + cypher = "MATCH (m) WHERE m.namespace = $namespace RETURN m" + params = {"namespace": namespace} + else: + cypher = "MATCH (m) RETURN m" + params = {} + records = self.find(cypher, params) + entities: List[Dict[str, Any]] = [] + for record in records: + node = record.get("m", record) + entities.append(self._normalize_node(node)) + return entities + 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/db/neo4j_driver.py b/meshmind/db/neo4j_driver.py new file mode 100644 index 0000000..1d76b73 --- /dev/null +++ b/meshmind/db/neo4j_driver.py @@ -0,0 +1,108 @@ +"""Neo4j implementation of :class:`GraphDriver` using the official driver.""" +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from meshmind.db.base_driver import GraphDriver + +try: # pragma: no cover - optional dependency + from neo4j import GraphDatabase # type: ignore +except ImportError: # pragma: no cover - optional dependency + GraphDatabase = None # type: ignore + + +class Neo4jGraphDriver(GraphDriver): + """GraphDriver backed by Neo4j via the ``neo4j`` Python driver.""" + + def __init__(self, uri: str, username: str = "neo4j", password: str = "") -> None: + if GraphDatabase is None: + raise ImportError("neo4j driver is required for Neo4jGraphDriver") + auth = None + if username or password: + auth = (username or None, password or None) + self._driver = GraphDatabase.driver(uri, auth=auth) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + def _run(self, cypher: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: + params = params or {} + with self._driver.session() as session: # type: ignore[attr-defined] + result = session.run(cypher, **params) + records = [] + for record in result: + records.append(record.data()) + return records + + @staticmethod + def _normalize_node(node: Any) -> Dict[str, Any]: + if hasattr(node, "_properties"): + return dict(node._properties) # type: ignore[attr-defined] + if isinstance(node, dict): + return dict(node) + return {k: v for k, v in getattr(node, "__dict__", {}).items() if not k.startswith("_")} + + # ------------------------------------------------------------------ + # GraphDriver API + # ------------------------------------------------------------------ + def upsert_entity(self, label: str, name: str, props: Dict[str, Any]) -> None: + cypher = ( + f"MERGE (n:{label} {{uuid: $uuid}})\n" + "SET n += $props" + ) + params = {"uuid": str(props.get("uuid")), "props": props} + self._run(cypher, params) + + def upsert_edge(self, subj: str, pred: str, obj: str, props: Dict[str, Any]) -> None: + predicate = pred.replace("`", "") + cypher = ( + "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} + self._run(cypher, params) + + def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: + return self._run(cypher, params) + + def get_entity(self, uid: str) -> Optional[Dict[str, Any]]: + records = self.find("MATCH (m) WHERE m.uuid = $uuid RETURN m", {"uuid": str(uid)}) + if not records: + return None + node = records[0].get("m", records[0]) + return self._normalize_node(node) + + def list_entities(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + if namespace: + cypher = "MATCH (m) WHERE m.namespace = $namespace RETURN m" + params = {"namespace": namespace} + else: + cypher = "MATCH (m) RETURN m" + params = {} + records = self.find(cypher, params) + return [self._normalize_node(rec.get("m", rec)) for rec in records] + + def delete(self, uuid: Any) -> None: + self._run("MATCH (m {uuid: $uuid}) DETACH DELETE m", {"uuid": str(uuid)}) + + def delete_triplet(self, subj: str, pred: str, obj: str) -> None: + predicate = pred.replace("`", "") + cypher = ( + f"MATCH (a {{uuid: $subj}})-[r:`{predicate}`]->(b {{uuid: $obj}})" + " DELETE r" + ) + self._run(cypher, {"subj": str(subj), "obj": str(obj)}) + + 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.find(cypher, params) + + def close(self) -> None: + self._driver.close() diff --git a/meshmind/db/sqlite_driver.py b/meshmind/db/sqlite_driver.py new file mode 100644 index 0000000..3c60c7f --- /dev/null +++ b/meshmind/db/sqlite_driver.py @@ -0,0 +1,208 @@ +"""SQLite implementation of :class:`GraphDriver` for lightweight persistence.""" +from __future__ import annotations + +import json +import sqlite3 +from pathlib import Path +from typing import Any, Dict, List, Optional + +from meshmind.db.base_driver import GraphDriver + + +class SQLiteGraphDriver(GraphDriver): + """GraphDriver backed by SQLite tables using simple JSON columns.""" + + def __init__(self, path: str | Path = ":memory:") -> None: + self._path = str(path) + self._conn = sqlite3.connect(self._path) + self._conn.row_factory = sqlite3.Row + self._ensure_schema() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + def _ensure_schema(self) -> None: + cur = self._conn.cursor() + cur.execute( + """ + CREATE TABLE IF NOT EXISTS entities ( + uuid TEXT PRIMARY KEY, + label TEXT NOT NULL, + name TEXT NOT NULL, + namespace TEXT, + props TEXT NOT NULL + ) + """ + ) + cur.execute( + """ + CREATE TABLE IF NOT EXISTS triplets ( + subject TEXT NOT NULL, + predicate TEXT NOT NULL, + object TEXT NOT NULL, + namespace TEXT, + metadata TEXT, + reference_time TEXT, + PRIMARY KEY (subject, predicate, object) + ) + """ + ) + self._conn.commit() + + def _row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]: + payload = dict(row) + if "props" in payload and payload["props"]: + payload.update(json.loads(payload["props"])) + payload.pop("props", None) + if "metadata" in payload and payload["metadata"]: + payload["metadata"] = json.loads(payload["metadata"]) + return payload + + # ------------------------------------------------------------------ + # GraphDriver API + # ------------------------------------------------------------------ + def upsert_entity(self, label: str, name: str, props: Dict[str, Any]) -> None: + payload = dict(props) + uid = str(payload.get("uuid")) + if not uid: + raise ValueError("Memory props must include a UUID for SQLiteGraphDriver") + payload.setdefault("entity_label", label) + payload.setdefault("name", name) + namespace = payload.get("namespace") + cur = self._conn.cursor() + cur.execute( + """ + INSERT INTO entities (uuid, label, name, namespace, props) + VALUES (:uuid, :label, :name, :namespace, :props) + ON CONFLICT(uuid) DO UPDATE SET + label=excluded.label, + name=excluded.name, + namespace=excluded.namespace, + props=excluded.props + """, + { + "uuid": uid, + "label": payload.get("entity_label", label), + "name": payload.get("name", name), + "namespace": namespace, + "props": json.dumps(payload), + }, + ) + self._conn.commit() + + def upsert_edge(self, subj: str, pred: str, obj: str, props: Dict[str, Any]) -> None: + payload = dict(props) + payload.setdefault("subject", subj) + payload.setdefault("predicate", pred) + payload.setdefault("object", obj) + metadata = payload.get("metadata") or {} + cur = self._conn.cursor() + cur.execute( + """ + INSERT INTO triplets (subject, predicate, object, namespace, metadata, reference_time) + VALUES (:subject, :predicate, :object, :namespace, :metadata, :reference_time) + ON CONFLICT(subject, predicate, object) DO UPDATE SET + namespace=excluded.namespace, + metadata=excluded.metadata, + reference_time=excluded.reference_time + """, + { + "subject": payload["subject"], + "predicate": payload["predicate"], + "object": payload["object"], + "namespace": payload.get("namespace"), + "metadata": json.dumps(metadata), + "reference_time": payload.get("reference_time"), + }, + ) + self._conn.commit() + + def find(self, cypher: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: + # Provide compatibility for simple MATCH queries used by MemoryManager. + cypher_lower = cypher.lower().strip() + cur = self._conn.cursor() + if "where m.uuid" in cypher_lower: + cur.execute("SELECT * FROM entities WHERE uuid = :uuid", {"uuid": params.get("uuid")}) + row = cur.fetchone() + if not row: + return [] + return [{"m": self._row_to_dict(row)}] + if "where m.namespace" in cypher_lower: + cur.execute( + "SELECT * FROM entities WHERE namespace = :namespace", + {"namespace": params.get("namespace")}, + ) + rows = cur.fetchall() + return [{"m": self._row_to_dict(row)} for row in rows] + if cypher_lower.startswith("match (m) return m"): + cur.execute("SELECT * FROM entities") + rows = cur.fetchall() + return [{"m": self._row_to_dict(row)} for row in rows] + if cypher_lower.startswith("match (a)-[r"): + namespace = params.get("namespace") + if namespace: + cur.execute( + "SELECT * FROM triplets WHERE namespace = :namespace", + {"namespace": namespace}, + ) + else: + cur.execute("SELECT * FROM triplets") + rows = cur.fetchall() + return [dict(row) for row in rows] + return [] + + def get_entity(self, uid: str) -> Optional[Dict[str, Any]]: + cur = self._conn.cursor() + cur.execute("SELECT * FROM entities WHERE uuid = :uuid", {"uuid": uid}) + row = cur.fetchone() + return self._row_to_dict(row) if row else None + + def list_entities(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + cur = self._conn.cursor() + if namespace: + cur.execute( + "SELECT * FROM entities WHERE namespace = :namespace", + {"namespace": namespace}, + ) + else: + cur.execute("SELECT * FROM entities") + rows = cur.fetchall() + return [self._row_to_dict(row) for row in rows] + + def delete(self, uuid: Any) -> None: + cur = self._conn.cursor() + cur.execute("DELETE FROM entities WHERE uuid = :uuid", {"uuid": str(uuid)}) + cur.execute( + "DELETE FROM triplets WHERE subject = :uuid OR object = :uuid", + {"uuid": str(uuid)}, + ) + self._conn.commit() + + def delete_triplet(self, subj: str, pred: str, obj: str) -> None: + cur = self._conn.cursor() + cur.execute( + "DELETE FROM triplets WHERE subject = :subject AND predicate = :predicate AND object = :object", + {"subject": str(subj), "predicate": pred, "object": str(obj)}, + ) + self._conn.commit() + + def list_triplets(self, namespace: Optional[str] = None) -> List[Dict[str, Any]]: + cur = self._conn.cursor() + if namespace: + cur.execute( + "SELECT * FROM triplets WHERE namespace = :namespace", + {"namespace": namespace}, + ) + else: + cur.execute("SELECT * FROM triplets") + rows = cur.fetchall() + result = [] + for row in rows: + payload = dict(row) + metadata = payload.get("metadata") + payload["metadata"] = json.loads(metadata) if metadata else {} + result.append(payload) + return result + + def close(self) -> None: + self._conn.close() 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..c7eb995 100644 --- a/meshmind/pipeline/compress.py +++ b/meshmind/pipeline/compress.py @@ -1,13 +1,11 @@ -""" -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.observability import log_event, telemetry +from meshmind.core.types import Memory +from meshmind.core.utils import get_token_encoder def compress_memories( @@ -20,19 +18,29 @@ 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') + log_event("pipeline.compress.start", items=len(memories)) + encoder = get_token_encoder("o200k_base", optional=True) + if encoder is None: + telemetry.increment("pipeline.compress.skipped", len(memories)) + return memories compressed = [] - for mem in memories: - content = mem.metadata.get('content') - if not isinstance(content, str): - compressed.append(mem) - continue - tokens = encoder.encode(content) - if len(tokens) <= max_tokens: + modified = 0 + with telemetry.track_duration("pipeline.compress.duration"): + for mem in memories: + content = mem.metadata.get('content') + if not isinstance(content, str): + compressed.append(mem) + continue + tokens = encoder.encode(content) + if len(tokens) <= max_tokens: + compressed.append(mem) + continue + # Truncate tokens and decode back to string + truncated = encoder.decode(tokens[:max_tokens]) + mem.metadata['content'] = truncated compressed.append(mem) - continue - # Truncate tokens and decode back to string - truncated = encoder.decode(tokens[:max_tokens]) - mem.metadata['content'] = truncated - compressed.append(mem) - return compressed \ No newline at end of file + modified += 1 + telemetry.increment("pipeline.compress.processed", len(memories)) + telemetry.increment("pipeline.compress.modified", modified) + log_event("pipeline.compress.complete", modified=modified) + return compressed diff --git a/meshmind/pipeline/extract.py b/meshmind/pipeline/extract.py index 613073b..24ba444 100644 --- a/meshmind/pipeline/extract.py +++ b/meshmind/pipeline/extract.py @@ -1,11 +1,13 @@ -from typing import Any, List, Type +from typing import Any, List, Sequence, Type + +from meshmind.core.observability import log_event, telemetry 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 +28,9 @@ 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 + + log_event("pipeline.extract.start", segments=len(content)) # Initialize default LLM client if not provided if llm_client is None: @@ -46,6 +51,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,15 +66,16 @@ 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( - model="gpt-4.1-mini", - messages=messages, - functions=[function_spec], - function_call={"name": "extract_memories"}, - ) + with telemetry.track_duration("pipeline.extract.duration"): + response = llm_client.responses.create( + model="gpt-4.1-mini", + messages=messages, + functions=[function_spec], + function_call={"name": "extract_memories"}, + ) msg = response.choices[0].message # Parse function call arguments or direct JSON if msg.get("function_call"): @@ -99,4 +108,8 @@ def extract_memories( emb = encoder.encode([mem.name])[0] mem.embedding = emb memories.append(mem) - return memories \ No newline at end of file + + telemetry.increment("pipeline.extract.segments", len(content)) + telemetry.increment("pipeline.extract.memories", len(memories)) + log_event("pipeline.extract.complete", memories=len(memories)) + return memories diff --git a/meshmind/pipeline/store.py b/meshmind/pipeline/store.py index 3d0b07c..502298b 100644 --- a/meshmind/pipeline/store.py +++ b/meshmind/pipeline/store.py @@ -1,23 +1,71 @@ +"""Persistence helpers for storing memories and triplets.""" +from __future__ import annotations + from typing import Any, Iterable + +from pydantic import BaseModel + +from meshmind.core.observability import log_event, telemetry +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. - - :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 - 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 + """Persist a sequence of Memory objects into the graph database.""" + + registry = entity_registry or EntityRegistry + stored = 0 + with telemetry.track_duration("pipeline.store.memories.duration"): + for mem in memories: + 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) + stored += 1 + telemetry.increment("pipeline.store.memories.stored", stored) + log_event("pipeline.store.memories", count=stored) + + +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 + stored = 0 + with telemetry.track_duration("pipeline.store.triplets.duration"): + for triplet in triplets: + registry.add(triplet.predicate) + props = _props(triplet) + graph_driver.upsert_edge( + triplet.subject, + triplet.predicate, + triplet.object, + props, + ) + stored += 1 + telemetry.increment("pipeline.store.triplets.stored", stored) + log_event("pipeline.store.triplets", count=stored) 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..152013b 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 @@ -9,25 +11,30 @@ _CELERY_BEAT = False def crontab(*args, **kwargs): # type: ignore return None -from meshmind.tasks.celery_app import app -from meshmind.pipeline.expire import expire_memories -from meshmind.pipeline.consolidate import consolidate_memories -from meshmind.pipeline.compress import compress_memories from meshmind.api.memory_manager import MemoryManager -from meshmind.db.memgraph_driver import MemgraphDriver from meshmind.core.config import settings +from meshmind.core.observability import log_event, telemetry +from meshmind.db.factory import create_graph_driver +from meshmind.pipeline.compress import compress_memories +from meshmind.pipeline.consolidate import consolidate_memories +from meshmind.pipeline.expire import expire_memories +from meshmind.tasks.celery_app import app -# 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 = create_graph_driver() + 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,30 +57,46 @@ 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) + log_event("task.expire.start") + with telemetry.track_duration("task.expire.duration"): + results = expire_memories(manager) + telemetry.increment("task.expire.runs") + log_event("task.expire.complete", removed=len(results)) + return results @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() - consolidated = consolidate_memories(memories) - for mem in consolidated: - manager.update_memory(mem) + log_event("task.consolidate.start", memories=len(memories)) + with telemetry.track_duration("task.consolidate.duration"): + consolidated = consolidate_memories(memories) + for mem in consolidated: + manager.update_memory(mem) + telemetry.increment("task.consolidate.runs") + log_event("task.consolidate.complete", merged=len(consolidated)) return len(consolidated) @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 + log_event("task.compress.start", memories=len(memories)) + with telemetry.track_duration("task.compress.duration"): + compressed = compress_memories(memories) + for mem in compressed: + manager.update_memory(mem) + telemetry.increment("task.compress.runs") + log_event("task.compress.complete", updated=len(compressed)) + return len(compressed) diff --git a/meshmind/tests/conftest.py b/meshmind/tests/conftest.py new file mode 100644 index 0000000..e096878 --- /dev/null +++ b/meshmind/tests/conftest.py @@ -0,0 +1,51 @@ +import pytest + +try: # pragma: no cover - allow test suite to skip when dependencies are unavailable + import pydantic # noqa: F401 +except ModuleNotFoundError: # pragma: no cover + pytest.skip("pydantic is required for the MeshMind test suite", allow_module_level=True) + +from meshmind.api.memory_manager import MemoryManager +from meshmind.api.service import MemoryService +from meshmind.core.embeddings import EncoderRegistry +from meshmind.core.types import Memory +from meshmind.db.in_memory_driver import InMemoryGraphDriver + + +@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() + + +@pytest.fixture +def in_memory_driver(): + return InMemoryGraphDriver() + + +@pytest.fixture +def memory_manager(in_memory_driver): + return MemoryManager(in_memory_driver) + + +@pytest.fixture +def memory_service(memory_manager): + return MemoryService(memory_manager) diff --git a/meshmind/tests/test_db_drivers.py b/meshmind/tests/test_db_drivers.py new file mode 100644 index 0000000..45d940b --- /dev/null +++ b/meshmind/tests/test_db_drivers.py @@ -0,0 +1,66 @@ +from uuid import uuid4 + +from meshmind.core.types import Memory, Triplet +from meshmind.db.in_memory_driver import InMemoryGraphDriver +from meshmind.db.sqlite_driver import SQLiteGraphDriver + + +def _memory_payload(name: str) -> dict: + return { + "uuid": str(uuid4()), + "namespace": "test", + "name": name, + "entity_label": "Note", + "metadata": {"content": name}, + } + + +def test_in_memory_driver_roundtrip(): + driver = InMemoryGraphDriver() + payload = _memory_payload("alpha") + driver.upsert_entity("Note", payload["name"], payload) + + entity = driver.get_entity(payload["uuid"]) + assert entity and entity["name"] == "alpha" + + triplet = Triplet( + subject=payload["uuid"], + predicate="related_to", + object=str(uuid4()), + namespace="test", + entity_label="Relation", + ) + driver.upsert_edge(triplet.subject, triplet.predicate, triplet.object, triplet.dict()) + records = driver.list_triplets("test") + assert records and records[0]["predicate"] == "related_to" + + driver.delete(uuid4()) # deleting unknown should not error + driver.delete_triplet(triplet.subject, triplet.predicate, triplet.object) + assert not driver.list_triplets("test") + + driver.delete(payload["uuid"]) + assert driver.get_entity(payload["uuid"]) is None + + +def test_sqlite_driver_roundtrip(): + driver = SQLiteGraphDriver(":memory:") + payload = _memory_payload("beta") + driver.upsert_entity("Note", payload["name"], payload) + + entity = driver.get_entity(payload["uuid"]) + assert entity and entity["name"] == "beta" + + triplet = Triplet( + subject=payload["uuid"], + predicate="mentions", + object=str(uuid4()), + namespace="test", + entity_label="Relation", + ) + driver.upsert_edge(triplet.subject, triplet.predicate, triplet.object, triplet.dict()) + records = driver.list_triplets("test") + assert records and records[0]["predicate"] == "mentions" + + driver.delete(payload["uuid"]) + assert driver.get_entity(payload["uuid"]) is None + assert not driver.list_triplets("test") diff --git a/meshmind/tests/test_driver_factory.py b/meshmind/tests/test_driver_factory.py new file mode 100644 index 0000000..3bc6c62 --- /dev/null +++ b/meshmind/tests/test_driver_factory.py @@ -0,0 +1,30 @@ +import pytest + +from meshmind.db.factory import create_graph_driver, graph_driver_factory +from meshmind.db.in_memory_driver import InMemoryGraphDriver +from meshmind.db.sqlite_driver import SQLiteGraphDriver + + +def test_create_graph_driver_memory(): + driver = create_graph_driver(backend="memory") + assert isinstance(driver, InMemoryGraphDriver) + + +def test_create_graph_driver_sqlite(tmp_path): + path = tmp_path / "graph.db" + driver = create_graph_driver(backend="sqlite", path=str(path)) + try: + assert isinstance(driver, SQLiteGraphDriver) + finally: + driver.close() + + +def test_graph_driver_factory_callable(): + factory = graph_driver_factory(backend="memory") + driver = factory() + assert isinstance(driver, InMemoryGraphDriver) + + +def test_create_graph_driver_invalid_backend(): + with pytest.raises(ValueError): + create_graph_driver(backend="unknown") 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_observability.py b/meshmind/tests/test_observability.py new file mode 100644 index 0000000..f191d7b --- /dev/null +++ b/meshmind/tests/test_observability.py @@ -0,0 +1,24 @@ +import pytest + +from meshmind.core.observability import log_event, telemetry +from meshmind.pipeline.store import store_memories + + +@pytest.fixture(autouse=True) +def reset_telemetry(): + telemetry.reset() + yield + telemetry.reset() + + +def test_log_event_increments_counter(): + log_event("unit.test", value=1) + snapshot = telemetry.snapshot() + assert snapshot["counters"]["events.unit.test"] == 1 + + +def test_store_memories_tracks_metrics(memory_factory, in_memory_driver): + memories = [memory_factory("one"), memory_factory("two")] + store_memories(memories, in_memory_driver) + snapshot = telemetry.snapshot() + assert snapshot["counters"]["pipeline.store.memories.stored"] == 2 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 diff --git a/meshmind/tests/test_service_interfaces.py b/meshmind/tests/test_service_interfaces.py new file mode 100644 index 0000000..4191829 --- /dev/null +++ b/meshmind/tests/test_service_interfaces.py @@ -0,0 +1,67 @@ +from meshmind.api.grpc import ( + GrpcServiceStub, + IngestMemoriesRequest, + IngestTripletsRequest, + SearchRequest, +) +from meshmind.api.rest import RestAPIStub +from meshmind.api.service import MemoryPayload, SearchPayload, TripletPayload +from meshmind.core.types import Memory + + +def _memory(name: str) -> MemoryPayload: + return MemoryPayload(namespace="test", name=name, entity_label="Note") + + +def test_memory_service_ingest_and_search(memory_service, dummy_encoder): + payloads = [_memory("apple"), _memory("banana")] + uuids = memory_service.ingest_memories(payloads) + assert len(uuids) == 2 + + request = SearchPayload(query="apple", namespace="test", encoder=dummy_encoder, top_k=1) + results = memory_service.search(request) + assert results and isinstance(results[0], Memory) + + +def test_rest_stub_routes(memory_service, dummy_encoder): + app = RestAPIStub(memory_service) + response = app.dispatch( + "POST", + "/memories", + {"memories": [_memory("alpha").model_dump()]}, + ) + assert "uuids" in response + + search_response = app.dispatch( + "POST", + "/search", + {"query": "alpha", "namespace": "test", "encoder": dummy_encoder, "top_k": 1}, + ) + assert search_response["results"] + + +def test_grpc_stub(memory_service, dummy_encoder): + grpc = GrpcServiceStub(memory_service) + request = IngestMemoriesRequest(memories=[_memory("cherry").model_dump()]) + resp = grpc.IngestMemories(request) + assert resp.uuids + + search_resp = grpc.Search( + SearchRequest(query="cherry", namespace="test", encoder=dummy_encoder, top_k=1) + ) + assert search_resp.results + + triplet_resp = grpc.IngestTriplets( + IngestTripletsRequest( + triplets=[ + TripletPayload( + subject=resp.uuids[0], + predicate="linked_to", + object=resp.uuids[0], + namespace="test", + entity_label="Relation", + ).model_dump() + ] + ) + ) + assert triplet_resp.stored == 1