Repository: superduper-io/superduper

Issue:
[GRAPH-MOD] [FEATURE] Graph based model-API
The current way of doing things such as RAG and similar models is restrictive. We might want to implement filters based on terms detected in the query string/ query image(s).

**Example**

```python
query = (
    collection()
        .find({Variable('character'): {'$in': '_outputs.image.characters.0'}})
        .like({'text': Variable('question')})
)

ner = NER()
query_model = QueryModel(query, variables={'X': 'character', 'Z': 'question'})

class QueryModel:
    def predict(self, X, Z):
        ...

m = GraphModel()
m.add_edge(m.input, ner.X)
m.add_edge(m.input, query_model.X)
m.add_edge(ner, query_model.Y)
```

In this example:
1. the input is passed to the `ner` model (recognizes import names).
2. the output of the `ner` model is passed to the `X` parameter of the `query_model`
3. the input is passed to the `Z` parameter of the `query_model`

- [ ] Bring created `Node` API inside `_Predictor`
- [x] #1555 


Discussion:
We would need a way to set the outputs of the model (the node with the greatest path length?). @kartik4949 to provide input.

Resolving patch:
diff --git a/superduperdb/components/graph.py b/superduperdb/components/graph.py
new file mode 100644
--- /dev/null
+++ b/superduperdb/components/graph.py
@@ -0,0 +1,287 @@
+import dataclasses as dc
+import typing as t
+
+import networkx as nx
+
+from superduperdb import Schema
+from superduperdb.backends.query_dataset import QueryDataset
+from superduperdb.components.model import Signature, _Predictor
+
+
+@dc.dataclass(kw_only=True)
+class Graph(_Predictor):
+    '''
+    Represents a directed acyclic graph composed of interconnected model nodes.
+
+    This class enables the creation of complex predictive models
+    by defining a computational graph structure where each node
+    represents a predictive model. Models can be connected via edges
+    to define dependencies and flow of data.
+
+    The predict() method executes predictions through the graph, ensuring
+    correct data flow and handling of dependencies
+
+    Example:
+    >>  g = Graph(
+    >>    identifier='simple-graph', input=model1, outputs=[model2], signature='*args'
+    >>  )
+    >>  g.connect(model1, model2)
+    >>  assert g.predict_one(1) == [(4, 2)]
+
+    '''
+
+    models: t.List[_Predictor] = dc.field(default_factory=list)
+    edges: t.List[t.Tuple[str, str, t.Tuple[int, str]]] = dc.field(default_factory=list)
+    input: _Predictor
+    outputs: t.List[t.Union[str, _Predictor]] = dc.field(default_factory=list)
+    _DEFAULT_ARG_WEIGHT: t.ClassVar[tuple] = (None, Signature.singleton)
+    signature: str = Signature.args_kwargs  # type: ignore[misc]
+    type_id: t.ClassVar[str] = 'graph'
+
+    def __post_init__(self, artifacts):
+        self.G = nx.DiGraph()
+        self.nodes = {}
+        self._node_output_cache = {}
+        self.version = 0
+        self._db = None
+
+        assert all([isinstance(o, _Predictor) for o in self.outputs])
+
+        self.output_schema = Schema(
+            identifier=self.identifier,
+            fields={k.identifier: k.datatype for k in self.outputs},
+        )
+
+        self.outputs = [
+            o.identifier if isinstance(o, _Predictor) else o for o in self.outputs
+        ]
+        self.input.signature = Signature.args
+
+        # Load the models and edges into a di graph
+        models = {m.identifier: m for m in self.models}
+        if self.edges and models:
+            for connection in self.edges:
+                u, v, on = connection
+                self.connect(
+                    models[u],
+                    models[v],
+                    on=on,
+                    update_edge=False,
+                )
+        super().__post_init__(artifacts=artifacts)
+
+    def connect(
+        self,
+        u: _Predictor,
+        v: _Predictor,
+        on: t.Optional[t.Tuple[int, str]] = None,
+        update_edge: t.Optional[bool] = True,
+    ):
+        '''
+        Connects two nodes `u` and `v` on edge where edge is a tuple with
+        first element describing outputs index (int or None)
+        and second describing input argument (str).
+
+        Note:
+        output index: None means all outputs of node u are connected to node v
+        '''
+        assert isinstance(u, _Predictor)
+        assert isinstance(v, _Predictor)
+
+        if u.identifier not in self.nodes:
+            u.signature = Signature.args  # type: ignore[misc]
+            self.nodes[u.identifier] = u
+            self.G.add_node(u.identifier)
+
+        if v.identifier not in self.nodes:
+            v.signature = Signature.args  # type: ignore[misc]
+            self.nodes[v.identifier] = v
+            self.G.add_node(v.identifier)
+
+        G_ = self.G.copy()
+        G_.add_edge(u.identifier, v.identifier, weight=on or self._DEFAULT_ARG_WEIGHT)
+
+        if not nx.is_directed_acyclic_graph(G_):
+            raise TypeError('The graph is not DAG with this edge')
+        self.G = G_
+
+        if update_edge:
+            self.edges.append(
+                (u.identifier, v.identifier, on or self._DEFAULT_ARG_WEIGHT)
+            )
+            if isinstance(u, _Predictor) and u not in self.models:
+                self.models.append(u)
+            if v not in self.models:
+                self.models.append(v)
+        return
+
+    def fetch_output(self, output, index: t.Optional[int] = None, one: bool = False):
+        if index is not None:
+            assert isinstance(index, int)
+            try:
+                if one:
+                    return output[index]
+                else:
+                    return [[o[index]] for o in output]
+
+            except KeyError:
+                raise KeyError("Model node does not have sufficient outputs")
+        else:
+            if one:
+                return output
+            return [[o] for o in output]
+
+    def validate(self, node):
+        '''
+        Validates the graph for any disconnection
+        '''
+        # TODO: Create a cache to reduce redundant validation in predict in db
+
+        predecessors = list(self.G.predecessors(node))
+        dependencies = [self.validate(node=p) for p in predecessors]
+        model = self.nodes[node]
+        if dependencies and len(model.inputs) != len(dependencies):
+            raise TypeError(
+                f'Graph disconnected at Node: {model.identifier} '
+                f'and is partially connected with {dependencies}\n'
+                f'Required connected node is {len(model.inputs)} '
+                f'but got only {len(dependencies)}, '
+                f'Node required params: {model.inputs.params}'
+            )
+        return node
+
+    def _fetch_inputs(self, dataset, edges=[], outputs=[]):
+        arg_inputs = []
+
+        def _length(lst):
+            count = 0
+            for item in lst:
+                if isinstance(item, (list, tuple)):
+                    count += _length(item)
+                else:
+                    count += 1
+            return count
+
+        if not _length(outputs):
+            outputs = dataset
+
+        for ix, edge in enumerate(edges):
+            output_key, input_key = edge['weight']
+
+            arg_input_dataset = self.fetch_output(outputs[ix], output_key, one=False)
+            if input_key == Signature.singleton:
+                return arg_input_dataset
+
+            arg_inputs.append(arg_input_dataset)
+
+        if not arg_inputs:
+            return outputs
+        return self._transpose(arg_inputs, args=True)
+
+    def _fetch_one_inputs(self, args, kwargs, edges=[], outputs=[]):
+        node_input = {}
+        for ix, edge in enumerate(edges):
+            output_key, input_key = edge['weight']
+
+            node_input[input_key] = self.fetch_output(outputs[ix], output_key, one=True)
+            kwargs = node_input
+            args = ()
+
+        if Signature.singleton in node_input:
+            args = (node_input[Signature.singleton],)
+            kwargs = {}
+        return args, kwargs
+
+    def _predict_on_node(self, *args, node=None, cache={}, one=True, **kwargs):
+        if node not in cache:
+            predecessors = list(self.G.predecessors(node))
+            outputs = [
+                self._predict_on_node(*args, **kwargs, one=one, node=p, cache=cache)
+                for p in predecessors
+            ]
+
+            edges = [self.G.get_edge_data(p, node) for p in predecessors]
+
+            if one is True:
+                args, kwargs = self._fetch_one_inputs(
+                    args, kwargs, edges=edges, outputs=outputs
+                )
+                cache[node] = self.nodes[node].predict_one(*args, **kwargs)
+            else:
+                dataset = self._fetch_inputs(args[0], edges=edges, outputs=outputs)
+                cache[node] = self.nodes[node].predict(dataset=dataset)
+            return cache[node]
+        return cache[node]
+
+    def predict_one(self, *args, **kwargs):
+        '''
+        Single data point prediction passes the args and kwargs to defined node flow
+        in the graph.
+        '''
+        # Validate the node for incompletion
+        list(map(self.validate, self.outputs))
+        cache = {}
+
+        outputs = [
+            self._predict_on_node(*args, node=output, cache=cache, **kwargs)
+            for output in self.outputs
+        ]
+
+        # TODO: check if output schema and datatype required
+        return outputs
+
+    def patch_dataset_to_args(self, dataset):
+        '''
+        Patch the dataset with args type as default, since all
+        corresponding nodes takes args as input type
+        '''
+        args_dataset = []
+        signature = self.signature
+
+        def mapping(x):
+            nonlocal signature
+            if signature == Signature.kwargs:
+                return list(x.values())
+            elif signature == Signature.args_kwargs:
+                return list(x[0]) + list(x[1].values())
+            else:
+                return x
+
+        for data in dataset:
+            data = mapping(data)
+            args_dataset.append(data)
+        return args_dataset
+
+    def predict(self, dataset: t.Union[t.List, QueryDataset]) -> t.List:
+        # Validate the node for incompletion
+        list(map(self.validate, self.outputs))
+
+        if isinstance(dataset, QueryDataset):
+            raise TypeError('QueryDataset is not supported in graph mode')
+        cache: t.Dict[str, t.Any] = {}
+        dataset = self.patch_dataset_to_args(dataset)
+
+        outputs = [
+            self._predict_on_node(dataset, node=output, cache=cache, one=False)
+            for output in self.outputs
+        ]
+
+        # TODO: check if output schema and datatype required
+        return outputs
+
+    def encode_outputs(self, outputs):
+        encoded_outputs = []
+        for o, n in zip(outputs, self.outputs):
+            encoded_outputs.append(self.nodes[n].encode_outputs(o))
+        outputs = self._transpose(outputs=encoded_outputs or outputs)
+        return self.encode_with_schema(outputs)
+
+    @staticmethod
+    def _transpose(outputs, args=False):
+        transposed_outputs = []
+        for i in range(len(outputs[0])):
+            batch_outs = []
+            for o in outputs:
+                batch_outs.append(o[i][0] if args else o[i])
+            transposed_outputs.append(batch_outs)
+        return transposed_outputs
diff --git a/superduperdb/components/model.py b/superduperdb/components/model.py
--- a/superduperdb/components/model.py
+++ b/superduperdb/components/model.py
@@ -612,7 +612,6 @@ def handle_input_type(data, signature):
                 f'{Signature.args}, '
                 f'{Signature.singleton}.'
             )
-            raise Exception('Unexpected signature')
 
     def _predict_with_select_and_ids(
         self,
@@ -642,14 +641,13 @@ def _predict_with_select_and_ids(
             X=X, db=db, select=select, ids=ids, in_memory=in_memory
         )
         outputs = self.predict(dataset)
-        outputs = self._encode_outputs(outputs)
+        outputs = self.encode_outputs(outputs)
 
         logging.info(f'Adding {len(outputs)} model outputs to `db`')
 
         assert isinstance(
             self.version, int
         ), 'Version has not been set, can\'t save outputs...'
-
         select.model_update(
             db=db,
             model=self.identifier,
@@ -661,7 +659,7 @@ def _predict_with_select_and_ids(
             **self.model_update_kwargs,
         )
 
-    def _encode_outputs(self, outputs):
+    def encode_outputs(self, outputs):
         if isinstance(self.datatype, DataType):
             if self.flatten:
                 outputs = [
@@ -670,14 +668,19 @@ def _encode_outputs(self, outputs):
             else:
                 outputs = [self.datatype(x).encode() for x in outputs]
         elif isinstance(self.output_schema, Schema):
-            encoded_outputs = []
-            for output in outputs:
-                if isinstance(output, dict):
-                    encoded_outputs.append(self.output_schema(output))
-                elif self.flatten:
-                    encoded_output = [self.output_schema(x) for x in output]
-                    encoded_outputs.append(encoded_output)
-            outputs = encoded_outputs if encoded_outputs else outputs
+            outputs = self.encode_with_schema(outputs)
+
+        return outputs
+
+    def encode_with_schema(self, outputs):
+        encoded_outputs = []
+        for output in outputs:
+            if isinstance(output, dict):
+                encoded_outputs.append(self.output_schema(output))
+            elif self.flatten:
+                encoded_output = [self.output_schema(x) for x in output]
+                encoded_outputs.append(encoded_output)
+        outputs = encoded_outputs if encoded_outputs else outputs
         return outputs
 
 
diff --git a/superduperdb/ext/torch/model.py b/superduperdb/ext/torch/model.py
--- a/superduperdb/ext/torch/model.py
+++ b/superduperdb/ext/torch/model.py
@@ -54,7 +54,7 @@ def factory(
         forward_method: str = '__call__',
         train_forward_method: str = '__call__',
         loader_kwargs: t.Dict = dc.field(default_factory=lambda: {}),
-        signature: str = Signature.singleton,
+        preprocess_signature: str = Signature.singleton,
         forward_signature: str = Signature.singleton,
         postprocess_signature: str = Signature.singleton,
         **kwargs,
@@ -69,7 +69,7 @@ def factory(
             forward_method=forward_method,
             train_forward_method=train_forward_method,
             loader_kwargs=loader_kwargs,
-            signature=signature,
+            preprocess_signature=preprocess_signature,
             forward_signature=forward_signature,
             postprocess_signature=postprocess_signature,
         )
@@ -154,7 +154,7 @@ class TorchModel(_Predictor, _Fittable, _DeviceManaged):
     forward_method: str = '__call__'
     train_forward_method: str = '__call__'
     loader_kwargs: t.Dict = dc.field(default_factory=lambda: {})
-    signature: str = Signature.singleton  # type: ignore[misc]
+    preprocess_signature: str = Signature.singleton
     forward_signature: str = Signature.singleton
     postprocess_signature: str = Signature.singleton
 
@@ -166,6 +166,19 @@ def __post_init__(self, artifacts):
 
         self._validation_set_cache = {}
 
+    @property
+    def signature(self):
+        if self.preprocess:
+            return self.preprocess_signature
+        return self.forward_signature
+
+    @signature.setter
+    def signature(self, signature):
+        if self.preprocess:
+            self.preprocess_signature = signature
+        else:
+            self.forward_signature = signature
+
     @property
     def inputs(self) -> CallableInputs:
         return CallableInputs(
@@ -251,7 +264,7 @@ def predict_one(self, *args, **kwargs):
         with torch.no_grad(), eval(self.object):
             if self.preprocess is not None:
                 out = self.preprocess(*args, **kwargs)
-                args, kwargs = self.handle_input_type(out, self.forward_signature)
+                args, kwargs = self.handle_input_type(out, self.signature)
 
             args, kwargs = to_device((args, kwargs), self.device)
             args, kwargs = create_batch((args, kwargs))
@@ -277,7 +290,7 @@ def predict(self, dataset: t.Union[t.List, QueryDataset]) -> t.List:
             out = []
             for batch in tqdm(loader, total=len(loader)):
                 batch = to_device(batch, device_of(self.object))
-                args, kwargs = self.handle_input_type(batch, self.forward_signature)
+                args, kwargs = self.handle_input_type(batch, self.signature)
                 method = getattr(self.object, self.forward_method)
                 tmp = method(*args, **kwargs, **self.predict_kwargs)
                 tmp = to_device(tmp, 'cpu')


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -221,6 +221,7 @@ def add_models(db: Datalayer):
             object=torch.nn.Linear(*weight_shape),
             identifier=identifier,
             datatype=datatype,
+            preferred_devices=['cpu'],
         )
         db.add(m)
 
diff --git a/test/unittest/component/test_graph.py b/test/unittest/component/test_graph.py
new file mode 100644
--- /dev/null
+++ b/test/unittest/component/test_graph.py
@@ -0,0 +1,127 @@
+from test.db_config import DBConfig
+
+import networkx as nx
+import pytest
+
+from superduperdb.components.graph import Graph
+from superduperdb.components.model import Model, Signature
+
+
+@pytest.fixture
+@pytest.mark.parametrize("db", [DBConfig.mongodb_empty], indirect=True)
+def model1(db):
+    def model_object(x):
+        return x + 1
+
+    model = Model(identifier='m1', object=model_object)
+    db.add(model)
+    yield model
+
+
+@pytest.fixture
+@pytest.mark.parametrize("db", [DBConfig.mongodb_empty], indirect=True)
+def model2(db):
+    def model_object(x):
+        return x + 2, x
+
+    model = Model(identifier='m2', object=model_object)
+    db.add(model)
+    yield model
+
+
+@pytest.fixture
+@pytest.mark.parametrize("db", [DBConfig.mongodb_empty], indirect=True)
+def model2_multi(db):
+    def model_object(x, y):
+        return x + y + 2
+
+    model = Model(identifier='m2_multi', object=model_object)
+    db.add(model)
+    yield model
+
+
+@pytest.fixture
+@pytest.mark.parametrize("db", [DBConfig.mongodb_empty], indirect=True)
+def model3(db):
+    def model_object(x, y):
+        return x + y + 3
+
+    model = Model(identifier='m3', object=model_object)
+    db.add(model)
+    yield model
+
+
+def test_simple_graph(model1, model2):
+    g = Graph(
+        identifier='simple-graph', input=model1, outputs=[model2], signature='*args'
+    )
+    g.connect(model1, model2)
+    assert g.predict_one(1) == [(4, 2)]
+
+    g = Graph(
+        identifier='simple-graph', input=model1, outputs=[model2], signature='*args'
+    )
+
+    g.connect(model1, model2)
+    assert g.predict([[1], [2], [3]]) == [[(4, 2), (5, 3), (6, 4)]]
+
+
+def test_complex_graph(model1, model2_multi, model3, model2):
+    g = Graph(
+        identifier='complex-graph',
+        input=model1,
+        outputs=[model2, model2_multi],
+        signature=Signature.kwargs,
+    )
+    g.connect(model1, model2_multi, on=(None, 'x'))
+    g.connect(model1, model2)
+    g.connect(model2, model2_multi, on=(0, 'y'))
+    g.connect(model2, model3, on=(1, 'x'))
+    g.connect(model2_multi, model3, on=(None, 'y'))
+    assert g.predict_one(1) == [(4, 2), 8]
+    assert g.predict([{'x': 1}, {'x': 2}, {'x': 3}]) == [
+        [(4, 2), (5, 3), (6, 4)],
+        [8, 10, 12],
+    ]
+    g.signature = Signature.args
+    assert g.predict([[1], [2], [3]]) == [[(4, 2), (5, 3), (6, 4)], [8, 10, 12]]
+
+
+def test_non_dag(model1, model2):
+    with pytest.raises(TypeError) as excinfo:
+        g = Graph(identifier='complex-graph', input=model1)
+        g.connect(model1, model2)
+        g.connect(model2, model1)
+        assert 'The graph is not DAG' in str(excinfo.value)
+
+
+def test_disconnected_edge(model1, model2_multi):
+    with pytest.raises(TypeError) as excinfo:
+        g = Graph(identifier='complex-graph', input=model1, outputs=[model2_multi])
+        g.connect(model1, model2_multi, on=(-1, 'x'))
+        g.predict_one(1)
+        assert 'Graph disconnected at Node: m2_multi' in str(excinfo.value)
+
+
+def test_complex_graph_with_select(db):
+    linear_a = db.load('model', 'linear_a')
+    linear_b = db.load('model', 'linear_b')
+    g = Graph(identifier='complex-graph', input=linear_a, outputs=[linear_b])
+    g.connect(linear_a, linear_b)
+
+    from superduperdb.backends.mongodb import Collection
+
+    select = Collection('documents').find({})
+    g.predict_in_db(X='x', select=select, db=db)
+    assert all(
+        ['complex-graph' in x['_outputs']['x'] for x in list(db.execute(select))]
+    )
+
+
+@pytest.mark.parametrize("db", [DBConfig.mongodb_empty], indirect=True)
+def test_serialization(db, model1):
+    g = Graph(identifier='complex-graph', input=model1)
+    original_g = g.G
+    db.add(g)
+    g = db.load('graph', 'complex-graph')
+    assert nx.utils.graphs_equal(original_g, g.G)
diff --git a/test/unittest/test_quality.py b/test/unittest/test_quality.py
--- a/test/unittest/test_quality.py
+++ b/test/unittest/test_quality.py
@@ -19,7 +19,7 @@
 ALLOWABLE_DEFECTS = {
     'cast': 12,  # Try to keep this down
     'noqa': 7,  # This should never change
-    'type_ignore': 40,  # This should only ever increase in obscure edge cases
+    'type_ignore': 42,  # This should only ever increase in obscure edge cases
 }