feat: add shared org memory and memory consolidation/archival (#125, #48)#187
feat: add shared org memory and memory consolidation/archival (#125, #48)#187
Conversation
Dependency Review✅ No vulnerabilities or license issues or OpenSSF Scorecard issues found.Scanned FilesNone |
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (50)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds a memory consolidation subsystem (retention, archival, strategies, orchestration) and a shared organizational memory subsystem (OrgMemoryBackend protocol, hybrid backend, SQLite store, access control, models). Integrates org_policies into system prompt rendering, extends config/schema, observability events, public exports, and comprehensive tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Service as MemoryConsolidationService
participant Backend as MemoryBackend
participant Strategy as ConsolidationStrategy
participant Retention as RetentionEnforcer
participant Archive as ArchivalStore
participant Logger as Observability
Note over Service,Logger: run_maintenance(agent_id)
Service->>Retention: cleanup_expired(agent_id)
Retention->>Backend: query(expired by category)
Backend-->>Retention: expired_entries
Retention->>Backend: delete(entry_id) per entry
Retention-->>Service: deleted_count
Service->>Backend: fetch_all_memories(agent_id)
Backend-->>Service: memories
Service->>Strategy: consolidate(memories, agent_id)
Strategy-->>Service: ConsolidationResult(removed_ids, summary_id)
alt archival enabled
Service->>Archive: archive(entry) for each removed_id
Archive-->>Service: archive_entry_id
Service->>Logger: emit ARCHIVAL events
end
Service->>Backend: delete(removed_id) prune
Service->>Backend: enforce_max_memories(agent_id)
Service->>Logger: emit MAINTENANCE_COMPLETE
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
✨ Simplify code
Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the agent memory capabilities by introducing a shared organizational memory and a robust memory consolidation and archival system. These features provide agents with access to company-wide knowledge and ensure efficient, managed memory usage over time. Additionally, organizational policies are now seamlessly integrated into agent prompts, reinforcing adherence to company guidelines. The extensive pre-review and testing ensure the stability and reliability of these new foundational memory components. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
Greptile SummaryThis PR introduces shared org memory ( One logic issue remains: Confidence Score: 3/5
Sequence DiagramsequenceDiagram
participant Caller
participant MCS as MemoryConsolidationService
participant RE as RetentionEnforcer
participant SCS as SimpleConsolidationStrategy
participant Backend as MemoryBackend
participant Archival as ArchivalStore
Caller->>MCS: run_maintenance(agent_id)
MCS->>RE: cleanup_expired(agent_id)
RE->>Backend: retrieve(agent_id, query[until=cutoff])
Backend-->>RE: expired entries
loop per expired entry
RE->>Backend: delete(agent_id, entry.id)
Backend-->>RE: bool (checked ✓)
end
RE-->>MCS: deleted_count
MCS->>MCS: run_consolidation(agent_id)
MCS->>Backend: retrieve(agent_id, MemoryQuery(limit=1000))
Backend-->>MCS: all entries
MCS->>SCS: consolidate(entries, agent_id)
loop per category group (>= threshold)
SCS->>Backend: store(agent_id, summary_request)
Backend-->>SCS: new_id
loop per entry to_remove
SCS->>Backend: delete(agent_id, entry.id)
Backend-->>SCS: bool (NOT checked ⚠️)
Note over SCS: entry.id always appended to removed_ids
end
end
SCS-->>MCS: ConsolidationResult(removed_ids, summary_id)
opt archival enabled
MCS->>Archival: archive(archival_entry) per removed_id
Note over MCS,Archival: May archive entries still in hot storage if delete returned False
end
MCS->>MCS: enforce_max_memories(agent_id)
MCS->>Backend: count(agent_id)
Backend-->>MCS: total
loop while remaining > 0
MCS->>Backend: retrieve(agent_id, MemoryQuery(limit=batch))
Backend-->>MCS: entries
loop per entry
MCS->>Backend: delete(agent_id, entry.id)
Backend-->>MCS: bool (checked ✓)
end
Note over MCS: remaining -= len(entries) not deleted count ⚠️
end
MCS-->>Caller: ConsolidationResult
Last reviewed commit: 73f4b63 |
There was a problem hiding this comment.
Actionable comments posted: 19
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unit/observability/test_events.py (1)
175-210:⚠️ Potential issue | 🟡 MinorAdd constant-level checks for the new event domains.
This only proves
events.consolidationandevents.org_memoryare discoverable. An empty module, or one missing its expected public constants, would still pass here. Please add at least one direct assertion per new domain so the new observability surface cannot regress silently.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/observability/test_events.py` around lines 175 - 210, The test test_all_domain_modules_discovered only checks pkgutil discovery for events.consolidation and events.org_memory but doesn't verify module contents; update the test to import events.consolidation and events.org_memory and assert each exposes at least one expected public constant (e.g., specific event name or __all__ entry) so empty modules will fail—add one direct assertion per new domain (referencing events.consolidation and events.org_memory and the test_all_domain_modules_discovered function) that checks a known constant or that getattr(module, "<EXPECTED_CONST>", None) is not None.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ai_company/engine/prompt_template.py`:
- Around line 122-128: The org_policies block in the Jinja prompt_template.py
renders multiline policies incorrectly because only the first line gets a "- "
prefix; update the template rendering of org_policies so each policy's
subsequent lines are indented rather than treated as new list items or raw
markdown. Specifically, locate the Jinja block that iterates over org_policies
(the section with "{% for policy in org_policies %}- {{ policy }}{% endfor %}")
and change it to normalize newlines (for example using a Jinja replace filter or
a split/join): render "- {{ policy|replace('\n', '\n ') }}" or split policy
into lines and emit the first line with "- " and prefix continuation lines with
two spaces, ensuring every policy becomes a single bulleted item with indented
continuation lines. Ensure you update any similar usage for core_policies if
present.
In `@src/ai_company/memory/consolidation/archival.py`:
- Around line 33-55: search() and restore() must be agent-scoped like count() to
prevent cross-agent reads: update the method signatures for search(self, query:
MemoryQuery) -> tuple[ArchivalEntry, ...] and restore(self, entry_id:
NotBlankStr) -> ArchivalEntry | None to accept an additional agent_id:
NotBlankStr parameter (i.e., search(self, agent_id: NotBlankStr, query:
MemoryQuery) and restore(self, agent_id: NotBlankStr, entry_id: NotBlankStr)),
update their docstrings to document the agent_id arg, and adjust any type hints
and return descriptions accordingly so implementations of ArchivalEntry-backed
stores can enforce ownership at the protocol boundary (count already shows the
pattern to follow).
In `@src/ai_company/memory/consolidation/retention.py`:
- Around line 71-89: The cleanup_expired function accepts a now parameter that
may be a naive datetime, leading to unsafe comparisons with timezone-aware
created_at values; add a defensive validation at the start of cleanup_expired to
ensure now is timezone-aware (or convert naive datetimes to UTC) before
computing cutoffs: if now is None set it to datetime.now(UTC) as before,
otherwise check now.tzinfo and either attach UTC or raise a ValueError, and then
use that tz-aware now for subsequent cutoff and comparisons (referencing
cleanup_expired, now, UTC, and created_at in the retention logic).
- Around line 97-101: The MemoryQuery call hardcodes limit=1000 which imposes an
undocumented batching behavior; update the enclosing function's signature (the
retention/consolidation function that builds MemoryQuery) to accept an optional
batch_limit (default 1000) and use that value in MemoryQuery(limit=batch_limit),
and also update the function docstring to document the batching behavior and the
meaning of batch_limit (e.g., that a single invocation will only process up to
batch_limit expired memories per category and callers can increase or loop until
none remain); reference MemoryQuery, the local variables category and cutoff,
and the new batch_limit parameter in the docstring.
In `@src/ai_company/memory/consolidation/service.py`:
- Around line 139-152: The current eviction uses MemoryQuery(limit=excess) and
_backend.retrieve(agent_id, query), but MemoryQuery.limit is capped at 1000 and
retrieve() doesn’t guarantee eviction order; replace this with a paging or
dedicated eviction API: iterate pages of candidates (using existing
cursor/offset parameters on _backend.retrieve if available) until you collect or
delete 'excess' items, respecting _config.max_memories_per_agent, or add a new
backend method (e.g., list_eviction_candidates(agent_id, limit, cursor) or
retrieve_eviction_victims(agent_id, limit, sort="oldest"/"least_relevant")) that
returns victims in the intended order and use that with _backend.delete for each
returned entry.id; ensure you handle pagination/cursor and keep deleting until
deleted == excess.
- Around line 90-106: The consolidation flow archives after calling
self._strategy.consolidate(), which can cause originals to be deleted before
archival and lead to data loss if archiving fails; change the ordering so that
when self._archival_store is set and self._config.archival.enabled is true you
call self._archive_entries(agent_id, entries, result_candidate_ids) before
invoking self._strategy.consolidate() (or alternately have the consolidation
strategy return the list of candidate removed IDs without performing deletions
so you can archive them first), then only call self._strategy.consolidate() to
perform destructive removals and construct the ConsolidationResult with
archived_count from the earlier archive step; use the existing symbols
_archive_entries, _strategy.consolidate, and ConsolidationResult to implement
this ordering change.
In `@src/ai_company/memory/consolidation/simple_strategy.py`:
- Around line 101-115: The current flow stores the consolidated summary via
MemoryStoreRequest and self._backend.store before deleting originals
(to_remove), which can leave orphaned summaries if any self._backend.delete
fails; change the sequence so deletions are attempted first (iterate to_remove
and call self._backend.delete, appending ids to removed_ids) and only call
self._backend.store to create the summary and set summary_id if all deletes
succeed, or implement a rollback pattern that deletes the newly created summary
(using summary_id) if any subsequent delete fails; reference MemoryStoreRequest,
MemoryMetadata, self._backend.store, self._backend.delete, summary_id,
to_remove, and removed_ids when making the change.
In `@src/ai_company/memory/org/access_control.py`:
- Around line 68-77: The WriteAccessConfig.rules mapping currently allows
missing OrgFactCategory entries which causes check_write_access to fallback to
permissive human-allowed rules; update validation so the rules map is closed: in
the model_validator _wrap_rules_readonly (or a new `@model_validator` pre/post
hook for WriteAccessConfig) ensure every OrgFactCategory key is present by
iterating OrgFactCategory and either (a) adding a deny-all CategoryWriteRule for
any missing category or (b) raise a validation error if any category is absent;
update the MappingProxyType assignment to use the completed map so rules is
immutable and fully covers all OrgFactCategory values before returning self.
In `@src/ai_company/memory/org/factory.py`:
- Around line 50-53: The error message in factory.py currently hardcodes the
valid backends list; update the construction of msg to derive the valid backends
from OrgMemoryConfig._VALID_BACKENDS instead of the static
"['hybrid_prompt_retrieval']" so it stays in sync with the source of truth—use
OrgMemoryConfig._VALID_BACKENDS (formatting as a repr or joined string) together
with config.backend in the message to produce the same informative output.
In `@src/ai_company/memory/org/hybrid_backend.py`:
- Around line 158-166: The handler in OrgMemoryBackend.query currently logs and
re-raises any Exception from self._store.query, which leaks backend errors;
instead catch the Exception, log it (using ORG_MEMORY_QUERY_FAILED and
logger.exception) and raise an OrgMemoryQueryError (from the org protocol) with
the original exception as the cause (e.g., "raise OrgMemoryQueryError(...) from
e") so callers only see the protocol-level error; update the except block around
self._store.query to construct and raise OrgMemoryQueryError while preserving
the original exception context.
- Around line 114-134: list_policies currently always returns in-memory
self._core_policies while write() persists OrgFactCategory.CORE_POLICY to
OrgFactStore, so new core policies never surface; fix by choosing one source of
truth — simplest is to reject CORE_POLICY writes in this backend: update the
write() implementation (the method named write in this class) to detect
OrgFactCategory.CORE_POLICY and raise/return an error (or log and ignore)
instead of persisting to OrgFactStore, and ensure logger messages (e.g.,
ORG_MEMORY_POLICIES_LISTED) reflect that core policies are immutable here;
alternatively, if you prefer persisted core policies, change list_policies to
read CORE_POLICY records from the OrgFactStore instead of using
self._core_policies so persisted writes appear in listings.
- Around line 207-233: The current flow computes version via
_compute_next_version() and then calls _store.save(fact) separately, which
allows lookup failures to surface as the wrong exception and permits concurrent
writers to get identical versions; change this by moving version allocation into
a single store-side transactional operation: add a store method (e.g.,
save_with_next_version or extend _store.save to accept an OrgFact without
version) that, inside a DB transaction, reads the current max version for the
given category, increments it, enforces a uniqueness constraint/index on
(category, version), assigns the version to the new OrgFact and persists it
atomically, and have the hybrid backend call that store method instead of
calling _compute_next_version() plus _store.save(); also ensure the store maps
lookup/constraint failures to OrgMemoryWriteError so callers see the correct
exception (apply the same change for the other write site that mirrors this
logic).
In `@src/ai_company/memory/org/models.py`:
- Around line 47-79: The human branch of
OrgFactAuthor._validate_author_consistency currently only rejects agent_id but
allows seniority; update the is_human branch in the _validate_author_consistency
model_validator to also enforce that seniority is None (mirror the logging and
raise ValueError behavior used for agent_id), using ORG_MEMORY_MODEL_INVALID
with model="OrgFactAuthor" and field="seniority" to log the violation before
raising.
In `@src/ai_company/memory/org/store.py`:
- Around line 421-435: The SQL query in the org facts retrieval (the sql
variable built in the function that calls db.execute) orders only by created_at,
causing recency to outrank relevance; modify the query to ORDER BY relevance
first (e.g., a computed expression that ranks exact matches highest,
substring/weak matches next) and then created_at as a tiebreaker so
HybridPromptRetrievalBackend.query() receives rows sorted by relevance then
recency; update the SQL construction where clauses and params accordingly and
keep raising OrgMemoryQueryError on sqlite3.Error as before.
- Around line 323-340: The current insert in the org fact store uses "INSERT OR
REPLACE" which destroys previous rows and audit history; update the write in the
function that calls db.execute for the org_facts INSERT (the block using
fact.id, fact.content, category, author_agent_id, author_seniority,
author_is_human, created_at, version) to use a plain "INSERT INTO org_facts ..."
and remove OR REPLACE, and add explicit duplicate-id handling (catch
IntegrityError or DBUniqueConstraint exceptions and surface/raise a clear error)
or implement a dedicated update/versioning path (e.g., separate
update_fact/version_fact function) instead of silently replacing rows. Ensure
any calling code expects the IntegrityError or uses the new update/version API
accordingly.
In `@tests/unit/memory/consolidation/test_archival.py`:
- Around line 23-36: The mock methods archive, search, restore, and count in the
test mock should use the same NotBlankStr types as the ArchivalStore protocol:
change parameter and return type annotations to use NotBlankStr where
appropriate (e.g., archive(self, entry: ArchivalEntry) -> NotBlankStr or archive
ID returns NotBlankStr, restore(self, entry_id: NotBlankStr) -> ArchivalEntry |
None, count(self, agent_id: NotBlankStr) -> int) and update search signature if
the protocol expects different types; import NotBlankStr from the types module
and adjust signatures for archive, search, restore, and count to match the
protocol exactly.
In `@tests/unit/memory/consolidation/test_retention.py`:
- Around line 87-109: In test_mixed_categories the call _make_entry("e1",
MemoryCategory.EPISODIC) creates an entry but its return is discarded (unused
variable e1); update test_mixed_categories to either assign the created entry to
a variable (e.g., episodic_entry = _make_entry("e1", MemoryCategory.EPISODIC))
if you intend to keep it for clarity/inspection, or remove the call entirely if
it’s a no-op, or add an explicit comment next to the call clarifying it is
intentionally unused — ensure the change references the _make_entry call and
MemoryCategory.EPISODIC so the test intent is clear.
In `@tests/unit/memory/consolidation/test_service.py`:
- Around line 61-77: The test test_run_consolidation_skipped_when_no_strategy
uses capsys to assert "consolidation.run.skipped" in stdout which is fragile;
update the test to capture logs reliably by using a logging capture or mocking
the logger used by MemoryConsolidationService/run_consolidation instead of
capsys — e.g., inject or patch the structlog logger (or use pytest's caplog
configured for the project's logger) to assert the "consolidation.run.skipped"
message was emitted when MemoryConsolidationService.run_consolidation is called
with strategy=None and validate the same result assertions (consolidated_count,
removed_ids, summary_id).
In `@tests/unit/memory/org/test_store.py`:
- Around line 335-349: The test uses AsyncMock with a custom __getitem__ lambda
to emulate a DB row for _row_to_org_fact but AsyncMock may not faithfully mimic
sqlite3.Row behavior; replace the AsyncMock with a simple, synchronous row-like
object (e.g., a dict subclass, types.SimpleNamespace, or a custom lightweight
class that implements __getitem__) that returns values from malformed_row so
_row_to_org_fact sees realistic item access; update the test name
test_row_parse_error_wraps_in_query_error to use that row-like object instead of
AsyncMock and keep the same malformed_row and the pytest.raises assertion.
---
Outside diff comments:
In `@tests/unit/observability/test_events.py`:
- Around line 175-210: The test test_all_domain_modules_discovered only checks
pkgutil discovery for events.consolidation and events.org_memory but doesn't
verify module contents; update the test to import events.consolidation and
events.org_memory and assert each exposes at least one expected public constant
(e.g., specific event name or __all__ entry) so empty modules will fail—add one
direct assertion per new domain (referencing events.consolidation and
events.org_memory and the test_all_domain_modules_discovered function) that
checks a known constant or that getattr(module, "<EXPECTED_CONST>", None) is not
None.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bb8d3a2e-4396-4348-9c09-b25e22452b96
📒 Files selected for processing (50)
CLAUDE.mdDESIGN_SPEC.mdREADME.mdsrc/ai_company/config/defaults.pysrc/ai_company/config/schema.pysrc/ai_company/core/enums.pysrc/ai_company/engine/prompt.pysrc/ai_company/engine/prompt_template.pysrc/ai_company/memory/__init__.pysrc/ai_company/memory/config.pysrc/ai_company/memory/consolidation/__init__.pysrc/ai_company/memory/consolidation/archival.pysrc/ai_company/memory/consolidation/config.pysrc/ai_company/memory/consolidation/models.pysrc/ai_company/memory/consolidation/retention.pysrc/ai_company/memory/consolidation/service.pysrc/ai_company/memory/consolidation/simple_strategy.pysrc/ai_company/memory/consolidation/strategy.pysrc/ai_company/memory/org/__init__.pysrc/ai_company/memory/org/access_control.pysrc/ai_company/memory/org/config.pysrc/ai_company/memory/org/errors.pysrc/ai_company/memory/org/factory.pysrc/ai_company/memory/org/hybrid_backend.pysrc/ai_company/memory/org/models.pysrc/ai_company/memory/org/protocol.pysrc/ai_company/memory/org/store.pysrc/ai_company/observability/events/consolidation.pysrc/ai_company/observability/events/org_memory.pytests/unit/config/conftest.pytests/unit/engine/test_prompt.pytests/unit/memory/consolidation/__init__.pytests/unit/memory/consolidation/test_archival.pytests/unit/memory/consolidation/test_config.pytests/unit/memory/consolidation/test_models.pytests/unit/memory/consolidation/test_retention.pytests/unit/memory/consolidation/test_service.pytests/unit/memory/consolidation/test_strategy.pytests/unit/memory/org/__init__.pytests/unit/memory/org/test_access_control.pytests/unit/memory/org/test_config.pytests/unit/memory/org/test_errors.pytests/unit/memory/org/test_factory.pytests/unit/memory/org/test_hybrid_backend.pytests/unit/memory/org/test_models.pytests/unit/memory/org/test_prompt_integration.pytests/unit/memory/org/test_protocol.pytests/unit/memory/org/test_store.pytests/unit/memory/test_init.pytests/unit/observability/test_events.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Agent
- GitHub Check: Greptile Review
🧰 Additional context used
📓 Path-based instructions (5)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Do NOT usefrom __future__ import annotations— Python 3.14 has PEP 649 native lazy annotations
Useexcept A, B:syntax (without parentheses) per PEP 758 — ruff enforces this on Python 3.14
All public functions must have type hints; use mypy strict mode for type-checking
Use Google-style docstrings on all public classes and functions; enforced by ruff D rules
Create new objects instead of mutating existing ones; for non-Pydantic internal collections (registries, BaseTool), use copy.deepcopy() at construction + MappingProxyType wrapping for read-only enforcement
For dict/list fields in frozen Pydantic models, rely on frozen=True for field reassignment prevention and copy.deepcopy() at system boundaries (tool execution, LLM provider serialization, inter-agent delegation, persistence serialization)
Use frozen Pydantic models for config/identity; use separate mutable-via-copy models (with model_copy(update=...)) for runtime state that evolves; never mix static config fields with mutable runtime fields in one model
Use Pydantic v2 (BaseModel, model_validator, computed_field, ConfigDict); use@computed_fieldfor derived values instead of storing redundant fields; use NotBlankStr for all identifier/name fields (including optional and tuple variants) instead of manual whitespace validators
Prefer asyncio.TaskGroup for fan-out/fan-in parallel operations in new code (multiple tool invocations, parallel agent calls); prefer structured concurrency over bare create_task
Keep functions under 50 lines and files under 800 lines
Handle errors explicitly, never silently swallow exceptions
Validate at system boundaries (user input, external APIs, config files)
Use line length of 88 characters (ruff)
Files:
src/ai_company/config/defaults.pytests/unit/memory/org/test_protocol.pytests/unit/memory/org/test_factory.pytests/unit/memory/consolidation/test_models.pysrc/ai_company/memory/org/factory.pysrc/ai_company/memory/consolidation/strategy.pytests/unit/observability/test_events.pytests/unit/memory/consolidation/test_service.pytests/unit/config/conftest.pysrc/ai_company/memory/consolidation/models.pysrc/ai_company/memory/org/models.pysrc/ai_company/engine/prompt_template.pysrc/ai_company/memory/consolidation/archival.pysrc/ai_company/memory/consolidation/service.pytests/unit/memory/org/test_errors.pysrc/ai_company/core/enums.pytests/unit/memory/org/test_models.pysrc/ai_company/config/schema.pysrc/ai_company/memory/org/hybrid_backend.pysrc/ai_company/memory/consolidation/retention.pytests/unit/memory/consolidation/test_config.pysrc/ai_company/engine/prompt.pysrc/ai_company/memory/org/errors.pysrc/ai_company/memory/__init__.pysrc/ai_company/memory/org/config.pysrc/ai_company/observability/events/consolidation.pysrc/ai_company/memory/org/access_control.pysrc/ai_company/memory/org/protocol.pytests/unit/memory/consolidation/test_archival.pytests/unit/engine/test_prompt.pysrc/ai_company/memory/config.pysrc/ai_company/observability/events/org_memory.pytests/unit/memory/consolidation/test_retention.pytests/unit/memory/org/test_prompt_integration.pytests/unit/memory/consolidation/test_strategy.pysrc/ai_company/memory/consolidation/config.pytests/unit/memory/test_init.pytests/unit/memory/org/test_store.pytests/unit/memory/org/test_hybrid_backend.pysrc/ai_company/memory/org/store.pytests/unit/memory/org/test_config.pytests/unit/memory/org/test_access_control.pysrc/ai_company/memory/consolidation/__init__.pysrc/ai_company/memory/org/__init__.pysrc/ai_company/memory/consolidation/simple_strategy.py
src/ai_company/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/ai_company/**/*.py: Every module with business logic must import and use get_logger(name) from ai_company.observability; never use import logging or logging.getLogger() or print() in application code
Always use 'logger' as the variable name (not '_logger', not 'log')
Always use event name constants from ai_company.observability.events domain modules (e.g., PROVIDER_CALL_START from events.provider) instead of string literals
Use structured logging with logger.info(EVENT, key=value) — never use logger.info('msg %s', val) string formatting
All error paths must log at WARNING or ERROR with context before raising
All state transitions must log at INFO level
Use DEBUG level logging for object creation, internal flow, and entry/exit of key functions
Files:
src/ai_company/config/defaults.pysrc/ai_company/memory/org/factory.pysrc/ai_company/memory/consolidation/strategy.pysrc/ai_company/memory/consolidation/models.pysrc/ai_company/memory/org/models.pysrc/ai_company/engine/prompt_template.pysrc/ai_company/memory/consolidation/archival.pysrc/ai_company/memory/consolidation/service.pysrc/ai_company/core/enums.pysrc/ai_company/config/schema.pysrc/ai_company/memory/org/hybrid_backend.pysrc/ai_company/memory/consolidation/retention.pysrc/ai_company/engine/prompt.pysrc/ai_company/memory/org/errors.pysrc/ai_company/memory/__init__.pysrc/ai_company/memory/org/config.pysrc/ai_company/observability/events/consolidation.pysrc/ai_company/memory/org/access_control.pysrc/ai_company/memory/org/protocol.pysrc/ai_company/memory/config.pysrc/ai_company/observability/events/org_memory.pysrc/ai_company/memory/consolidation/config.pysrc/ai_company/memory/org/store.pysrc/ai_company/memory/consolidation/__init__.pysrc/ai_company/memory/org/__init__.pysrc/ai_company/memory/consolidation/simple_strategy.py
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Never use real vendor names (Anthropic, OpenAI, Claude, GPT, etc.) in project-owned code, docstrings, comments, tests, or config examples; use generic names (example-provider, example-large-001, example-medium-001, example-small-001, large/medium/small aliases)
Files:
src/ai_company/config/defaults.pysrc/ai_company/memory/org/factory.pysrc/ai_company/memory/consolidation/strategy.pysrc/ai_company/memory/consolidation/models.pysrc/ai_company/memory/org/models.pysrc/ai_company/engine/prompt_template.pysrc/ai_company/memory/consolidation/archival.pysrc/ai_company/memory/consolidation/service.pysrc/ai_company/core/enums.pysrc/ai_company/config/schema.pysrc/ai_company/memory/org/hybrid_backend.pysrc/ai_company/memory/consolidation/retention.pysrc/ai_company/engine/prompt.pysrc/ai_company/memory/org/errors.pysrc/ai_company/memory/__init__.pysrc/ai_company/memory/org/config.pysrc/ai_company/observability/events/consolidation.pysrc/ai_company/memory/org/access_control.pysrc/ai_company/memory/org/protocol.pysrc/ai_company/memory/config.pysrc/ai_company/observability/events/org_memory.pysrc/ai_company/memory/consolidation/config.pysrc/ai_company/memory/org/store.pysrc/ai_company/memory/consolidation/__init__.pysrc/ai_company/memory/org/__init__.pysrc/ai_company/memory/consolidation/simple_strategy.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Mark tests with@pytest.mark.unit,@pytest.mark.integration,@pytest.mark.e2e, or@pytest.mark.slow
Prefer@pytest.mark.parametrizefor testing similar cases
In tests, use test-provider, test-small-001, etc. instead of real vendor names
Files:
tests/unit/memory/org/test_protocol.pytests/unit/memory/org/test_factory.pytests/unit/memory/consolidation/test_models.pytests/unit/observability/test_events.pytests/unit/memory/consolidation/test_service.pytests/unit/config/conftest.pytests/unit/memory/org/test_errors.pytests/unit/memory/org/test_models.pytests/unit/memory/consolidation/test_config.pytests/unit/memory/consolidation/test_archival.pytests/unit/engine/test_prompt.pytests/unit/memory/consolidation/test_retention.pytests/unit/memory/org/test_prompt_integration.pytests/unit/memory/consolidation/test_strategy.pytests/unit/memory/test_init.pytests/unit/memory/org/test_store.pytests/unit/memory/org/test_hybrid_backend.pytests/unit/memory/org/test_config.pytests/unit/memory/org/test_access_control.py
src/ai_company/engine/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
RetryExhaustedError signals that all retries failed — the engine layer catches this to trigger fallback chains
Files:
src/ai_company/engine/prompt_template.pysrc/ai_company/engine/prompt.py
🧠 Learnings (9)
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : Always use event name constants from ai_company.observability.events domain modules (e.g., PROVIDER_CALL_START from events.provider) instead of string literals
Applied to files:
tests/unit/observability/test_events.pyCLAUDE.mdsrc/ai_company/config/schema.pysrc/ai_company/observability/events/consolidation.pysrc/ai_company/observability/events/org_memory.py
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/providers/**/*.py : RetryConfig and RateLimiterConfig are set per-provider in ProviderConfig
Applied to files:
tests/unit/config/conftest.py
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to **/*.py : Use frozen Pydantic models for config/identity; use separate mutable-via-copy models (with model_copy(update=...)) for runtime state that evolves; never mix static config fields with mutable runtime fields in one model
Applied to files:
src/ai_company/memory/consolidation/models.pysrc/ai_company/memory/org/models.pysrc/ai_company/memory/org/config.pysrc/ai_company/memory/consolidation/config.py
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : Every module with business logic must import and use get_logger(__name__) from ai_company.observability; never use import logging or logging.getLogger() or print() in application code
Applied to files:
CLAUDE.mdsrc/ai_company/config/schema.py
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : Always use 'logger' as the variable name (not '_logger', not 'log')
Applied to files:
CLAUDE.md
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : Use structured logging with logger.info(EVENT, key=value) — never use logger.info('msg %s', val) string formatting
Applied to files:
CLAUDE.md
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : All error paths must log at WARNING or ERROR with context before raising
Applied to files:
CLAUDE.md
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : All state transitions must log at INFO level
Applied to files:
CLAUDE.md
📚 Learning: 2026-03-09T12:14:21.716Z
Learnt from: CR
Repo: Aureliolo/ai-company PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-09T12:14:21.716Z
Learning: Applies to src/ai_company/**/*.py : Use DEBUG level logging for object creation, internal flow, and entry/exit of key functions
Applied to files:
CLAUDE.md
🧬 Code graph analysis (28)
tests/unit/memory/org/test_protocol.py (2)
src/ai_company/memory/org/access_control.py (1)
WriteAccessConfig(59-77)src/ai_company/memory/org/protocol.py (1)
OrgMemoryBackend(20-110)
tests/unit/memory/org/test_factory.py (3)
src/ai_company/memory/org/config.py (1)
OrgMemoryConfig(60-108)src/ai_company/memory/org/errors.py (1)
OrgMemoryConfigError(28-29)src/ai_company/memory/org/factory.py (1)
create_org_memory_backend(21-59)
tests/unit/memory/consolidation/test_models.py (3)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/models.py (3)
ArchivalEntry(47-82)ConsolidationResult(16-44)RetentionRule(85-101)src/ai_company/memory/models.py (1)
MemoryMetadata(20-52)
src/ai_company/memory/org/factory.py (6)
src/ai_company/memory/org/config.py (1)
OrgMemoryConfig(60-108)src/ai_company/memory/org/errors.py (1)
OrgMemoryConfigError(28-29)src/ai_company/memory/org/hybrid_backend.py (1)
HybridPromptRetrievalBackend(39-262)src/ai_company/memory/org/protocol.py (1)
OrgMemoryBackend(20-110)src/ai_company/memory/org/store.py (1)
OrgFactStore(60-157)src/ai_company/observability/_logger.py (1)
get_logger(8-28)
src/ai_company/memory/consolidation/strategy.py (3)
src/ai_company/memory/consolidation/models.py (1)
ConsolidationResult(16-44)src/ai_company/memory/models.py (1)
MemoryEntry(82-150)src/ai_company/memory/consolidation/simple_strategy.py (1)
consolidate(63-130)
tests/unit/memory/consolidation/test_service.py (5)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/config.py (3)
ArchivalConfig(47-65)ConsolidationConfig(68-96)RetentionConfig(15-44)src/ai_company/memory/consolidation/models.py (1)
ConsolidationResult(16-44)src/ai_company/memory/consolidation/service.py (5)
MemoryConsolidationService(37-268)run_consolidation(66-122)enforce_max_memories(124-168)cleanup_retention(170-182)run_maintenance(184-213)src/ai_company/memory/models.py (2)
MemoryEntry(82-150)MemoryMetadata(20-52)
tests/unit/config/conftest.py (1)
src/ai_company/memory/org/config.py (1)
OrgMemoryConfig(60-108)
src/ai_company/memory/consolidation/models.py (2)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/models.py (1)
MemoryMetadata(20-52)
src/ai_company/memory/org/models.py (2)
src/ai_company/core/enums.py (2)
OrgFactCategory(119-129)SeniorityLevel(6-21)src/ai_company/observability/_logger.py (1)
get_logger(8-28)
src/ai_company/memory/consolidation/archival.py (2)
src/ai_company/memory/consolidation/models.py (1)
ArchivalEntry(47-82)src/ai_company/memory/models.py (1)
MemoryQuery(153-230)
src/ai_company/memory/consolidation/service.py (9)
src/ai_company/memory/consolidation/archival.py (3)
ArchivalStore(15-64)count(55-64)archive(22-31)src/ai_company/memory/consolidation/config.py (1)
ConsolidationConfig(68-96)src/ai_company/memory/consolidation/models.py (2)
ArchivalEntry(47-82)ConsolidationResult(16-44)src/ai_company/memory/consolidation/retention.py (2)
RetentionEnforcer(25-129)cleanup_expired(71-129)src/ai_company/memory/consolidation/strategy.py (2)
ConsolidationStrategy(15-38)consolidate(23-38)src/ai_company/memory/models.py (2)
MemoryEntry(82-150)MemoryQuery(153-230)src/ai_company/memory/protocol.py (1)
MemoryBackend(20-182)src/ai_company/memory/consolidation/simple_strategy.py (1)
consolidate(63-130)tests/unit/memory/consolidation/test_archival.py (2)
count(35-36)archive(23-27)
tests/unit/memory/org/test_models.py (2)
src/ai_company/core/enums.py (2)
OrgFactCategory(119-129)SeniorityLevel(6-21)src/ai_company/memory/org/models.py (4)
OrgFact(82-101)OrgFactAuthor(19-79)OrgFactWriteRequest(104-115)OrgMemoryQuery(118-142)
src/ai_company/config/schema.py (1)
src/ai_company/memory/org/config.py (1)
OrgMemoryConfig(60-108)
src/ai_company/memory/org/hybrid_backend.py (5)
src/ai_company/core/enums.py (1)
OrgFactCategory(119-129)src/ai_company/memory/org/access_control.py (2)
WriteAccessConfig(59-77)require_write_access(112-145)src/ai_company/memory/org/errors.py (2)
OrgMemoryConnectionError(12-13)OrgMemoryWriteError(16-17)src/ai_company/memory/org/models.py (4)
OrgFact(82-101)OrgFactAuthor(19-79)OrgFactWriteRequest(104-115)OrgMemoryQuery(118-142)src/ai_company/memory/org/store.py (13)
OrgFactStore(60-157)connect(63-65)connect(245-276)disconnect(67-69)disconnect(286-298)is_connected(72-74)is_connected(502-504)query(103-124)query(382-435)save(76-86)save(312-348)list_by_category(126-142)list_by_category(437-468)
src/ai_company/memory/consolidation/retention.py (4)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/config.py (1)
RetentionConfig(15-44)src/ai_company/memory/models.py (1)
MemoryQuery(153-230)src/ai_company/memory/protocol.py (1)
MemoryBackend(20-182)
tests/unit/memory/consolidation/test_config.py (3)
src/ai_company/core/enums.py (2)
ConsolidationInterval(110-116)MemoryCategory(100-107)src/ai_company/memory/consolidation/config.py (3)
ArchivalConfig(47-65)ConsolidationConfig(68-96)RetentionConfig(15-44)src/ai_company/memory/consolidation/models.py (1)
RetentionRule(85-101)
src/ai_company/memory/org/config.py (3)
src/ai_company/memory/org/access_control.py (1)
WriteAccessConfig(59-77)src/ai_company/observability/_logger.py (1)
get_logger(8-28)src/ai_company/memory/config.py (1)
_validate_backend_name(181-195)
tests/unit/memory/consolidation/test_archival.py (3)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/archival.py (5)
ArchivalStore(15-64)archive(22-31)search(33-42)restore(44-53)count(55-64)src/ai_company/memory/consolidation/models.py (1)
ArchivalEntry(47-82)
src/ai_company/memory/config.py (1)
src/ai_company/memory/consolidation/config.py (1)
ConsolidationConfig(68-96)
tests/unit/memory/consolidation/test_retention.py (5)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/config.py (1)
RetentionConfig(15-44)src/ai_company/memory/consolidation/models.py (1)
RetentionRule(85-101)src/ai_company/memory/consolidation/retention.py (2)
RetentionEnforcer(25-129)cleanup_expired(71-129)src/ai_company/memory/models.py (2)
MemoryEntry(82-150)MemoryMetadata(20-52)
tests/unit/memory/org/test_prompt_integration.py (4)
src/ai_company/core/agent.py (2)
AgentIdentity(265-323)ModelConfig(149-178)src/ai_company/core/enums.py (1)
SeniorityLevel(6-21)src/ai_company/engine/prompt.py (1)
build_system_prompt(150-232)tests/unit/engine/test_prompt.py (1)
_make_agent(309-323)
src/ai_company/memory/consolidation/config.py (2)
src/ai_company/core/enums.py (1)
ConsolidationInterval(110-116)src/ai_company/memory/consolidation/models.py (1)
RetentionRule(85-101)
tests/unit/memory/org/test_store.py (3)
src/ai_company/memory/org/errors.py (3)
OrgMemoryConnectionError(12-13)OrgMemoryQueryError(20-21)OrgMemoryWriteError(16-17)src/ai_company/memory/org/models.py (2)
OrgFact(82-101)OrgFactAuthor(19-79)src/ai_company/memory/org/store.py (16)
SQLiteOrgFactStore(227-509)_row_to_org_fact(182-224)connect(63-65)connect(245-276)is_connected(72-74)is_connected(502-504)save(76-86)save(312-348)get(88-101)get(350-380)query(103-124)query(382-435)list_by_category(126-142)list_by_category(437-468)delete(144-157)delete(470-499)
tests/unit/memory/org/test_config.py (1)
src/ai_company/memory/org/config.py (2)
ExtendedStoreConfig(19-57)OrgMemoryConfig(60-108)
tests/unit/memory/org/test_access_control.py (4)
src/ai_company/core/enums.py (2)
OrgFactCategory(119-129)SeniorityLevel(6-21)src/ai_company/memory/org/access_control.py (4)
CategoryWriteRule(25-43)WriteAccessConfig(59-77)check_write_access(80-109)require_write_access(112-145)src/ai_company/memory/org/errors.py (1)
OrgMemoryAccessDeniedError(24-25)src/ai_company/memory/org/models.py (1)
OrgFactAuthor(19-79)
src/ai_company/memory/consolidation/__init__.py (7)
src/ai_company/memory/consolidation/archival.py (1)
ArchivalStore(15-64)src/ai_company/memory/consolidation/config.py (3)
ArchivalConfig(47-65)ConsolidationConfig(68-96)RetentionConfig(15-44)src/ai_company/memory/consolidation/models.py (3)
ArchivalEntry(47-82)ConsolidationResult(16-44)RetentionRule(85-101)src/ai_company/memory/consolidation/retention.py (1)
RetentionEnforcer(25-129)src/ai_company/memory/consolidation/service.py (1)
MemoryConsolidationService(37-268)src/ai_company/memory/consolidation/simple_strategy.py (1)
SimpleConsolidationStrategy(31-176)src/ai_company/memory/consolidation/strategy.py (1)
ConsolidationStrategy(15-38)
src/ai_company/memory/org/__init__.py (8)
src/ai_company/memory/org/access_control.py (4)
CategoryWriteRule(25-43)WriteAccessConfig(59-77)check_write_access(80-109)require_write_access(112-145)src/ai_company/memory/org/config.py (2)
ExtendedStoreConfig(19-57)OrgMemoryConfig(60-108)src/ai_company/memory/org/errors.py (6)
OrgMemoryAccessDeniedError(24-25)OrgMemoryConfigError(28-29)OrgMemoryConnectionError(12-13)OrgMemoryError(8-9)OrgMemoryQueryError(20-21)OrgMemoryWriteError(16-17)src/ai_company/memory/org/factory.py (1)
create_org_memory_backend(21-59)src/ai_company/memory/org/hybrid_backend.py (1)
HybridPromptRetrievalBackend(39-262)src/ai_company/memory/org/models.py (4)
OrgFact(82-101)OrgFactAuthor(19-79)OrgFactWriteRequest(104-115)OrgMemoryQuery(118-142)src/ai_company/memory/org/protocol.py (1)
OrgMemoryBackend(20-110)src/ai_company/memory/org/store.py (2)
OrgFactStore(60-157)SQLiteOrgFactStore(227-509)
src/ai_company/memory/consolidation/simple_strategy.py (4)
src/ai_company/core/enums.py (1)
MemoryCategory(100-107)src/ai_company/memory/consolidation/models.py (1)
ConsolidationResult(16-44)src/ai_company/memory/models.py (3)
MemoryEntry(82-150)MemoryMetadata(20-52)MemoryStoreRequest(55-79)src/ai_company/memory/protocol.py (1)
MemoryBackend(20-182)
🪛 LanguageTool
README.md
[typographical] ~24-~24: To join two clauses or introduce examples, consider using an em dash.
Context: ...a migrations - Memory Interface (M5) - Pluggable MemoryBackend protocol with ...
(DASH_RULE)
[typographical] ~25-~25: To join two clauses or introduce examples, consider using an em dash.
Context: ...t - Coordination Error Taxonomy (M5) - Post-execution classification pipeline d...
(DASH_RULE)
[typographical] ~26-~26: To join two clauses or introduce examples, consider using an em dash.
Context: ...n failures - Budget Enforcement (M5) - BudgetEnforcer service with pre-flight...
(DASH_RULE)
[grammar] ~26-~26: Please add a punctuation mark at the end of paragraph.
Context: ...cking; CFO agent and advanced reporting pending ### Not implemented yet (planned miles...
(PUNCTUATION_PARAGRAPH_END)
[typographical] ~30-~30: To join two clauses or introduce examples, consider using an em dash.
Context: ... milestones) - Memory Backends (M5) - Initial adapter backend ([ADR-001](docs/...
(DASH_RULE)
CLAUDE.md
[style] ~86-~86: A comma is missing here.
Context: ...nder ai_company.observability.events (e.g. PROVIDER_CALL_START from `events.prov...
(EG_NO_COMMA)
🔇 Additional comments (32)
README.md (1)
24-24: LGTM!The README updates accurately reflect the new memory subsystems introduced in this PR: shared organizational memory with hybrid prompt+retrieval backend, and memory consolidation/archival with pluggable strategies. The formatting is consistent with the existing document style.
Also applies to: 30-30
src/ai_company/core/enums.py (1)
119-130: LGTM!The new
OrgFactCategoryenum follows the established patterns in this file: usesStrEnum, includes a Google-style docstring with design spec reference (§7.4), and provides appropriate category values for organizational knowledge classification. The placement afterConsolidationIntervalmaintains logical grouping.src/ai_company/memory/consolidation/retention.py (1)
1-129: LGTM on observability and structure.The module correctly follows coding guidelines:
- Uses
get_logger(__name__)fromai_company.observability- Uses event constants (
RETENTION_CLEANUP_START,RETENTION_CLEANUP_COMPLETE, etc.) instead of string literals- Structured logging with
key=valueformat- Error paths log at WARNING level with context before continuing
- State transitions (start/complete) log at INFO level
The per-category error isolation with
try/except/continueprovides good resilience.src/ai_company/memory/consolidation/config.py (1)
1-96: LGTM!The configuration models are well-designed:
- All three models (
RetentionConfig,ArchivalConfig,ConsolidationConfig) correctly usefrozen=Truefor immutability- The duplicate category validation in
RetentionConfigproperly detects and reports duplicates- Sensible defaults:
DAILYconsolidation interval, 10,000 max memories per agent, archival disabled by default with 90-day threshold- Uses
default_factoryfor nested configs to avoid mutable default issues- Follows the "frozen Pydantic models for config" pattern. Based on learnings: "Use frozen Pydantic models for config/identity".
src/ai_company/memory/config.py (1)
17-17: LGTM!Clean integration of
ConsolidationConfigintoCompanyMemoryConfig:
- Import placed correctly with other memory-related imports
- Field uses
default_factory=ConsolidationConfigto avoid mutable default issues- Docstring updated to document the new attribute
- Follows the existing patterns in the file for nested configuration
Also applies to: 148-148, 175-178
src/ai_company/memory/org/config.py (1)
1-108: LGTM!The organizational memory configuration models are well-implemented:
- Both
ExtendedStoreConfigandOrgMemoryConfigcorrectly usefrozen=True- Backend validation follows the established pattern from
src/ai_company/memory/config.py(using_VALID_BACKENDSClassVar and logging before raising)- Uses
get_logger(__name__)andCONFIG_VALIDATION_FAILEDevent constant for structured logging in validatorscore_policiesusestuple[NotBlankStr, ...]for immutability- Nested configs use
default_factoryappropriatelyBased on learnings: "Use frozen Pydantic models for config/identity".
src/ai_company/config/schema.py (1)
23-23: LGTM!Clean integration of
OrgMemoryConfigintoRootConfig:
- Import placed correctly alongside other memory-related imports
- Field follows the established pattern for nested configs (
default_factory=OrgMemoryConfig)- Docstring updated to document the new attribute
- Consistent with how other config sections (
memory,persistence,cost_tiers) are integratedAlso applies to: 487-488, 567-570
src/ai_company/config/defaults.py (1)
36-36: LGTM!The
org_memorydefault follows the established pattern: an empty dict that allowsOrgMemoryConfig's Pydantic defaults to take effect during validation. Consistent with how other nested configs (memory,persistence,cost_tiers) are handled in this file.tests/unit/memory/org/test_models.py (5)
1-18: LGTM on module setup and imports.Module structure follows conventions with proper pytest markers, type hints, and organized imports. The
@pytest.mark.unitmarker is applied correctly to all test classes.
21-66: LGTM on OrgFactAuthor test coverage.The tests comprehensively validate:
- Human vs agent author consistency rules
- Required fields for non-human authors (agent_id, seniority)
- Model immutability (frozen behavior)
68-108: LGTM on OrgFact test coverage.Tests validate fact creation, version constraints (must be positive), and immutability correctly.
110-128: LGTM on OrgFactWriteRequest tests.Both blank and whitespace-only content rejection are tested, validating the
NotBlankStrconstraint.
131-153: LGTM on OrgMemoryQuery tests.Tests validate defaults, limit bounds (1-100), and category filtering. Good coverage of field constraints.
tests/unit/memory/org/test_errors.py (1)
1-76: LGTM on error hierarchy tests.Excellent use of
@pytest.mark.parametrizeto test multiple error classes efficiently, validating inheritance, catchability with the base class, and message preservation. This follows the coding guidelines for preferring parametrize for testing similar cases.tests/unit/memory/org/test_config.py (1)
1-56: LGTM on org memory configuration tests.Tests comprehensively cover:
- Default values for both
ExtendedStoreConfigandOrgMemoryConfig- Unknown backend rejection with specific error messages
- Bounds validation for
max_retrieved_per_query- Core policies tuple handling
Well-structured with clear, focused assertions.
tests/unit/memory/consolidation/test_config.py (1)
1-98: LGTM on consolidation configuration tests.Comprehensive test coverage for:
RetentionConfig: defaults, rules with retention days, and duplicate category validationArchivalConfig: defaults, enabled state, and zero threshold rejectionConsolidationConfig: defaults, custom values, and zero max memories rejectionThe separation of
TestRetentionConfigValidationfor duplicate category testing is reasonable for organizational clarity.tests/unit/memory/consolidation/test_retention.py (1)
35-127: Async tests missing explicit@pytest.mark.asynciomarker.Same observation as in
test_archival.py— async test methods don't have explicit markers. Assuming this relies on pytest-asyncio auto mode configuration.src/ai_company/memory/org/factory.py (1)
1-59: LGTM on factory implementation.The factory follows coding guidelines:
- Uses
get_logger(__name__)with the correct variable namelogger- Uses event constants from
observability.events.org_memory- Structured logging with
key=valuepattern- Error path logs at ERROR level before raising
- Google-style docstring with Args/Returns/Raises sections
tests/unit/memory/org/test_prompt_integration.py (2)
15-27: LGTM on helper and model configuration.The
_make_agenthelper correctly uses"test-provider"and"test-model-001"instead of real vendor names, following the coding guidelines.
30-61: LGTM on org_policies prompt integration tests.Tests comprehensively validate:
- No section when no policies provided
- Policies appear in rendered content and sections
- Template version assertion (1.2.0)
- Trimming behavior under token budget
All assertions are clear and well-structured.
tests/unit/memory/consolidation/test_archival.py (1)
52-100: No changes needed — async tests are correctly configured withasyncio_mode = "auto".The test file already marks async test classes with
@pytest.mark.unit, which follows the coding guidelines. Withasyncio_mode = "auto"configured inpyproject.toml, explicit@pytest.mark.asynciodecorators are not required and would introduce unnecessary boilerplate.> Likely an incorrect or invalid review comment.tests/unit/memory/consolidation/test_models.py (1)
1-134: LGTM!Well-structured unit tests covering the consolidation domain models comprehensively. The tests validate:
- Minimal and full construction scenarios
- Field constraints (negative values, zero values)
- Immutability (frozen models)
- Temporal ordering validation for
ArchivalEntryThe test organization by class and clear naming conventions make the test intent clear.
tests/unit/memory/consolidation/test_service.py (1)
46-225: LGTM!The service tests provide comprehensive coverage of:
- Consolidation with/without strategy
- Exception propagation
- Archival integration
- Max memories enforcement with partial delete failures
- Retention cleanup delegation
- Full maintenance orchestration
The mock setup is clean and tests verify both return values and side effects appropriately.
tests/unit/memory/org/test_access_control.py (1)
1-113: LGTM!Comprehensive access control tests validating:
- Default write rules for all category/author combinations (human, senior, junior, lead)
- Custom rule configuration that can deny human access
- Unknown category fallback behavior (defaults to human-only)
require_write_accessraisesOrgMemoryAccessDeniedErrorappropriatelyWriteAccessConfig.rulesimmutability viaMappingProxyTypeThe test fixtures (
_HUMAN,_SENIOR_AGENT, etc.) are well-designed for reuse across tests.tests/unit/memory/consolidation/test_strategy.py (1)
1-175: LGTM!Thorough test coverage for
SimpleConsolidationStrategyincluding:
- Empty input handling
- Threshold-based consolidation triggers
- Multi-category consolidation with proper relevance-based selection
- Summary content truncation for long entries
- Handling of
Nonerelevance scores (treated as 0.0)- Constructor validation for
group_thresholdminimumThe
_make_entryhelper provides good flexibility for creating test fixtures with varying attributes.src/ai_company/memory/consolidation/simple_strategy.py (2)
31-62: LGTM!The class is well-documented with clear Google-style docstrings, proper initialization validation, and appropriate use of structured logging with event constants. The constructor validation for
group_threshold >= 2is correctly implemented.
132-176: LGTM!The helper methods are cleanly implemented:
_select_entriescorrectly handlesNonerelevance by treating it as0.0and usescreated_atas a tiebreaker_build_summaryproperly truncates long content with ellipsis indicationtests/unit/memory/org/test_store.py (1)
1-334: LGTM!Comprehensive test coverage for
SQLiteOrgFactStoreincluding:
- Connection lifecycle (connect, disconnect, double connect)
- All CRUD operations with proper cleanup via
try/finally- Query filtering by category, text, and combined
- Agent author handling with seniority
- Operations when not connected raise
OrgMemoryConnectionError- SQLite error wrapping into domain-specific errors
- Path traversal protection
- LIKE special character escaping (SQL injection prevention)
src/ai_company/memory/org/errors.py (1)
1-29: LGTM!Clean and well-organized error hierarchy that:
- Provides a common base class (
OrgMemoryError) for catching all org memory exceptions- Defines specific exception types for different failure modes (connection, write, query, access denied, config)
- Includes clear docstrings for each exception type
This enables callers to handle errors at the appropriate granularity.
src/ai_company/memory/consolidation/models.py (1)
1-101: LGTM!Well-designed frozen Pydantic models following project conventions:
ConsolidationResult: Immutable result with appropriate defaults and non-negative constraintsArchivalEntry: Comprehensive archived memory record with temporal ordering validation (archived_at >= created_at)RetentionRule: Simple per-category retention policy with minimum 1-day retentionAll models use
frozen=Trueandallow_inf_nan=Falseappropriately, with clear Google-style docstrings. Based on learnings: frozen Pydantic models are correctly used for these config/identity-type structures.tests/unit/memory/org/test_hybrid_backend.py (1)
232-280: Good pinning of the write-error contract.These two mock-store tests lock in the subtle behavior that matters here: unexpected store exceptions get wrapped, while a preexisting
OrgMemoryWriteErroris re-raised unchanged. That should make backend refactors much safer.src/ai_company/engine/prompt.py (1)
139-144: Good trim order fororg_policies.Putting
_SECTION_ORG_POLICIESafter company/tools/task preserves the “trim last” requirement for org policies when the prompt is over budget.
| async def cleanup_expired( | ||
| self, | ||
| agent_id: NotBlankStr, | ||
| now: datetime | None = None, | ||
| ) -> int: | ||
| """Delete memories that have exceeded their retention period. | ||
|
|
||
| Processes each category independently so that a failure in one | ||
| category does not prevent cleanup of the remaining categories. | ||
|
|
||
| Args: | ||
| agent_id: Agent whose memories to clean up. | ||
| now: Current time (defaults to UTC now). | ||
|
|
||
| Returns: | ||
| Number of expired memories deleted. | ||
| """ | ||
| if now is None: | ||
| now = datetime.now(UTC) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider validating timezone awareness for the now parameter.
If a caller passes a naive datetime (without timezone info), the cutoff calculation at line 96 will produce a naive datetime that gets compared against created_at timestamps which may be timezone-aware. This could lead to incorrect retention behavior.
🛡️ Proposed defensive check
if now is None:
now = datetime.now(UTC)
+ elif now.tzinfo is None:
+ msg = "now must be timezone-aware"
+ raise ValueError(msg)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/memory/consolidation/retention.py` around lines 71 - 89, The
cleanup_expired function accepts a now parameter that may be a naive datetime,
leading to unsafe comparisons with timezone-aware created_at values; add a
defensive validation at the start of cleanup_expired to ensure now is
timezone-aware (or convert naive datetimes to UTC) before computing cutoffs: if
now is None set it to datetime.now(UTC) as before, otherwise check now.tzinfo
and either attach UTC or raise a ValueError, and then use that tz-aware now for
subsequent cutoff and comparisons (referencing cleanup_expired, now, UTC, and
created_at in the retention logic).
| query = MemoryQuery( | ||
| categories=frozenset({category}), | ||
| until=cutoff, | ||
| limit=1000, | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider documenting or making the batch limit configurable.
The hardcoded limit=1000 means that if an agent has more than 1,000 expired memories in a single category, not all will be cleaned up in one invocation. This batching behavior may be intentional for performance reasons, but it's not documented in the docstring.
📝 Suggested docstring update
async def cleanup_expired(
self,
agent_id: NotBlankStr,
now: datetime | None = None,
) -> int:
"""Delete memories that have exceeded their retention period.
Processes each category independently so that a failure in one
category does not prevent cleanup of the remaining categories.
+
+ Note:
+ Retrieves up to 1000 expired memories per category per run.
+ Subsequent runs will continue cleanup if more remain.
Args:
agent_id: Agent whose memories to clean up.
now: Current time (defaults to UTC now).
Returns:
Number of expired memories deleted.
"""📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| query = MemoryQuery( | |
| categories=frozenset({category}), | |
| until=cutoff, | |
| limit=1000, | |
| ) | |
| async def cleanup_expired( | |
| self, | |
| agent_id: NotBlankStr, | |
| now: datetime | None = None, | |
| ) -> int: | |
| """Delete memories that have exceeded their retention period. | |
| Processes each category independently so that a failure in one | |
| category does not prevent cleanup of the remaining categories. | |
| Note: | |
| Retrieves up to 1000 expired memories per category per run. | |
| Subsequent runs will continue cleanup if more remain. | |
| Args: | |
| agent_id: Agent whose memories to clean up. | |
| now: Current time (defaults to UTC now). | |
| Returns: | |
| Number of expired memories deleted. | |
| """ |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/memory/consolidation/retention.py` around lines 97 - 101, The
MemoryQuery call hardcodes limit=1000 which imposes an undocumented batching
behavior; update the enclosing function's signature (the retention/consolidation
function that builds MemoryQuery) to accept an optional batch_limit (default
1000) and use that value in MemoryQuery(limit=batch_limit), and also update the
function docstring to document the batching behavior and the meaning of
batch_limit (e.g., that a single invocation will only process up to batch_limit
expired memories per category and callers can increase or loop until none
remain); reference MemoryQuery, the local variables category and cutoff, and the
new batch_limit parameter in the docstring.
| result = await self._strategy.consolidate( | ||
| entries, | ||
| agent_id=agent_id, | ||
| ) | ||
|
|
||
| if self._archival_store is not None and self._config.archival.enabled: | ||
| archived = await self._archive_entries( | ||
| agent_id, | ||
| entries, | ||
| result.removed_ids, | ||
| ) | ||
| result = ConsolidationResult( | ||
| consolidated_count=result.consolidated_count, | ||
| removed_ids=result.removed_ids, | ||
| summary_id=result.summary_id, | ||
| archived_count=archived, | ||
| ) |
There was a problem hiding this comment.
Archive before the strategy removes hot memories.
run_consolidation() invokes the strategy first and only archives afterward. With the current strategy contract, that means originals may already be deleted by the time _archive_entries() runs. If archival then fails, the maintenance call raises after destructive work, and the removed memories are gone from both hot and cold storage.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/ai_company/memory/consolidation/service.py` around lines 90 - 106, The
consolidation flow archives after calling self._strategy.consolidate(), which
can cause originals to be deleted before archival and lead to data loss if
archiving fails; change the ordering so that when self._archival_store is set
and self._config.archival.enabled is true you call
self._archive_entries(agent_id, entries, result_candidate_ids) before invoking
self._strategy.consolidate() (or alternately have the consolidation strategy
return the list of candidate removed IDs without performing deletions so you can
archive them first), then only call self._strategy.consolidate() to perform
destructive removals and construct the ConsolidationResult with archived_count
from the earlier archive step; use the existing symbols _archive_entries,
_strategy.consolidate, and ConsolidationResult to implement this ordering
change.
| async def archive(self, entry: ArchivalEntry) -> str: | ||
| self._next_id += 1 | ||
| archive_id = f"archive-{self._next_id}" | ||
| self._entries[archive_id] = entry | ||
| return archive_id | ||
|
|
||
| async def search(self, query: object) -> tuple[ArchivalEntry, ...]: | ||
| return tuple(self._entries.values()) | ||
|
|
||
| async def restore(self, entry_id: str) -> ArchivalEntry | None: | ||
| return self._entries.get(entry_id) | ||
|
|
||
| async def count(self, agent_id: str) -> int: | ||
| return sum(1 for e in self._entries.values() if e.agent_id == agent_id) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Mock return types should match protocol signatures.
The mock methods use str for parameters and return types, but the ArchivalStore protocol (from relevant snippets) specifies NotBlankStr. While this works at runtime due to structural typing, aligning the mock signatures with the protocol improves type safety and documentation clarity.
♻️ Suggested type alignment
- async def archive(self, entry: ArchivalEntry) -> str:
+ async def archive(self, entry: ArchivalEntry) -> NotBlankStr:
self._next_id += 1
- archive_id = f"archive-{self._next_id}"
+ archive_id = NotBlankStr(f"archive-{self._next_id}")
self._entries[archive_id] = entry
return archive_idNote: This would require importing NotBlankStr from the types module.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def archive(self, entry: ArchivalEntry) -> str: | |
| self._next_id += 1 | |
| archive_id = f"archive-{self._next_id}" | |
| self._entries[archive_id] = entry | |
| return archive_id | |
| async def search(self, query: object) -> tuple[ArchivalEntry, ...]: | |
| return tuple(self._entries.values()) | |
| async def restore(self, entry_id: str) -> ArchivalEntry | None: | |
| return self._entries.get(entry_id) | |
| async def count(self, agent_id: str) -> int: | |
| return sum(1 for e in self._entries.values() if e.agent_id == agent_id) | |
| async def archive(self, entry: ArchivalEntry) -> NotBlankStr: | |
| self._next_id += 1 | |
| archive_id = f"archive-{self._next_id}" | |
| self._entries[archive_id] = entry | |
| return archive_id | |
| async def search(self, query: object) -> tuple[ArchivalEntry, ...]: | |
| return tuple(self._entries.values()) | |
| async def restore(self, entry_id: str) -> ArchivalEntry | None: | |
| return self._entries.get(entry_id) | |
| async def count(self, agent_id: str) -> int: | |
| return sum(1 for e in self._entries.values() if e.agent_id == agent_id) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/memory/consolidation/test_archival.py` around lines 23 - 36, The
mock methods archive, search, restore, and count in the test mock should use the
same NotBlankStr types as the ArchivalStore protocol: change parameter and
return type annotations to use NotBlankStr where appropriate (e.g.,
archive(self, entry: ArchivalEntry) -> NotBlankStr or archive ID returns
NotBlankStr, restore(self, entry_id: NotBlankStr) -> ArchivalEntry | None,
count(self, agent_id: NotBlankStr) -> int) and update search signature if the
protocol expects different types; import NotBlankStr from the types module and
adjust signatures for archive, search, restore, and count to match the protocol
exactly.
| async def test_row_parse_error_wraps_in_query_error(self) -> None: | ||
| malformed_row = { | ||
| "id": "f1", | ||
| "content": "test", | ||
| "category": "INVALID_CATEGORY", | ||
| "author_agent_id": None, | ||
| "author_seniority": None, | ||
| "author_is_human": 1, | ||
| "created_at": _NOW.isoformat(), | ||
| "version": 1, | ||
| } | ||
| mock_row = AsyncMock() | ||
| mock_row.__getitem__ = lambda self, key: malformed_row[key] | ||
| with pytest.raises(OrgMemoryQueryError, match="Failed to deserialize"): | ||
| _row_to_org_fact(mock_row) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider using a proper mock Row object.
The test creates a dict-based mock for _row_to_org_fact, but AsyncMock might not behave exactly like a sqlite3.Row or aiosqlite.Row. The __getitem__ lambda should work, but consider using a more faithful mock or a named tuple for clearer test intent.
💡 Alternative using SimpleNamespace or dict subclass
async def test_row_parse_error_wraps_in_query_error(self) -> None:
- malformed_row = {
+ # Use a class that supports dict-like access
+ class MockRow(dict):
+ def __getitem__(self, key: str) -> object:
+ return super().__getitem__(key)
+
+ malformed_row = MockRow({
"id": "f1",
"content": "test",
"category": "INVALID_CATEGORY",
"author_agent_id": None,
"author_seniority": None,
"author_is_human": 1,
"created_at": _NOW.isoformat(),
"version": 1,
- }
- mock_row = AsyncMock()
- mock_row.__getitem__ = lambda self, key: malformed_row[key]
+ })
with pytest.raises(OrgMemoryQueryError, match="Failed to deserialize"):
- _row_to_org_fact(mock_row)
+ _row_to_org_fact(malformed_row)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def test_row_parse_error_wraps_in_query_error(self) -> None: | |
| malformed_row = { | |
| "id": "f1", | |
| "content": "test", | |
| "category": "INVALID_CATEGORY", | |
| "author_agent_id": None, | |
| "author_seniority": None, | |
| "author_is_human": 1, | |
| "created_at": _NOW.isoformat(), | |
| "version": 1, | |
| } | |
| mock_row = AsyncMock() | |
| mock_row.__getitem__ = lambda self, key: malformed_row[key] | |
| with pytest.raises(OrgMemoryQueryError, match="Failed to deserialize"): | |
| _row_to_org_fact(mock_row) | |
| async def test_row_parse_error_wraps_in_query_error(self) -> None: | |
| # Use a class that supports dict-like access | |
| class MockRow(dict): | |
| def __getitem__(self, key: str) -> object: | |
| return super().__getitem__(key) | |
| malformed_row = MockRow({ | |
| "id": "f1", | |
| "content": "test", | |
| "category": "INVALID_CATEGORY", | |
| "author_agent_id": None, | |
| "author_seniority": None, | |
| "author_is_human": 1, | |
| "created_at": _NOW.isoformat(), | |
| "version": 1, | |
| }) | |
| with pytest.raises(OrgMemoryQueryError, match="Failed to deserialize"): | |
| _row_to_org_fact(malformed_row) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unit/memory/org/test_store.py` around lines 335 - 349, The test uses
AsyncMock with a custom __getitem__ lambda to emulate a DB row for
_row_to_org_fact but AsyncMock may not faithfully mimic sqlite3.Row behavior;
replace the AsyncMock with a simple, synchronous row-like object (e.g., a dict
subclass, types.SimpleNamespace, or a custom lightweight class that implements
__getitem__) that returns values from malformed_row so _row_to_org_fact sees
realistic item access; update the test name
test_row_parse_error_wraps_in_query_error to use that row-like object instead of
AsyncMock and keep the same malformed_row and the pytest.raises assertion.
There was a problem hiding this comment.
Code Review
This pull request introduces significant new functionality for shared organizational memory and memory consolidation, including new modules for org memory storage, access control, configuration, a hybrid backend, and a complete memory consolidation pipeline with retention, archival, and pluggable strategies. While the code is well-structured and has comprehensive test coverage, a high-severity prompt injection vulnerability was identified in the system prompt template, where untrusted organizational policies are rendered without delimiters or sanitization. This must be addressed to prevent hijacking of agent behavior. Additionally, there are opportunities for performance optimization in version calculation and a minor issue with incorrect exception handling syntax.
| {% if org_policies %} | ||
| ## Organizational Policies | ||
|
|
||
| These are company-wide rules that must always be followed: | ||
| {% for policy in org_policies %} | ||
| - {{ policy }} | ||
| {% endfor %} | ||
|
|
||
| {% endif %} |
There was a problem hiding this comment.
Prompt Injection Risk
The system prompt template directly incorporates organizational policies without any sanitization or delimiters. An attacker who can influence these policies (e.g., by writing a malicious organizational fact) can perform a prompt injection attack to hijack the behavior of all agents in the company. Since agents have access to tools and authority, this is a high-severity risk.
Remediation:
Use clear delimiters around each policy and update the instructions to the agent to treat the content within those delimiters as data.
| {% if org_policies %} | |
| ## Organizational Policies | |
| These are company-wide rules that must always be followed: | |
| {% for policy in org_policies %} | |
| - {{ policy }} | |
| {% endfor %} | |
| {% endif %} | |
| {% if org_policies %} | |
| ## Organizational Policies | |
| These are company-wide rules that must always be followed. Each policy is enclosed in [POLICY START] and [POLICY END] delimiters: | |
| {% for policy in org_policies %} | |
| - [POLICY START] {{ policy }} [POLICY END] | |
| {% endfor %} | |
| {% endif %} |
src/ai_company/memory/org/store.py
Outdated
| if self._db is not None: | ||
| try: | ||
| await self._db.close() | ||
| except sqlite3.Error, OSError: |
There was a problem hiding this comment.
This except syntax with a comma is from Python 2 and will cause a SyntaxError in Python 3. To catch multiple exceptions, they should be grouped in a tuple.
except (sqlite3.Error, OSError):
src/ai_company/memory/org/store.py
Outdated
| return | ||
| try: | ||
| await self._db.close() | ||
| except sqlite3.Error, OSError: |
There was a problem hiding this comment.
This except syntax with a comma is from Python 2 and will cause a SyntaxError in Python 3. To catch multiple exceptions, they should be grouped in a tuple.
| except sqlite3.Error, OSError: | |
| except (sqlite3.Error, OSError): |
| async def _compute_next_version( | ||
| self, | ||
| category: OrgFactCategory, | ||
| ) -> int: | ||
| """Compute the next version number for facts in this category. | ||
|
|
||
| Uses ``MAX(version)`` query via the store to determine the next | ||
| version. Note: concurrent writers in the same category may | ||
| produce duplicate versions — callers requiring strict uniqueness | ||
| should wrap the version+save in a transaction. | ||
|
|
||
| Args: | ||
| category: The fact category. | ||
|
|
||
| Returns: | ||
| Next version number (max existing + 1, or 1 if none). | ||
| """ | ||
| existing = await self._store.list_by_category(category) | ||
| if not existing: | ||
| return 1 | ||
| return max(f.version for f in existing) + 1 |
There was a problem hiding this comment.
The current implementation of _compute_next_version fetches all existing facts for a category just to find the maximum version number. This can be inefficient if a category contains a large number of facts, as it transfers all fact data and loads it into memory.
To optimize this, I recommend adding a new method to the OrgFactStore protocol, like get_max_version(category: OrgFactCategory) -> int | None, and implementing it in SQLiteOrgFactStore with a SELECT MAX(version) FROM org_facts WHERE category = ? query. This would be much more efficient as the database would only return a single number.
Then, this method could be simplified to:
async def _compute_next_version(
self,
category: OrgFactCategory,
) -> int:
max_version = await self._store.get_max_version(category)
return (max_version or 0) + 1There was a problem hiding this comment.
Pull request overview
This PR adds two major features to the memory subsystem: shared organizational memory (memory/org/) and memory consolidation/archival (memory/consolidation/). Organizational memory provides a protocol-based OrgMemoryBackend with a hybrid prompt+retrieval backend, seniority-based write access control, and SQLite-backed fact storage. Memory consolidation adds retention enforcement, a pluggable consolidation strategy, archival to cold storage, and a maintenance orchestrator. Core policies are injected into agent system prompts via a new Jinja2 template section.
Changes:
- New
memory/org/package withOrgMemoryBackendprotocol,HybridPromptRetrievalBackend,SQLiteOrgFactStore, access control, config, errors, factory, and domain models - New
memory/consolidation/package withConsolidationStrategyprotocol,SimpleConsolidationStrategy,RetentionEnforcer,ArchivalStoreprotocol,MemoryConsolidationService, and config/models - System prompt integration (
org_policiesparameter inbuild_system_prompt), trimmable under token budget;RootConfigextended withorg_memoryfield
Reviewed changes
Copilot reviewed 48 out of 50 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/ai_company/memory/org/store.py |
OrgFactStore protocol + SQLiteOrgFactStore implementation with path traversal protection and LIKE escaping |
src/ai_company/memory/org/protocol.py |
OrgMemoryBackend runtime-checkable protocol |
src/ai_company/memory/org/hybrid_backend.py |
Hybrid prompt+retrieval backend combining static policies with dynamic store |
src/ai_company/memory/org/models.py |
OrgFact, OrgFactAuthor, OrgFactWriteRequest, OrgMemoryQuery frozen Pydantic models |
src/ai_company/memory/org/access_control.py |
Seniority-based write access control with MappingProxyType immutability |
src/ai_company/memory/org/config.py |
OrgMemoryConfig, ExtendedStoreConfig with backend validation |
src/ai_company/memory/org/errors.py |
OrgMemoryError hierarchy |
src/ai_company/memory/org/factory.py |
create_org_memory_backend() factory function |
src/ai_company/memory/org/__init__.py |
Package re-exports |
src/ai_company/memory/consolidation/service.py |
MemoryConsolidationService orchestrating retention→consolidation→max enforcement |
src/ai_company/memory/consolidation/simple_strategy.py |
SimpleConsolidationStrategy grouping by category |
src/ai_company/memory/consolidation/retention.py |
RetentionEnforcer with per-category age-based cleanup |
src/ai_company/memory/consolidation/models.py |
ConsolidationResult, ArchivalEntry, RetentionRule models |
src/ai_company/memory/consolidation/config.py |
ConsolidationConfig, RetentionConfig, ArchivalConfig |
src/ai_company/memory/consolidation/archival.py |
ArchivalStore protocol |
src/ai_company/memory/consolidation/strategy.py |
ConsolidationStrategy protocol |
src/ai_company/memory/consolidation/__init__.py |
Package re-exports |
src/ai_company/memory/__init__.py |
Updated exports with org memory and consolidation types |
src/ai_company/memory/config.py |
Added consolidation field to CompanyMemoryConfig |
src/ai_company/engine/prompt.py |
Added org_policies parameter threading through prompt building |
src/ai_company/engine/prompt_template.py |
Added org_policies Jinja2 template section, bumped version to 1.2.0 |
src/ai_company/core/enums.py |
Added OrgFactCategory enum |
src/ai_company/config/schema.py |
Added org_memory field to RootConfig |
src/ai_company/config/defaults.py |
Added org_memory to defaults |
src/ai_company/observability/events/org_memory.py |
Org memory event constants |
src/ai_company/observability/events/consolidation.py |
Consolidation event constants |
tests/unit/memory/org/* |
9 test files covering store, protocol, backend, models, config, access control, errors, factory, prompt |
tests/unit/memory/consolidation/* |
6 test files covering strategy, service, retention, models, config, archival |
tests/unit/engine/test_prompt.py |
Updated template version assertion |
tests/unit/observability/test_events.py |
Added new event module names |
tests/unit/memory/test_init.py |
Updated __all__ assertions |
tests/unit/config/conftest.py |
Added OrgMemoryConfig to factory |
README.md |
Updated feature descriptions |
DESIGN_SPEC.md |
Updated file tree and conventions table |
CLAUDE.md |
Updated memory directory description and event naming examples |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
src/ai_company/memory/org/store.py
Outdated
| if self._db is not None: | ||
| try: | ||
| await self._db.close() | ||
| except sqlite3.Error, OSError: |
There was a problem hiding this comment.
except sqlite3.Error, OSError: uses Python 2-style exception syntax which is a SyntaxError in Python 3. This should be except (sqlite3.Error, OSError): (with parentheses), consistent with line 259 in this same file and with src/ai_company/persistence/sqlite/backend.py:93,102,129.
src/ai_company/memory/org/store.py
Outdated
| return | ||
| try: | ||
| await self._db.close() | ||
| except sqlite3.Error, OSError: |
There was a problem hiding this comment.
except sqlite3.Error, OSError: uses Python 2-style exception syntax which is a SyntaxError in Python 3. This should be except (sqlite3.Error, OSError): (with parentheses), consistent with line 259 in this same file and with src/ai_company/persistence/sqlite/backend.py:93,102,129.
) Implement organizational memory subsystem with access-controlled fact storage (SQLite-backed), hybrid prompt+retrieval backend, and system prompt integration for company-wide policies. Add memory consolidation pipeline with simple grouping strategy, per-category retention enforcement, archival store protocol, and orchestrating service for periodic maintenance. Closes #125 Closes #48
Pre-reviewed by 10 agents, 70 findings addressed: - Hardened error handling: try/except + structured logging on all error paths - Security: path traversal validation, LIKE escape, MappingProxyType for access rules - Type safety: NotBlankStr for all identifiers, allow_inf_nan=False on all models - Validation: ArchivalEntry temporal order, RetentionConfig unique categories - Resilience: per-category retention cleanup continues on failure - Prompt: org_policies before autonomy, trimmed last not second - Docs: updated CLAUDE.md, README.md, DESIGN_SPEC.md for new modules - Tests: 25 new tests across 9 files (4870 total, 96% coverage)
…iewers (#125, #48) - Archival resilience: continue on per-entry failure instead of aborting - Fail-closed access control for unknown category rules - Batch loop for enforce_max_memories (MemoryQuery.limit cap at 1000) - INSERT instead of INSERT OR REPLACE for append-only audit trail - consolidated_count as @computed_field derived from removed_ids - ArchivalStore protocol signatures with explicit agent_id params - Relevance-based query ordering for text searches - Prompt injection warning and multiline policy sanitization - Event constant tests for consolidation and org_memory modules - Documentation updates: DESIGN_SPEC.md §1.4, §7.4, §7.5; README.md - Test coverage: archival resilience, retention continuity, error propagation, connect failure, relevance tiebreaker, event constants
_trim_sections now returns the final (content, estimated) tuple so _render_with_trimming can reuse it instead of rendering a second time with identical inputs.
bf15761 to
73f4b63
Compare
| for entry in to_remove: | ||
| await self._backend.delete(agent_id, entry.id) | ||
| removed_ids.append(entry.id) |
There was a problem hiding this comment.
delete() return value is ignored—falsely marks entries as removed.
self._backend.delete(agent_id, entry.id) returns bool per the MemoryBackend protocol (True = deleted, False = not found). The return value is never inspected, so entry.id is unconditionally appended to removed_ids even when the backend reports the entry was not present (e.g., race condition with another concurrent delete).
This has two cascading consequences:
ConsolidationResult.removed_idsover-counts what was actually removed, makingconsolidated_countincorrect.- When archival is enabled,
_archive_entriesiteratesremoved_idsand archives every entry in that set. An entry that was never deleted from hot storage will therefore appear in both hot storage and cold storage simultaneously, and will be re-processed on the next consolidation run, creating duplicate summary entries over time.
Compare with RetentionEnforcer.cleanup_expired(), which correctly handles the bool return:
deleted = await self._backend.delete(agent_id, entry.id)
if deleted:
total_deleted += 1
else:
logger.debug(RETENTION_DELETE_SKIPPED, ...)| for entry in to_remove: | |
| await self._backend.delete(agent_id, entry.id) | |
| removed_ids.append(entry.id) | |
| deleted = await self._backend.delete(agent_id, entry.id) | |
| if deleted: | |
| removed_ids.append(entry.id) |
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/ai_company/memory/consolidation/simple_strategy.py
Line: 116-118
Comment:
`delete()` return value is ignored—falsely marks entries as removed.
`self._backend.delete(agent_id, entry.id)` returns `bool` per the `MemoryBackend` protocol (`True` = deleted, `False` = not found). The return value is never inspected, so `entry.id` is unconditionally appended to `removed_ids` even when the backend reports the entry was not present (e.g., race condition with another concurrent delete).
This has two cascading consequences:
1. `ConsolidationResult.removed_ids` over-counts what was actually removed, making `consolidated_count` incorrect.
2. When archival is enabled, `_archive_entries` iterates `removed_ids` and archives every entry in that set. An entry that was never deleted from hot storage will therefore appear in **both** hot storage and cold storage simultaneously, and will be re-processed on the next consolidation run, creating duplicate summary entries over time.
Compare with `RetentionEnforcer.cleanup_expired()`, which correctly handles the `bool` return:
```python
deleted = await self._backend.delete(agent_id, entry.id)
if deleted:
total_deleted += 1
else:
logger.debug(RETENTION_DELETE_SKIPPED, ...)
```
```suggestion
deleted = await self._backend.delete(agent_id, entry.id)
if deleted:
removed_ids.append(entry.id)
```
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Pull request overview
Copilot reviewed 48 out of 50 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| class OrgFactCategory(StrEnum): | ||
| """Category of organizational fact (§7.4). | ||
|
|
||
| Categorizes shared organizational knowledge entries by their nature | ||
| and purpose within the company. | ||
| """ | ||
|
|
||
| CORE_POLICY = "core_policy" | ||
| ADR = "adr" | ||
| PROCEDURE = "procedure" | ||
| CONVENTION = "convention" | ||
|
|
There was a problem hiding this comment.
OrgFactCategory is defined in core/enums.py but is not re-exported from core/__init__.py, unlike every other enum in the module (e.g., ConsolidationInterval, MemoryCategory, SeniorityLevel, etc. are all imported and listed in __all__). While current consumers import directly from core.enums, this breaks the pattern established by the rest of the codebase. Add OrgFactCategory to the import list and __all__ in core/__init__.py for consistency.
Summary
memory/org/): Protocol-based organizational fact storage with SQLite backend, seniority-based write access control, hybrid prompt+retrieval backend, and factory for config-driven instantiationmemory/consolidation/): Retention enforcement, pluggable consolidation strategy (simple category-grouping included), archival to cold storage, and full maintenance orchestration (retention → consolidation → max-memories enforcement)Closes #125, closes #48
Test plan
Review coverage
Pre-reviewed by 10 agents: code-reviewer, python-reviewer, pr-test-analyzer, silent-failure-hunter, comment-analyzer, type-design-analyzer, logging-audit, resilience-audit, security-reviewer, docs-consistency. 70 findings triaged and implemented + polish pass by code-simplifier.
🤖 Generated with Claude Code