diff --git a/datamind/capabilities/db/base.py b/datamind/capabilities/db/base.py index d0bfeab..0ebfc32 100644 --- a/datamind/capabilities/db/base.py +++ b/datamind/capabilities/db/base.py @@ -106,6 +106,8 @@ async def execute_readonly( ) -> QueryResult: if not sql or not sql.strip(): raise CapabilityError("db", "Empty SQL") + if row_limit < 1: + raise CapabilityError("db", "row_limit must be at least 1") if contains_multiple_statements(sql): raise MultiStatementSQLError( "multiple statements are not allowed (use a single SELECT)" diff --git a/datamind/capabilities/hooks/audit.py b/datamind/capabilities/hooks/audit.py index b1b69dd..268bd0e 100644 --- a/datamind/capabilities/hooks/audit.py +++ b/datamind/capabilities/hooks/audit.py @@ -16,9 +16,9 @@ Secret redaction: arg values whose KEY matches a redaction regex (`api_key`, `password`, `token`, `secret`, ...) are replaced with -`"[REDACTED]"` before logging. Values are not scanned (we trust the -caller to put secrets in well-named fields, which is true across -DataMind's own tools). +`"[REDACTED]"` before logging. Credential-shaped values in exception +text are redacted as well because provider and transport errors commonly +include the rejected request detail. Concurrency: writes are serialised via an asyncio.Lock per AuditLogHook instance. For multi-process deployments use one process per profile @@ -45,6 +45,13 @@ r"(?i)(api[_-]?key|password|passwd|token|secret|authorization|bearer|access[_-]?key|client[_-]?secret)" ) _REDACTED = "[REDACTED]" +_BEARER_TEXT_RE = re.compile(r"(?i)\bBearer\s+([^\s,;]+)") +_SECRET_TEXT_RE = re.compile( + r"(?i)([\"']?" + r"(?:api[_-]?key|password|passwd|token|secret|authorization|access[_-]?key|client[_-]?secret)" + r"[\"']?\s*[:=]\s*)([\"']?)([^\s,;}\"']+)([\"']?)" +) +_RAW_SECRET_RE = re.compile(r"\bsk-[A-Za-z0-9_-]{8,}\b") _HASH_HEX_LEN = 16 # truncated SHA-256 hex; 64 bits of collision resistance @@ -74,6 +81,16 @@ def _redact(value: Any) -> Any: return value +def _redact_text(value: str) -> str: + """Redact common credential formats embedded in diagnostic text.""" + value = _BEARER_TEXT_RE.sub(f"Bearer {_REDACTED}", value) + value = _SECRET_TEXT_RE.sub( + lambda match: f"{match.group(1)}{match.group(2)}{_REDACTED}{match.group(4)}", + value, + ) + return _RAW_SECRET_RE.sub(_REDACTED, value) + + def _decision_to_record(decision: HookDecision) -> dict[str, Any]: if isinstance(decision, Allow): return {"kind": "allow"} @@ -163,7 +180,7 @@ async def post_tool_use( "error": ( None if error is None - else f"{type(error).__name__}: {error}" + else _redact_text(f"{type(error).__name__}: {error}") ), } await self._append(record) diff --git a/datamind/capabilities/ingest/service.py b/datamind/capabilities/ingest/service.py index b04387e..124cd3c 100644 --- a/datamind/capabilities/ingest/service.py +++ b/datamind/capabilities/ingest/service.py @@ -725,6 +725,9 @@ async def kb_add_file(self, *, path: str, copy_to_profile: bool = True) -> dict[ )) if not chunks: + await self._upsert_chunks( + chunks, replace_sources={source, str(resolved)}, + ) return {"file": str(resolved), "chunks_added": 0, "note": "file was empty"} parsed_chunks = None @@ -818,8 +821,6 @@ async def _upsert_chunks( raise CapabilityError("ingest", "KB surface is disabled") provider = self._kb.embedding store = self._kb.vector_store - texts = [c.text for c in chunks] - vectors = await provider.embed_texts(texts) stale_ids: list[str] = [] if replace_sources: new_ids = {chunk.id for chunk in chunks} @@ -830,6 +831,13 @@ async def _upsert_chunks( } if recorded_sources & replace_sources and chunk_id not in new_ids: stale_ids.append(chunk_id) + if not chunks: + if stale_ids: + await store.delete(stale_ids) + await self._kb.record_incremental_ingest() + return + texts = [c.text for c in chunks] + vectors = await provider.embed_texts(texts) await store.add( ids=[c.id for c in chunks], texts=texts, @@ -882,12 +890,25 @@ async def db_import_csv( # Sanitise column names: same rule as table names. safe_cols: list[str] = [] - for raw in header: + used_cols: set[str] = set() + for index, raw in enumerate(header, start=1): col = raw.strip() if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,63}", col): # Fall back to col_ if header is unusable. - col = f"col_{len(safe_cols) + 1}" + col = f"col_{index}" + if col.casefold() in used_cols: + # SQL tables cannot contain duplicate column names. Keep the + # first header unchanged and give later occurrences a stable + # fallback name without losing their values in the row dict. + base = f"col_{index}" + col = base + suffix = 2 + while col.casefold() in used_cols: + suffix_text = f"_{suffix}" + col = f"{base[:64 - len(suffix_text)]}{suffix_text}" + suffix += 1 safe_cols.append(col) + used_cols.add(col.casefold()) rows: list[dict[str, str]] = [] for raw_row in reader: diff --git a/datamind/capabilities/memory/providers/sqlite_store.py b/datamind/capabilities/memory/providers/sqlite_store.py index b9cd08e..5dadad8 100644 --- a/datamind/capabilities/memory/providers/sqlite_store.py +++ b/datamind/capabilities/memory/providers/sqlite_store.py @@ -287,14 +287,32 @@ async def recall( session_id: str | None = None, top_k: int = 8, kinds: Sequence[str] | None = None, + scope_filter: Sequence[str] | None = None, include_archived: bool = False, # advanced — let callers tune the per-scope budget for ablation per_scope: dict[str, int] | None = None, ) -> list[MemoryItem]: + valid_scopes = {"session", "profile", "global"} + allowed_scopes = valid_scopes if scope_filter is None else set(scope_filter) + invalid_scopes = allowed_scopes - valid_scopes + if invalid_scopes: + raise CapabilityError( + "memory", f"unsupported scope filter: {sorted(invalid_scopes)!r}" + ) + if not allowed_scopes: + return [] + # Default per-scope budget: 2 (session) + 4 (profile) + 2 (global) = 8. - budgets: dict[str, int] = {"session": 2, "profile": 4, "global": 2} + budgets: dict[str, int] = { + "session": 2 if "session" in allowed_scopes else 0, + "profile": 4 if "profile" in allowed_scopes else 0, + "global": 2 if "global" in allowed_scopes else 0, + } if per_scope: - budgets.update({k: v for k, v in per_scope.items() if k in budgets}) + budgets.update({ + k: v for k, v in per_scope.items() + if k in budgets and k in allowed_scopes + }) # If no embedding, fall back to lexical scoring within the same # scope filters so the contract stays identical. diff --git a/datamind/capabilities/memory/service.py b/datamind/capabilities/memory/service.py index 882b858..078304d 100644 --- a/datamind/capabilities/memory/service.py +++ b/datamind/capabilities/memory/service.py @@ -92,14 +92,25 @@ async def recall( session_id: str | None = None, top_k: int = 8, kinds: Sequence[str] | None = None, + scope_filter: Sequence[str] | None = None, include_archived: bool = False, ) -> list[dict[str, Any]]: + allowed = None if scope_filter is None else set(scope_filter) + effective_profile = ( + profile or self._default_profile + if allowed is None or "profile" in allowed + else None + ) + effective_session = ( + session_id if allowed is None or "session" in allowed else None + ) hits = await self.long_term.recall( query, - profile=profile or self._default_profile, - session_id=session_id, + profile=effective_profile, + session_id=effective_session, top_k=top_k, kinds=kinds, + scope_filter=scope_filter, include_archived=include_archived, ) return [h.model_dump() for h in hits] diff --git a/datamind/capabilities/memory/tools.py b/datamind/capabilities/memory/tools.py index a328fa3..dc5eb50 100644 --- a/datamind/capabilities/memory/tools.py +++ b/datamind/capabilities/memory/tools.py @@ -82,6 +82,7 @@ async def _recall( session_id=eff_session, top_k=top_k, kinds=kinds, + scope_filter=scope_filter, ) return {"query": query, "count": len(hits), "results": hits} diff --git a/datamind/config.py b/datamind/config.py index 8f182d3..cc5e304 100644 --- a/datamind/config.py +++ b/datamind/config.py @@ -113,7 +113,7 @@ class DBConfig(BaseModel): dialect: str = "sqlite" # sqlite | mysql | postgres | ... dsn: str | None = None # e.g. mysql+pymysql://user:pw@host/db read_only: bool = True - row_limit: int = 1000 + row_limit: int = Field(default=1000, ge=1) query_timeout_s: float = 10.0 diff --git a/datamind/core/protocols.py b/datamind/core/protocols.py index e680186..c64c747 100644 --- a/datamind/core/protocols.py +++ b/datamind/core/protocols.py @@ -374,6 +374,7 @@ async def recall( session_id: str | None = None, top_k: int = 8, kinds: Sequence[str] | None = None, + scope_filter: Sequence[str] | None = None, include_archived: bool = False, ) -> list[MemoryItem]: ... diff --git a/datamind/tests/test_config.py b/datamind/tests/test_config.py index dab841e..e5c1d13 100644 --- a/datamind/tests/test_config.py +++ b/datamind/tests/test_config.py @@ -14,7 +14,7 @@ import pytest from pydantic import ValidationError -from datamind.config import DataConfig, Settings +from datamind.config import DBConfig, DataConfig, Settings def test_nested_env_hydrates_required_llm(monkeypatch, tmp_path): @@ -116,3 +116,8 @@ def test_profile_assignment_keeps_path_boundary(tmp_path): config = DataConfig(base_dir=tmp_path) with pytest.raises(ValidationError): config.profile = "../escape" + + +def test_db_config_rejects_nonpositive_row_limit(): + with pytest.raises(ValidationError): + DBConfig(row_limit=0) diff --git a/datamind/tests/test_db_safeguard.py b/datamind/tests/test_db_safeguard.py index bda0bd6..874043f 100644 --- a/datamind/tests/test_db_safeguard.py +++ b/datamind/tests/test_db_safeguard.py @@ -12,6 +12,8 @@ leading_verb, strip_comments, ) +from datamind.capabilities.db.providers.sqlite import SQLiteDialect +from datamind.core.errors import CapabilityError @pytest.mark.parametrize( @@ -67,3 +69,17 @@ def test_ensure_row_limit_strips_trailing_semicolon(): def test_strip_comments(): assert strip_comments("SELECT /* x */ 1") == "SELECT 1" assert strip_comments("-- line\nSELECT 1") == "\nSELECT 1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("row_limit", [0, -1]) +async def test_execute_readonly_rejects_nonpositive_row_limit(tmp_path, row_limit): + dialect = SQLiteDialect() + engine = dialect.build_engine( + None, default_path=str(tmp_path / "db.sqlite") + ) + + with pytest.raises(CapabilityError, match="row_limit"): + await dialect.execute_readonly( + engine, "SELECT 1", row_limit=row_limit + ) diff --git a/datamind/tests/test_hooks.py b/datamind/tests/test_hooks.py index 55b5567..6a72719 100644 --- a/datamind/tests/test_hooks.py +++ b/datamind/tests/test_hooks.py @@ -322,6 +322,48 @@ async def test_audit_log_redacts_secret_keys(tmp_path: Path): assert rec["args"]["username"] == "ann" +@pytest.mark.asyncio +async def test_audit_log_redacts_secrets_in_errors(tmp_path: Path): + audit = tmp_path / "audit.jsonl" + hook = AuditLogHook(audit_path=audit) + await hook.post_tool_use( + _ctx(), + "fake_tool", + {}, + result=None, + error=RuntimeError( + "request failed: api_key=sk-live-secret Authorization: Bearer bearer-secret" + ), + ) + + rec = json.loads(audit.read_text().splitlines()[-1]) + assert "sk-live-secret" not in rec["error"] + assert "bearer-secret" not in rec["error"] + assert "[REDACTED]" in rec["error"] + + +@pytest.mark.asyncio +async def test_audit_log_redacts_quoted_url_and_raw_secrets_in_errors(tmp_path: Path): + audit = tmp_path / "audit.jsonl" + hook = AuditLogHook(audit_path=audit) + secrets = ["sk-json-secret1", "sk-query-secret2", "sk-raw-secret3"] + await hook.post_tool_use( + _ctx(), + "fake_tool", + {}, + result=None, + error=RuntimeError( + '{"api_key": "sk-json-secret1"} ' + "https://example.test?access_token=sk-query-secret2 " + "provider rejected sk-raw-secret3" + ), + ) + + rec = json.loads(audit.read_text().splitlines()[-1]) + assert all(secret not in rec["error"] for secret in secrets) + assert rec["error"].count("[REDACTED]") >= 3 + + @pytest.mark.asyncio async def test_audit_log_records_denied_decision(tmp_path: Path): audit = tmp_path / "audit.jsonl" diff --git a/datamind/tests/test_ingest_revisions.py b/datamind/tests/test_ingest_revisions.py index bcf4ed7..926790d 100644 --- a/datamind/tests/test_ingest_revisions.py +++ b/datamind/tests/test_ingest_revisions.py @@ -7,6 +7,8 @@ class _Embedding: async def embed_texts(self, texts): + if not texts: + raise AssertionError("empty revisions must not call the embedding provider") return [[float(len(text))] for text in texts] @@ -43,6 +45,34 @@ async def generate_text(self, prompt, **kwargs): return "[]" +@pytest.mark.asyncio +async def test_reingesting_empty_path_removes_old_kb_chunks(tmp_path: Path): + profile = tmp_path / "profile" + uploads = profile / "uploads" + uploads.mkdir(parents=True) + source = uploads / "release_note.txt" + kb = _KB() + service = IngestService( + kb=kb, + db=None, + graph=None, + llm_client=_Model(), + llm_model="test", + profile_data_dir=profile, + chunk_size=512, + chunk_overlap=64, + ) + + source.write_text("old release note", encoding="utf-8") + await service.kb_add_file(path=str(source)) + source.write_text("\n\n", encoding="utf-8") + + result = await service.kb_add_file(path=str(source)) + + assert result["chunks_added"] == 0 + assert kb.vector_store.items == {} + + @pytest.mark.asyncio async def test_reingesting_same_path_replaces_old_kb_chunks(tmp_path: Path): profile = tmp_path / "profile" diff --git a/datamind/tests/test_memory.py b/datamind/tests/test_memory.py index 9b9c452..3a1b495 100644 --- a/datamind/tests/test_memory.py +++ b/datamind/tests/test_memory.py @@ -8,8 +8,10 @@ from datamind.capabilities.memory import ( MemoryService, ShortTermMemory, + build_memory_tools, ) from datamind.capabilities.memory.providers.sqlite_store import SQLiteMemoryStore +from datamind.core.errors import CapabilityError from datamind.core.protocols import MemoryStore @@ -276,6 +278,52 @@ async def test_service_combines_short_and_long(tmp_path): assert await svc.forget(rid) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope_filter", "expected"), + [ + (None, {"global", "profile", "session"}), + (["global"], {"global"}), + (["profile"], {"profile"}), + (["session"], {"session"}), + (["global", "profile"], {"global", "profile"}), + ([], set()), + ], +) +async def test_memory_recall_scope_filter_is_exact(tmp_path, scope_filter, expected): + lt = SQLiteMemoryStore(db_path=str(tmp_path / "m.db"), embedding=_FakeEmbed()) + svc = MemoryService( + short_term=ShortTermMemory(max_turns=5), + long_term=lt, + default_profile="acme", + ) + tools = {tool.name: tool.handler for tool in build_memory_tools(svc)} + + await tools["memory_save"]("global fact", scope="global") + await tools["memory_save"]("profile fact", scope="profile", profile="acme") + await tools["memory_save"]( + "session fact", scope="session", session_id="chat-1" + ) + + result = await tools["memory_recall"]( + "fact", + profile="acme", + session_id="chat-1", + scope_filter=scope_filter, + ) + assert {item["scope"] for item in result["results"]} == expected + + +@pytest.mark.asyncio +async def test_memory_recall_rejects_unknown_scope_filter(tmp_path): + lt = SQLiteMemoryStore(db_path=str(tmp_path / "m.db"), embedding=_FakeEmbed()) + svc = MemoryService( + short_term=ShortTermMemory(max_turns=5), long_term=lt, default_profile="acme" + ) + with pytest.raises(CapabilityError, match="unsupported scope filter"): + await svc.recall("fact", scope_filter=["other"]) + + @pytest.mark.asyncio async def test_service_default_profile_used_when_omitted(tmp_path): """save without explicit profile should land under default_profile.""" diff --git a/datamind/tests/test_replace_receipts.py b/datamind/tests/test_replace_receipts.py index ee4b2d6..34c56de 100644 --- a/datamind/tests/test_replace_receipts.py +++ b/datamind/tests/test_replace_receipts.py @@ -77,3 +77,33 @@ async def test_append_retry_remains_deduplicated(database, tmp_path): assert second["results"][0]["status"] == "unchanged" assert second["revision"] == first["revision"] assert (await db.query_sql("SELECT amount FROM sales")).rows == [["200"]] + + +@pytest.mark.asyncio +async def test_csv_import_keeps_duplicate_headers_as_distinct_columns(database, tmp_path): + db, raw = database + source = tmp_path / "duplicate-columns.csv" + source.write_text("a,a\nfirst,second\n", encoding="utf-8") + + receipt = await raw.get("db_import_csv").handler( + path=str(source), table="duplicate_columns", if_exists="replace" + ) + + assert receipt["columns"] == ["a", "col_2"] + assert (await db.query_sql("SELECT * FROM duplicate_columns")).rows == [["first", "second"]] + + +@pytest.mark.asyncio +async def test_csv_import_deduplicates_case_insensitive_headers(database, tmp_path): + db, raw = database + source = tmp_path / "case-columns.csv" + source.write_text("a,A,col_3,col_3\n1,2,3,4\n", encoding="utf-8") + + receipt = await raw.get("db_import_csv").handler( + path=str(source), table="case_columns", if_exists="replace" + ) + + assert receipt["columns"] == ["a", "col_2", "col_3", "col_4"] + assert (await db.query_sql("SELECT * FROM case_columns")).rows == [ + ["1", "2", "3", "4"] + ]