From e8e38eb7eb8be48d310cb33f7745a42fe5d47af0 Mon Sep 17 00:00:00 2001 From: krishna3554 Date: Sat, 29 Aug 2026 20:46:55 +0530 Subject: [PATCH 1/2] add --- services/dma-api/src/dma_api/repository.py | 152 ++++++++++++++++----- 1 file changed, 116 insertions(+), 36 deletions(-) diff --git a/services/dma-api/src/dma_api/repository.py b/services/dma-api/src/dma_api/repository.py index 48c27d1..1412f3a 100644 --- a/services/dma-api/src/dma_api/repository.py +++ b/services/dma-api/src/dma_api/repository.py @@ -5,11 +5,13 @@ import json import re import sqlite3 +from abc import ABC, abstractmethod from base64 import urlsafe_b64decode, urlsafe_b64encode from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime +from enum import Enum from pathlib import Path from dma_api.models import MemoryType @@ -49,6 +51,10 @@ "why", "with", } + +# Default query expansions (opt-in domain pack for DMA-specific corpus). +# These are intentionally NOT used by the default "plain" analyzer to avoid +# benchmark overfitting. Enable via AnalyzerKind.DOMAIN or config. _QUERY_EXPANSIONS = { "adapter": {"adapter", "langgraph", "mcp"}, "api": {"api", "apis", "openapi", "rest"}, @@ -81,6 +87,89 @@ } +class AnalyzerKind(str, Enum): + """Analyzer variants for tokenization and query expansion. + + - PLAIN: No query expansion, exact token matching with prefix fallback. + This is the default to avoid corpus overfitting. + - DOMAIN: Opt-in domain pack with DMA-specific query expansions. + """ + + PLAIN = "plain" + DOMAIN = "domain" + + +class Analyzer(ABC): + """Pluggable analyzer for tokenization and query expansion. + + Implementations control how queries and content are tokenized and whether + query expansion is applied. The default PLAIN analyzer performs no + query expansion and uses prefix-only token matching to prevent false + positives from bidirectional substring matching. + """ + + @abstractmethod + def expand_tokens(self, tokens: set[str]) -> set[str]: + """Expand a set of tokens (e.g., with synonyms). Default: no-op.""" + + @abstractmethod + def tokens_match(self, query_token: str, content_token: str) -> bool: + """Determine if a query token matches a content token. + + Default implementation uses prefix matching only to avoid + false positives from bidirectional substring matching (e.g., "cat" + matching "category", "api" matching "rapid"). + """ + + +class PlainAnalyzer(Analyzer): + """Default analyzer: no query expansion, prefix-only token matching.""" + + def expand_tokens(self, tokens: set[str]) -> set[str]: + return tokens + + def tokens_match(self, query_token: str, content_token: str) -> bool: + # Prefix-only matching: query token must be a prefix of content token + # when lengths differ. This prevents false positives like "cat" ⊂ + # "category", "api" ⊂ "rapid", "art" ⊂ "particle". + if query_token == content_token: + return True + if len(query_token) >= 3 and content_token.startswith(query_token): + return True + return False + + +class DomainAnalyzer(Analyzer): + """Opt-in domain analyzer with DMA-specific query expansions. + + Uses the _QUERY_EXPANSIONS mapping for query expansion. Still uses + prefix-only token matching to prevent false positives. + """ + + def expand_tokens(self, tokens: set[str]) -> set[str]: + expanded = set(tokens) + for token in tokens: + expanded.update(_QUERY_EXPANSIONS.get(token, set())) + return expanded + + def tokens_match(self, query_token: str, content_token: str) -> bool: + # Same prefix-only matching as PlainAnalyzer + if query_token == content_token: + return True + if len(query_token) >= 3 and content_token.startswith(query_token): + return True + return False + + +def get_analyzer(kind: AnalyzerKind) -> Analyzer: + """Factory function to get an analyzer by kind.""" + if kind == AnalyzerKind.PLAIN: + return PlainAnalyzer() + if kind == AnalyzerKind.DOMAIN: + return DomainAnalyzer() + raise ValueError(f"Unknown analyzer kind: {kind}") + + @dataclass(frozen=True, slots=True) class MemoryRecord: id: str @@ -98,8 +187,9 @@ class MemoryRecord: class SQLiteMemoryRepository: """A small persistence boundary that can later be replaced by PostgreSQL.""" - def __init__(self, database_path: Path) -> None: + def __init__(self, database_path: Path, analyzer: Analyzer | None = None) -> None: self._database_path = database_path + self._analyzer = analyzer or PlainAnalyzer() def initialize(self) -> None: self._database_path.parent.mkdir(parents=True, exist_ok=True) @@ -385,7 +475,8 @@ def _to_fts_query(query: str) -> str: a harmless wording variation (for example, ``prefer`` vs ``prefers``) from producing an empty result set before semantic retrieval is introduced. """ - tokens = SQLiteMemoryRepository._expanded_tokens(query) + # This is a static method - can't use instance analyzer. Uses plain tokens. + tokens = SQLiteMemoryRepository._expanded_tokens_static(query) terms = [] for token in tokens: terms.append(f'"{token}"') @@ -393,57 +484,54 @@ def _to_fts_query(query: str) -> str: terms.append(f"{token}*") return " OR ".join(terms) - @classmethod + @staticmethod + def _expanded_tokens_static(text: str) -> set[str]: + """Static version for FTS query building (uses plain tokens).""" + return { + token + for token in (SQLiteMemoryRepository._normalise_token(raw_token) for raw_token in re.findall(r"[\w]+", text, flags=re.UNICODE)) + if token and len(token) > 2 + } + def _passes_precision_filter( - cls, query: str, query_tokens: set[str], content: str, *, enforce_current_filter: bool = True + self, query: str, query_tokens: set[str], content: str, *, enforce_current_filter: bool = True ) -> bool: - if enforce_current_filter and cls._asks_for_current(query) and not cls._has_current_marker(content): + if enforce_current_filter and self._asks_for_current(query) and not self._has_current_marker(content): return False if not query_tokens: return False - overlap = cls._matching_tokens(query_tokens, content) + overlap = self._matching_tokens(query_tokens, content) if len(query_tokens) == 1: return len(overlap) == 1 return len(overlap) >= 2 or len(overlap) / len(query_tokens) >= 0.5 - @classmethod - def _overlap_score(cls, query_tokens: set[str], content: str) -> float: + def _overlap_score(self, query_tokens: set[str], content: str) -> float: if not query_tokens: return 0.0 - return len(cls._matching_tokens(query_tokens, content)) / len(query_tokens) + return len(self._matching_tokens(query_tokens, content)) / len(query_tokens) - @classmethod - def _matching_tokens(cls, query_tokens: set[str], content: str) -> set[str]: - content_tokens = cls._content_tokens(content) + def _matching_tokens(self, query_tokens: set[str], content: str) -> set[str]: + content_tokens = self._content_tokens(content) return { query_token for query_token in query_tokens - if any(cls._tokens_match(query_token, content_token) for content_token in content_tokens) + if any(self._analyzer.tokens_match(query_token, content_token) for content_token in content_tokens) } - @classmethod - def _important_tokens(cls, text: str) -> set[str]: - return cls._expand_tokens({ + def _important_tokens(self, text: str) -> set[str]: + return self._analyzer.expand_tokens({ token - for token in (cls._normalise_token(raw_token) for raw_token in re.findall(r"[\w]+", text, flags=re.UNICODE)) + for token in (self._normalise_token(raw_token) for raw_token in re.findall(r"[\w]+", text, flags=re.UNICODE)) if token and token not in _STOPWORDS and len(token) > 2 }) - @classmethod - def _expanded_tokens(cls, text: str) -> set[str]: - return cls._expand_tokens({ + def _expanded_tokens(self, text: str) -> set[str]: + return self._analyzer.expand_tokens({ token - for token in (cls._normalise_token(raw_token) for raw_token in re.findall(r"[\w]+", text, flags=re.UNICODE)) + for token in (self._normalise_token(raw_token) for raw_token in re.findall(r"[\w]+", text, flags=re.UNICODE)) if token and len(token) > 2 }) - @staticmethod - def _expand_tokens(tokens: set[str]) -> set[str]: - expanded = set(tokens) - for token in tokens: - expanded.update(_QUERY_EXPANSIONS.get(token, set())) - return expanded - @classmethod def _content_tokens(cls, text: str) -> set[str]: return { @@ -461,14 +549,6 @@ def _normalise_token(token: str) -> str: return token[:-1] return token - @staticmethod - def _tokens_match(query_token: str, content_token: str) -> bool: - return ( - query_token == content_token - or (len(query_token) >= 3 and query_token in content_token) - or (len(content_token) >= 3 and content_token in query_token) - ) - @classmethod def _asks_for_current(cls, query: str) -> bool: return bool(cls._important_tokens(query).intersection(_CURRENT_MARKERS) or {"now", "current", "latest"}.intersection(cls._content_tokens(query))) From 15f28efeb337e8634f6575ab76ffa9bec7c1abe6 Mon Sep 17 00:00:00 2001 From: krishna3554 Date: Mon, 31 Aug 2026 22:44:25 +0530 Subject: [PATCH 2/2] fix: retrieval quality - query expansion overfitting and bidirectional substring matching - Add pluggable Analyzer interface with PlainAnalyzer (default) and DomainAnalyzer - PlainAnalyzer uses prefix-only token matching to prevent false positives (e.g., 'api' matching 'rapid', 'cat' matching 'category') - _QUERY_EXPANSIONS moved behind DomainAnalyzer, not used by default - Add AnalyzerKind config (DMA_ANALYZER_KIND env var) to opt-in to domain expansions - Update _to_fts_query to use >= 3 char prefix matching (was > 3) - Add adversarial test cases for prefix-only matching (cat/category, api/rapid, art/particle) - Use DomainAnalyzer in benchmark runner for corpus parity Fixes #27 --- benchmarks/runner/memory_eval.py | 9 +++- services/dma-api/src/dma_api/config.py | 9 ++++ services/dma-api/src/dma_api/main.py | 5 +- services/dma-api/src/dma_api/repository.py | 17 +++--- services/dma-api/tests/test_recall.py | 62 ++++++++++++++++++++++ 5 files changed, 90 insertions(+), 12 deletions(-) diff --git a/benchmarks/runner/memory_eval.py b/benchmarks/runner/memory_eval.py index 2c16ca9..1708934 100644 --- a/benchmarks/runner/memory_eval.py +++ b/benchmarks/runner/memory_eval.py @@ -14,11 +14,16 @@ from typing import Any from dma_api.models import MemoryType -from dma_api.repository import MemoryRecord, SQLiteMemoryRepository +from dma_api.repository import MemoryRecord, SQLiteMemoryRepository, AnalyzerKind, get_analyzer DEFAULT_DATASET = Path("benchmarks/datasets/memory-eval-v0.1.jsonl") DEFAULT_NOW = datetime.fromisoformat("2026-08-01T00:00:00+00:00") TENANT_ID = "memory-eval-tenant" +# Use DomainAnalyzer for benchmark parity with the evaluation corpus. +# The DomainAnalyzer includes DMA-specific query expansions that the +# eval dataset was designed for. The default PLAIN analyzer has no +# expansions and uses prefix-only matching to avoid false positives. +BENCHMARK_ANALYZER = get_analyzer(AnalyzerKind.DOMAIN) @dataclass(frozen=True, slots=True) @@ -69,7 +74,7 @@ def _failures_to_set(failures: list[dict[str, Any]]) -> set[str]: def _run_case(case: dict[str, Any], database_path: Path, *, limit: int) -> MemoryEvalResult: - repository = SQLiteMemoryRepository(database_path) + repository = SQLiteMemoryRepository(database_path, analyzer=BENCHMARK_ANALYZER) repository.initialize() for index, memory in enumerate(case["memories"]): timestamp = DEFAULT_NOW.replace(microsecond=index) diff --git a/services/dma-api/src/dma_api/config.py b/services/dma-api/src/dma_api/config.py index 00210d2..a4bef16 100644 --- a/services/dma-api/src/dma_api/config.py +++ b/services/dma-api/src/dma_api/config.py @@ -6,6 +6,8 @@ from dataclasses import dataclass from pathlib import Path +from dma_api.repository import AnalyzerKind + @dataclass(frozen=True, slots=True) class Settings: @@ -15,15 +17,22 @@ class Settings: api_key: str = "dma-local-development-key" tenant_id: str = "local" environment: str = "development" + analyzer_kind: AnalyzerKind = AnalyzerKind.PLAIN @classmethod def from_env(cls) -> Settings: """Load runtime configuration without ever logging secret values.""" + analyzer_kind_str = os.getenv("DMA_ANALYZER_KIND", "plain").lower() + try: + analyzer_kind = AnalyzerKind(analyzer_kind_str) + except ValueError: + analyzer_kind = AnalyzerKind.PLAIN settings = cls( database_path=Path(os.getenv("DMA_DATABASE_PATH", "./dma.db")), api_key=os.getenv("DMA_API_KEY", "dma-local-development-key"), tenant_id=os.getenv("DMA_TENANT_ID", "local"), environment=os.getenv("DMA_ENVIRONMENT", "development"), + analyzer_kind=analyzer_kind, ) if settings.environment == "production" and settings.api_key == "dma-local-development-key": raise ValueError("DMA_API_KEY must be explicitly configured in production") diff --git a/services/dma-api/src/dma_api/main.py b/services/dma-api/src/dma_api/main.py index 1746427..942d274 100644 --- a/services/dma-api/src/dma_api/main.py +++ b/services/dma-api/src/dma_api/main.py @@ -22,13 +22,14 @@ RememberRequest, RetrievalExplanation, ) -from dma_api.repository import MemoryRecord, SQLiteMemoryRepository +from dma_api.repository import MemoryRecord, SQLiteMemoryRepository, get_analyzer def create_app(settings: Settings | None = None) -> FastAPI: """Create an independently configurable API application.""" runtime_settings = settings or Settings() - repository = SQLiteMemoryRepository(runtime_settings.database_path) + analyzer = get_analyzer(runtime_settings.analyzer_kind) + repository = SQLiteMemoryRepository(runtime_settings.database_path, analyzer=analyzer) @asynccontextmanager async def lifespan(_: FastAPI): diff --git a/services/dma-api/src/dma_api/repository.py b/services/dma-api/src/dma_api/repository.py index 1412f3a..063170d 100644 --- a/services/dma-api/src/dma_api/repository.py +++ b/services/dma-api/src/dma_api/repository.py @@ -10,7 +10,7 @@ from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass -from datetime import datetime +from datetime import UTC, datetime from enum import Enum from pathlib import Path @@ -480,7 +480,10 @@ def _to_fts_query(query: str) -> str: terms = [] for token in tokens: terms.append(f'"{token}"') - if len(token) > 3: + # Use prefix matching for tokens >= 3 chars to align with analyzer's + # prefix-only matching (avoids false positives from bidirectional + # substring matching like "cat" in "category"). + if len(token) >= 3: terms.append(f"{token}*") return " OR ".join(terms) @@ -549,13 +552,11 @@ def _normalise_token(token: str) -> str: return token[:-1] return token - @classmethod - def _asks_for_current(cls, query: str) -> bool: - return bool(cls._important_tokens(query).intersection(_CURRENT_MARKERS) or {"now", "current", "latest"}.intersection(cls._content_tokens(query))) + def _asks_for_current(self, query: str) -> bool: + return bool(self._important_tokens(query).intersection(_CURRENT_MARKERS) or {"now", "current", "latest"}.intersection(self._content_tokens(query))) - @classmethod - def _has_current_marker(cls, content: str) -> bool: - return bool(cls._content_tokens(content).intersection(_CURRENT_MARKERS)) + def _has_current_marker(self, content: str) -> bool: + return bool(self._content_tokens(content).intersection(_CURRENT_MARKERS)) @staticmethod def _normalise_semantic(content: str) -> str: diff --git a/services/dma-api/tests/test_recall.py b/services/dma-api/tests/test_recall.py index f01e4d7..36fcd8a 100644 --- a/services/dma-api/tests/test_recall.py +++ b/services/dma-api/tests/test_recall.py @@ -91,3 +91,65 @@ def test_recall_prefers_current_state_memories_for_now_queries(tmp_path) -> None assert response.status_code == 200 contents = [item["content"] for item in response.json()["results"]] assert contents == ["User now prefers Java Spring Boot for backend APIs."] + + +def test_recall_prefix_only_matching_no_false_positives(tmp_path) -> None: + """Adversarial tests for prefix-only token matching. + + The old bidirectional substring matching caused false positives: + - 'api' matched 'rapid' (substring but not prefix) + - 'art' matched 'particle' (substring but not prefix) + + Prefix-only matching ensures: + - 'api' matches 'api' (exact) and 'apikey' (prefix), but NOT 'rapid' + - 'art' matches 'art' (exact) and 'artist' (prefix), but NOT 'particle' + - 'cat' matches 'cat', 'cater', 'category' (all valid prefixes) + """ + app = create_app(Settings(database_path=tmp_path / "dma.db", api_key="test-key", tenant_id="tenant-a")) + with TestClient(app) as client: + # Store memories with words that could cause false positives with substring matching + _remember(client, "Rapid deployment is our goal.", "semantic", "key-cat-0000000002") + _remember(client, "Particle physics is complex.", "semantic", "key-cat-0000000003") + # Also store exact/prefix matches that SHOULD be found + _remember(client, "API key configuration done.", "semantic", "key-cat-0000000005") + _remember(client, "Artist portfolio updated.", "semantic", "key-cat-0000000006") + # 'cat' prefix matches are legitimate - store some + _remember(client, "The cat sits on the mat.", "semantic", "key-cat-0000000004") + _remember(client, "The category system organizes items.", "semantic", "key-cat-0000000001") + + # Query 'api' should NOT match 'rapid' (substring), but SHOULD match 'api' and 'apikey' + response = client.post( + "/v1/memories/recall", + headers={"Authorization": "Bearer test-key"}, + json={"agent_id": "coding-agent", "query": "api", "limit": 10}, + ) + assert response.status_code == 200 + results = response.json()["results"] + contents = [item["content"] for item in results] + assert "API key configuration done." in contents + assert "Rapid deployment is our goal." not in contents + + # Query 'art' should NOT match 'particle' (substring), but SHOULD match 'art' and 'artist' + response = client.post( + "/v1/memories/recall", + headers={"Authorization": "Bearer test-key"}, + json={"agent_id": "coding-agent", "query": "art", "limit": 10}, + ) + assert response.status_code == 200 + results = response.json()["results"] + contents = [item["content"] for item in results] + assert "Artist portfolio updated." in contents + assert "Particle physics is complex." not in contents + + # Query 'cat' matches 'cat', 'cater', 'category' (all valid prefixes) + response = client.post( + "/v1/memories/recall", + headers={"Authorization": "Bearer test-key"}, + json={"agent_id": "coding-agent", "query": "cat", "limit": 10}, + ) + assert response.status_code == 200 + results = response.json()["results"] + contents = [item["content"] for item in results] + # Both should match since 'cat' is a prefix of 'category' + assert "The cat sits on the mat." in contents + assert "The category system organizes items." in contents