Towards better inference: bits → nibbles

octopoes/.ci/docker-compose.yml

@@ -17,7 +17,7 @@ services:
       args:
         ENVIRONMENT: dev
       context: .
-    command: pytest tests/integration --timeout=300
+    command: pytest -s tests/integration/test_nibbles.py --timeout=300
     depends_on:
       - xtdb
       - ci_octopoes

octopoes/nibbles/definitions.py

@@ -0,0 +1,96 @@
+import importlib
+import pkgutil
+from collections.abc import Iterable
+from pathlib import Path
+from types import MethodType, ModuleType
+
+import structlog
+from pydantic import BaseModel
+
+from octopoes.models import OOI
+
+NIBBLES_DIR = Path(__file__).parent
+NIBBLE_ATTR_NAME = "NIBBLE"
+NIBBLE_FUNC_NAME = "nibble"
+logger = structlog.get_logger(__name__)
+
+
+class NibbleParameter(BaseModel):
+    object_type: type
+    parser: str = "[]"
+
+    def __eq__(self, other):
+        if isinstance(other, NibbleParameter):
+            return vars(self) == vars(other)
+        elif isinstance(other, type):
+            return self.object_type == other
+        else:
+            return False
+
+
+class NibbleDefinition:
+    id: str
+    signature: list[NibbleParameter]
+    query: str | None = None
+    min_scan_level: int = 1
+    default_enabled: bool = True
+    config_ooi_relation_path: str | None = None
+    payload: MethodType | None = None
+
+    def __init__(
+        self,
+        name: str,
+        signature: list[NibbleParameter],
+        query: str | None = None,
+        min_scan_level: int = 1,
+        default_enabled: bool = True,
+        config_ooi_relation_path: str | None = None,
+    ):
+        self.id = name
+        self.signature = signature
+        self.query = query
+        self.min_scan_level = min_scan_level
+        self.default_enabled = default_enabled
+        self.config_ooi_relation_path = config_ooi_relation_path
+
+    def __call__(self, args: Iterable[OOI]) -> OOI | Iterable[OOI | None] | None:
+        if self.payload is None:
+            raise NotImplementedError
+        else:
+            return self.payload(*args)
+
+
+def get_nibble_definitions() -> dict[str, NibbleDefinition]:
+    nibble_definitions = {}
+
+    for package in pkgutil.walk_packages([str(NIBBLES_DIR)]):
+        if package.name in ["definitions", "runner"]:
+            continue
+
+        try:
+            module: ModuleType = importlib.import_module(".nibble", f"{NIBBLES_DIR.name}.{package.name}")
+
+            if hasattr(module, NIBBLE_ATTR_NAME):
+                nibble_definition: NibbleDefinition = getattr(module, NIBBLE_ATTR_NAME)
+
+                try:
+                    payload: ModuleType = importlib.import_module(
+                        f".{package.name}", f"{NIBBLES_DIR.name}.{package.name}"
+                    )
+                    if hasattr(payload, NIBBLE_FUNC_NAME):
+                        nibble_definition.payload = getattr(payload, NIBBLE_FUNC_NAME)
+                    else:
+                        logger.warning('module "%s" has no function %s', package.name, NIBBLE_FUNC_NAME)
+
+                except ModuleNotFoundError:
+                    logger.warning('package "%s" has no function nibble', package.name)
+
+                nibble_definitions[nibble_definition.id] = nibble_definition
+
+            else:
+                logger.warning('module "%s" has no attribute %s', package.name, NIBBLE_ATTR_NAME)
+
+        except ModuleNotFoundError:
+            logger.warning('package "%s" has no module nibble', package.name)
+
+    return nibble_definitions

octopoes/nibbles/runner.py

@@ -0,0 +1,138 @@
+import json
+from collections.abc import Callable, Iterable
+from datetime import datetime
+from typing import TypeVar
+
+from xxhash import xxh3_128_hexdigest as xxh3  # INFO: xxh3_64_hexdigest is faster but hash more collision probabilities
+
+from nibbles.definitions import NibbleDefinition, get_nibble_definitions
+from octopoes.models import OOI
+from octopoes.models.origin import Origin, OriginType
+from octopoes.models.types import type_by_name
+from octopoes.repositories.ooi_repository import OOIRepository
+from octopoes.repositories.origin_repository import OriginRepository
+from octopoes.repositories.scan_profile_repository import ScanProfileRepository
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def ooi_type(ooi: OOI) -> type[OOI]:
+    return type_by_name(ooi.get_ooi_type())
+
+
+def merge_with(func: Callable[[set[T], set[T]], set[T]], d1: dict[U, set[T]], d2: dict[U, set[T]]) -> dict[U, set[T]]:
+    return {k: func(d1.get(k, set()), d2.get(k, set())) for k in set(d1) | set(d2)}
+
+
+def flatten(items: Iterable[OOI | Iterable[OOI | None] | None]) -> Iterable[OOI]:
+    for item in items:
+        if isinstance(item, OOI):
+            yield item
+        elif item is None:
+            continue
+        else:
+            yield from flatten(item)
+
+
+def nibble_hasher(data: Iterable) -> str:
+    return xxh3(
+        "".join(
+            [
+                json.dumps(json.loads(ooi.model_dump_json()), sort_keys=True)
+                if isinstance(ooi, OOI)
+                else json.dumps(ooi, sort_keys=True)
+                for ooi in data
+            ]
+        )
+    )
+
+
+class NibblesRunner:
+    def __init__(
+        self,
+        ooi_repository: OOIRepository,
+        origin_repository: OriginRepository,
+        scan_profile_repository: ScanProfileRepository,
+        perform_writes: bool = True,
+    ):
+        self.ooi_repository = ooi_repository
+        self.origin_repository = origin_repository
+        self.scan_profile_repository = scan_profile_repository
+        self.perform_writes = perform_writes
+        self.update_nibbles()
+
+    def update_nibbles(self):
+        self.nibbles: dict[str, NibbleDefinition] = get_nibble_definitions()
+
+    def _run(self, ooi: OOI, valid_time: datetime) -> dict[str, dict[tuple, set[OOI]]]:
+        return_value: dict[str, dict[tuple, set[OOI]]] = {}
+        nibblets = self.origin_repository.list_origins(
+            valid_time, origin_type=OriginType.NIBBLET, parameters_references=[ooi.reference]
+        )
+        if nibblets:
+            for nibblet in nibblets:
+                # INFO: we do not strictly need this if statement because OriginType.NIBBLETS \
+                # always have parameters_references but it makes the linters super happy
+                if nibblet.parameters_references:
+                    nibble = self.nibbles[nibblet.method]
+                    args = self.ooi_repository.nibble_query(
+                        ooi,
+                        nibble,
+                        valid_time,
+                        nibblet.parameters_references
+                        if nibble.query is not None and nibble.query.count("$") > 0
+                        else [],
+                    )
+                    results = {
+                        tuple(arg): set(flatten([nibble(arg)]))
+                        for arg in args
+                        if nibblet.parameters_hash != nibble_hasher(arg)
+                    }
+                    return_value |= {nibble.id: results}
+        nibblet_nibbles = {self.nibbles[nibblet.method] for nibblet in nibblets}
+        for nibble in filter(lambda x: type(ooi) in x.signature and x not in nibblet_nibbles, self.nibbles.values()):
+            args = self.ooi_repository.nibble_query(ooi, nibble, valid_time)
+            results = {tuple(arg): set(flatten([nibble(arg)])) for arg in args}
+            return_value |= {nibble.id: results}
+        # TODO: we could cache the writes for single OOI nibbles
+        self._write({ooi: return_value}, valid_time)
+        return return_value
+
+    def _cleared(self, ooi: OOI, valid_time: datetime) -> bool:
+        ooi_level = self.scan_profile_repository.get(ooi.reference, valid_time).level.value
+        target_nibbles = filter(lambda x: type(ooi) in x.signature, self.nibbles.values())
+        return any(nibble.min_scan_level < ooi_level for nibble in target_nibbles)
+
+    def _write(self, inferences: dict[OOI, dict[str, dict[tuple, set[OOI]]]], valid_time: datetime):
+        if self.perform_writes:
+            for source_ooi, results in inferences.items():
+                self.ooi_repository.save(source_ooi, valid_time)
+                for nibble_id, run_result in results.items():
+                    for arg, result in run_result.items():
+                        nibble_origin = Origin(
+                            method=nibble_id,
+                            origin_type=OriginType.NIBBLET,
+                            source=source_ooi.reference,
+                            result=[ooi.reference for ooi in result],
+                            parameters_hash=nibble_hasher(arg),
+                            # TODO: What to do if a is not an OOI?
+                            parameters_references=[a.reference for a in arg if isinstance(a, OOI)],
+                        )
+                        for ooi in result:
+                            self.ooi_repository.save(ooi, valid_time=valid_time)
+                        self.origin_repository.save(nibble_origin, valid_time=valid_time)
+
+    def infer(self, stack: list[OOI], valid_time: datetime) -> dict[OOI, dict[str, dict[tuple, set[OOI]]]]:
+        inferences: dict[OOI, dict[str, dict[tuple, set[OOI]]]] = {}
+        blockset = set(stack)
+        if stack and self._cleared(stack[-1], valid_time):
+            while stack:
+                ooi = stack.pop()
+                results = self._run(ooi, valid_time)
+                if results:
+                    blocks = set.union(set(), *[ooiset for result in results.values() for _, ooiset in result.items()])
+                    stack += [o for o in blocks if o not in blockset]
+                    blockset |= blocks
+                    inferences |= {ooi: results}
+        return inferences

octopoes/octopoes/core/service.py

@@ -8,6 +8,7 @@
 import structlog
 from bits.definitions import get_bit_definitions
 from bits.runner import BitRunner
+from nibbles.runner import NibblesRunner
 from pydantic import TypeAdapter
 
 from octopoes.config.settings import (
@@ -76,6 +77,7 @@ def __init__(
         self.origin_repository = origin_repository
         self.origin_parameter_repository = origin_parameter_repository
         self.scan_profile_repository = scan_profile_repository
+        self.nibbler = NibblesRunner(ooi_repository, origin_repository, scan_profile_repository)
         self.session = session
 
     @overload
@@ -170,10 +172,10 @@ def save_origin(
             self.ooi_repository.get(origin.source, valid_time)
         except ObjectNotFoundException:
             if (
-                origin.origin_type not in [OriginType.DECLARATION, OriginType.AFFIRMATION]
+                origin.origin_type not in [OriginType.DECLARATION, OriginType.AFFIRMATION, OriginType.NIBBLET]
                 and origin.source not in origin.result
             ):
-                raise ValueError("Origin source of observation does not exist")
+                raise ValueError(f"Origin source [{origin.source}] does not exist")
             elif origin.origin_type == OriginType.AFFIRMATION:
                 logger.debug("Affirmation source %s already deleted", origin.source)
                 return
@@ -200,6 +202,7 @@ def save_origin(
             self.origin_repository.delete(origin, valid_time=valid_time)
 
     def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
+        # The bit part of inferring
         bit_definition = get_bit_definitions().get(origin.method, None)
 
         if bit_definition is None:
@@ -234,6 +237,7 @@ def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
             if len(configs) != 0:
                 config = configs[-1].config
 
+        resulting_oois: list[OOI] = []
         try:
             if isinstance(self.session, XTDBSession):
                 start = perf_counter()
@@ -252,10 +256,14 @@ def _run_inference(self, origin: Origin, valid_time: datetime) -> None:
                 self.session.client.submit_transaction(ops)
             else:
                 resulting_oois = BitRunner(bit_definition).run(source, parameters, config=config)
-            self.save_origin(origin, resulting_oois, valid_time)
         except Exception as e:
             logger.exception("Error running inference", exc_info=e)
 
+        self.save_origin(origin, resulting_oois, valid_time)
+
+        # The nibble part of inferring
+        self.nibbler.infer([source], valid_time)
+
     @staticmethod
     def check_path_level(path_level: int | None, current_level: int):
         return path_level is not None and path_level >= current_level

octopoes/octopoes/models/origin.py

@@ -11,6 +11,7 @@ class OriginType(Enum):
     OBSERVATION = "observation"
     INFERENCE = "inference"
     AFFIRMATION = "affirmation"
+    NIBBLET = "nibblet"
 
 
 class Origin(BaseModel):
@@ -19,6 +20,8 @@ class Origin(BaseModel):
     source: Reference
     source_method: str | None = None  # None for bits and normalizers
     result: list[Reference] = Field(default_factory=list)
+    parameters_hash: str | None = None  # None for anything other than Nibblette
+    parameters_references: list[Reference] | None = None  # None for anything other than Nibblette
     task_id: UUID | None = None
 
     def __sub__(self, other) -> set[Reference]:
@@ -29,7 +32,22 @@ def __sub__(self, other) -> set[Reference]:
 
     @property
     def id(self) -> str:
-        if self.source_method is not None:
+        if self.origin_type == OriginType.NIBBLET:
+            return "|".join(
+                map(
+                    str,
+                    [
+                        self.__class__.__name__,
+                        self.origin_type.value,
+                        self.method,
+                        self.source,
+                        self.result,
+                        self.parameters_hash,
+                        self.parameters_references,
+                    ],
+                )
+            )
+        elif self.source_method is not None:
             return (
                 f"{self.__class__.__name__}|{self.origin_type.value}|{self.method}|{self.source_method}|{self.source}"
             )