Merge pull request stanfordnlp#967 from fsndzomga/graphviz-of-dspy-modules

feature(dspy): generate a graphical representation of modules

Author: arnavsinghvi11

dspy/experimental/__init__.py

@@ -1,2 +1,4 @@
+from module_graph import *
+
 from .synthesizer import *
 from .synthetic_data import *

dspy/experimental/module_graph.py

@@ -0,0 +1,142 @@
+import dspy
+
+try:
+    import graphviz  # type: ignore
+    graphviz_available = True
+except ImportError:
+    graphviz_available = False
+
+
+class ModuleGraph:
+    def __init__(self, module_name, module):
+        if graphviz_available is False:
+            raise ImportError(
+                """Please install graphviz to use this feature.
+                Run 'pip install graphviz'""")
+
+        self.graph = graphviz.Digraph(format='png')
+        self.nodes = set()
+        self.module_name = module_name
+        self.module = module
+        self.inspect_settings(dspy.settings)
+        self.add_module(self.module_name, self.module)
+
+    def inspect_settings(self, settings):
+        """Check for the existence and configuration of LM and RM and add them to the graph."""
+        components = {'lm': settings.lm, 'rm': settings.rm}
+        for component_name, component in components.items():
+            if component:
+                details = {attr: getattr(component, attr) for attr in dir(component)
+                           if not attr.startswith('_') and not callable(getattr(component, attr))}
+                component_details = f"{component_name.upper()} Details: " + ', '.join(f"{k}: {v}" for k, v in details.items() if k!='history')
+                self.graph.node(component_name, label=component_details, shape='box')
+                self.nodes.add(component_name)
+
+    def add_module(self, module_name, module):
+        """Add a module to the graph"""
+
+        module_type = type(module)
+
+        if 'dspy.predict' in str(module_type):
+            module_name = self.generate_module_name(module_name, module_type)
+            self.process_submodule(module_name, module)
+        else:
+            self.process_submodules(module_name, module)
+
+    def generate_module_name(self, base_name, module_type):  # noqa: ANN201
+        """ Generate a module name based on the module type"""
+        type_map = {
+            'Predict': '__Predict',
+            'ReAct': '__ReAct',
+            'ChainOfThought': '__ChainOfThought',
+            'ProgramOfThought': '__ProgramOfThought',
+            'MultiChainComparison': '__MultiChainComparison',
+            'majority': '__majority',
+        }
+
+        for key, suffix in type_map.items():
+            if key in str(module_type):
+                return base_name + suffix
+        return base_name
+
+    def process_submodules(self, module_name, module):
+        """ Process submodules of a module and add them to the graph"""
+
+        for sub_module_name, sub_module in module.__dict__.items():
+            if isinstance(sub_module, dspy.Predict):
+                sub_module_name = self.generate_module_name(sub_module_name, type(sub_module))
+                self.process_submodule(f"{module_name}__{sub_module_name}", sub_module)
+            elif isinstance(sub_module, (dspy.Module, dspy.Retrieve)):
+                self.add_module(f"{module_name}__{sub_module_name}", sub_module)
+                if isinstance(sub_module, dspy.Retrieve):
+                    self.graph.edge("rm", 'lm', label='RM used in Module')
+
+    def process_submodule(self, sub_module_name, sub_module):
+        """Process a submodule and add it to the graph"""
+
+        for field_type, fields in [('input', sub_module.signature.input_fields),
+                                   ('output', sub_module.signature.output_fields)]:
+            for field_name, field in fields.items():
+                node_id = f"{sub_module_name}_{field_type}_{field_name}"
+                if node_id not in self.nodes:
+                    label = f"{field_name}: ({field.json_schema_extra['desc']})"
+                    self.graph.node(node_id, label=label, shape='ellipse')
+                    self.nodes.add(node_id)
+                edge_direction = (node_id, sub_module_name) if field_type == 'input' else (sub_module_name, node_id)
+                self.graph.edge(*edge_direction)
+
+        # Add node for the submodule itself
+        self.graph.node(sub_module_name, label=sub_module_name, shape='box')
+        self.nodes.add(sub_module_name)
+
+        # Connect submodule to LM if configured
+        if 'lm' in self.nodes:
+            self.graph.edge('lm', sub_module_name, label='LM used in Module')
+
+    def render_graph(self, filename=None):
+        """Render the graph to a file(png)"""
+        if filename is None:
+            filename = self.module_name
+        self.graph.render(filename)
+
+
+# Example usage of the ModuleGraph class:
+# import dspy
+# import os
+# from dotenv import load_dotenv
+# from dspy.experimental import ModuleGraph
+
+# load_dotenv()
+
+# # Configuration of dspy models
+# llm = dspy.OpenAI(
+#     model='gpt-3.5-turbo',
+#     api_key=os.environ['OPENAI_API_KEY'],
+#     max_tokens=100
+# )
+
+# colbertv2_wiki = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')
+
+# dspy.settings.configure(lm=llm, rm=colbertv2_wiki)
+
+# class GenerateAnswer(dspy.Signature):
+#   "Answer with long and detailed answers"
+#   context = dspy.InputField(desc="may content relevant facts")
+#   question = dspy.InputField()
+#   answer = dspy.OutputField(desc="often between 10 and 50 words")
+
+# class RAG(dspy.Module):
+#   def __init__(self, num_passages=3):
+#     super().__init__()
+#     self.retrieve = dspy.Retrieve(k=num_passages)
+#     self.generate_answer = dspy.ChainOfThought(GenerateAnswer)
+
+#   def forward(self, question):
+#     context = self.retrieve(question).passages
+#     prediction = self.generate_answer(context=context, question=question)
+#     return dspy.Prediction(context=context, answer=prediction.answer)
+
+# rag_system = RAG()
+# graph = ModuleGraph("RAG", rag_system)
+
+# graph.render_graph()