mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
feat: Schema system for Basic Memory (#549)
Signed-off-by: phernandez <paul@basicmachines.co> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Paul Hernandez
@@ -0,0 +1,323 @@
|
||||
"""Schema inference engine for Basic Memory.
|
||||
|
||||
Analyzes notes of a given type and suggests a schema based on observation
|
||||
and relation frequency. Instead of requiring users to define schemas upfront,
|
||||
schemas emerge from actual usage patterns:
|
||||
|
||||
Write notes freely -> Patterns emerge -> Crystallize into schema
|
||||
|
||||
Frequency thresholds:
|
||||
- 95%+ present -> required field
|
||||
- 25%+ present -> optional field
|
||||
- Below 25% -> excluded from suggestion (but noted)
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
# --- Result Data Model ---
|
||||
|
||||
|
||||
@dataclass
|
||||
class FieldFrequency:
|
||||
"""Frequency analysis for a single field across notes of a type."""
|
||||
|
||||
name: str
|
||||
source: str # "observation" | "relation"
|
||||
count: int # notes containing this field
|
||||
total: int # total notes analyzed
|
||||
percentage: float
|
||||
sample_values: list[str] = field(default_factory=list)
|
||||
is_array: bool = False # True if typically appears multiple times per note
|
||||
target_type: str | None = None # For relations, the most common target entity type
|
||||
|
||||
|
||||
@dataclass
|
||||
class InferenceResult:
|
||||
"""Complete inference result with frequency analysis and suggested schema."""
|
||||
|
||||
entity_type: str
|
||||
notes_analyzed: int
|
||||
field_frequencies: list[FieldFrequency]
|
||||
suggested_schema: dict # Ready-to-use Picoschema YAML dict
|
||||
suggested_required: list[str]
|
||||
suggested_optional: list[str]
|
||||
excluded: list[str] # Below threshold
|
||||
|
||||
|
||||
# --- Note Data Abstraction ---
|
||||
# Instead of depending on the ORM Entity model, we accept simple data structures.
|
||||
# This keeps the inference engine decoupled from the data access layer.
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObservationData:
|
||||
"""Lightweight observation for schema analysis. Decoupled from ORM."""
|
||||
|
||||
category: str
|
||||
content: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class RelationData:
|
||||
"""Lightweight relation for schema analysis. Decoupled from ORM."""
|
||||
|
||||
relation_type: str
|
||||
target_name: str
|
||||
target_entity_type: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class NoteData:
|
||||
"""Minimal note representation for inference analysis.
|
||||
|
||||
Decoupled from ORM models so the inference engine can work with
|
||||
any data source (database, files, API responses).
|
||||
"""
|
||||
|
||||
identifier: str
|
||||
observations: list[ObservationData]
|
||||
relations: list[RelationData]
|
||||
|
||||
|
||||
# --- Inference Logic ---
|
||||
|
||||
|
||||
def infer_schema(
|
||||
entity_type: str,
|
||||
notes: list[NoteData],
|
||||
required_threshold: float = 0.95,
|
||||
optional_threshold: float = 0.25,
|
||||
max_sample_values: int = 5,
|
||||
) -> InferenceResult:
|
||||
"""Analyze notes and suggest a Picoschema definition.
|
||||
|
||||
Examines observation categories and relation types across all provided notes.
|
||||
Fields that appear in a high percentage of notes become required; those that
|
||||
appear less frequently become optional.
|
||||
|
||||
Args:
|
||||
entity_type: The entity type being analyzed (e.g., "Person").
|
||||
notes: List of NoteData objects to analyze.
|
||||
required_threshold: Frequency at or above which a field is required (default 0.95).
|
||||
optional_threshold: Frequency at or above which a field is optional (default 0.25).
|
||||
max_sample_values: Maximum number of sample values to include per field.
|
||||
|
||||
Returns:
|
||||
An InferenceResult with frequency analysis and suggested Picoschema dict.
|
||||
"""
|
||||
total = len(notes)
|
||||
if total == 0:
|
||||
return InferenceResult(
|
||||
entity_type=entity_type,
|
||||
notes_analyzed=0,
|
||||
field_frequencies=[],
|
||||
suggested_schema={},
|
||||
suggested_required=[],
|
||||
suggested_optional=[],
|
||||
excluded=[],
|
||||
)
|
||||
|
||||
# --- Analyze observation frequencies ---
|
||||
obs_frequencies = analyze_observations(notes, total, max_sample_values)
|
||||
|
||||
# --- Analyze relation frequencies ---
|
||||
rel_frequencies = analyze_relations(notes, total, max_sample_values)
|
||||
|
||||
# --- Classify fields by threshold ---
|
||||
all_frequencies = obs_frequencies + rel_frequencies
|
||||
suggested_required: list[str] = []
|
||||
suggested_optional: list[str] = []
|
||||
excluded: list[str] = []
|
||||
|
||||
for freq in all_frequencies:
|
||||
if freq.percentage >= required_threshold:
|
||||
suggested_required.append(freq.name)
|
||||
elif freq.percentage >= optional_threshold:
|
||||
suggested_optional.append(freq.name)
|
||||
else:
|
||||
excluded.append(freq.name)
|
||||
|
||||
# --- Build suggested Picoschema dict ---
|
||||
suggested_schema = _build_picoschema_dict(
|
||||
all_frequencies, required_threshold, optional_threshold
|
||||
)
|
||||
|
||||
return InferenceResult(
|
||||
entity_type=entity_type,
|
||||
notes_analyzed=total,
|
||||
field_frequencies=all_frequencies,
|
||||
suggested_schema=suggested_schema,
|
||||
suggested_required=suggested_required,
|
||||
suggested_optional=suggested_optional,
|
||||
excluded=excluded,
|
||||
)
|
||||
|
||||
|
||||
# --- Observation Analysis ---
|
||||
|
||||
|
||||
def analyze_observations(
|
||||
notes: list[NoteData],
|
||||
total: int,
|
||||
max_sample_values: int,
|
||||
) -> list[FieldFrequency]:
|
||||
"""Count observation category frequencies across notes.
|
||||
|
||||
A category is counted once per note (presence), not per occurrence.
|
||||
Array detection: if a category appears multiple times in a single note
|
||||
in more than half the notes where it appears, it's flagged as an array.
|
||||
"""
|
||||
# Count how many notes contain each category (presence per note)
|
||||
category_note_count: Counter[str] = Counter()
|
||||
# Count how many notes have multiple occurrences (for array detection)
|
||||
category_multi_count: Counter[str] = Counter()
|
||||
# Collect sample values
|
||||
category_samples: dict[str, list[str]] = {}
|
||||
|
||||
for note in notes:
|
||||
# Group observations by category within this note
|
||||
note_categories: dict[str, list[str]] = {}
|
||||
for obs in note.observations:
|
||||
note_categories.setdefault(obs.category, []).append(obs.content)
|
||||
|
||||
for category, values in note_categories.items():
|
||||
category_note_count[category] += 1
|
||||
if len(values) > 1:
|
||||
category_multi_count[category] += 1
|
||||
|
||||
# Collect sample values (deduplicated)
|
||||
samples = category_samples.setdefault(category, [])
|
||||
for v in values:
|
||||
if v not in samples and len(samples) < max_sample_values:
|
||||
samples.append(v)
|
||||
|
||||
# Build FieldFrequency objects
|
||||
frequencies: list[FieldFrequency] = []
|
||||
for category, count in category_note_count.most_common():
|
||||
# Array detection: if more than half of notes with this category have
|
||||
# multiple occurrences, treat it as an array field
|
||||
multi_count = category_multi_count.get(category, 0)
|
||||
is_array = multi_count > (count / 2)
|
||||
|
||||
frequencies.append(
|
||||
FieldFrequency(
|
||||
name=category,
|
||||
source="observation",
|
||||
count=count,
|
||||
total=total,
|
||||
percentage=count / total,
|
||||
sample_values=category_samples.get(category, []),
|
||||
is_array=is_array,
|
||||
)
|
||||
)
|
||||
|
||||
return frequencies
|
||||
|
||||
|
||||
# --- Relation Analysis ---
|
||||
|
||||
|
||||
def analyze_relations(
|
||||
notes: list[NoteData],
|
||||
total: int,
|
||||
max_sample_values: int,
|
||||
) -> list[FieldFrequency]:
|
||||
"""Count relation type frequencies across notes.
|
||||
|
||||
Similar to observations, a relation type is counted once per note.
|
||||
Array detection follows the same logic.
|
||||
"""
|
||||
rel_note_count: Counter[str] = Counter()
|
||||
rel_multi_count: Counter[str] = Counter()
|
||||
rel_samples: dict[str, list[str]] = {}
|
||||
# Track target entity types to suggest the type in the schema
|
||||
rel_target_types: dict[str, Counter[str]] = {}
|
||||
|
||||
for note in notes:
|
||||
note_rels: dict[str, list[str]] = {}
|
||||
note_rel_objects: dict[str, list[RelationData]] = {}
|
||||
for rel in note.relations:
|
||||
note_rels.setdefault(rel.relation_type, []).append(rel.target_name)
|
||||
note_rel_objects.setdefault(rel.relation_type, []).append(rel)
|
||||
|
||||
for rel_type, targets in note_rels.items():
|
||||
rel_note_count[rel_type] += 1
|
||||
if len(targets) > 1:
|
||||
rel_multi_count[rel_type] += 1
|
||||
|
||||
samples = rel_samples.setdefault(rel_type, [])
|
||||
for t in targets:
|
||||
if t not in samples and len(samples) < max_sample_values:
|
||||
samples.append(t)
|
||||
|
||||
# Track target entity types from individual relations (not the source note)
|
||||
target_counter = rel_target_types.setdefault(rel_type, Counter())
|
||||
for rel in note_rel_objects[rel_type]:
|
||||
if rel.target_entity_type:
|
||||
target_counter[rel.target_entity_type] += 1
|
||||
|
||||
frequencies: list[FieldFrequency] = []
|
||||
for rel_type, count in rel_note_count.most_common():
|
||||
multi_count = rel_multi_count.get(rel_type, 0)
|
||||
is_array = multi_count > (count / 2)
|
||||
|
||||
# Determine most common target type
|
||||
target_counter = rel_target_types.get(rel_type, Counter())
|
||||
most_common_target = target_counter.most_common(1)[0][0] if target_counter else None
|
||||
|
||||
frequencies.append(
|
||||
FieldFrequency(
|
||||
name=rel_type,
|
||||
source="relation",
|
||||
count=count,
|
||||
total=total,
|
||||
percentage=count / total,
|
||||
sample_values=rel_samples.get(rel_type, []),
|
||||
is_array=is_array,
|
||||
target_type=most_common_target,
|
||||
)
|
||||
)
|
||||
|
||||
return frequencies
|
||||
|
||||
|
||||
# --- Schema Generation ---
|
||||
|
||||
|
||||
def _build_picoschema_dict(
|
||||
frequencies: list[FieldFrequency],
|
||||
required_threshold: float,
|
||||
optional_threshold: float,
|
||||
) -> dict:
|
||||
"""Build a Picoschema YAML dict from field frequencies.
|
||||
|
||||
Only includes fields at or above the optional threshold.
|
||||
"""
|
||||
schema: dict = {}
|
||||
|
||||
for freq in frequencies:
|
||||
if freq.percentage < optional_threshold:
|
||||
continue
|
||||
|
||||
is_required = freq.percentage >= required_threshold
|
||||
|
||||
# --- Build the field key ---
|
||||
key = freq.name
|
||||
if not is_required:
|
||||
key += "?"
|
||||
if freq.is_array:
|
||||
key += "(array)"
|
||||
|
||||
# --- Build the field value ---
|
||||
if freq.source == "relation":
|
||||
# Relations become entity reference fields
|
||||
target = freq.target_type or "string"
|
||||
# Capitalize first letter for entity ref convention
|
||||
target = target[0].upper() + target[1:] if target != "string" else "string"
|
||||
schema[key] = target
|
||||
else:
|
||||
schema[key] = "string"
|
||||
|
||||
return schema
|
||||
Reference in New Issue
Block a user