It's working

docs/docs_skeleton/link_generator.py

@@ -11,7 +11,9 @@ logger = logging.getLogger(__name__)
 # Base URL for all class documentation
 _BASE_URL = "https://api.python.langchain.com/en/latest/"
 
-# Regular expression to match Python import lines
+# Regular expression to match Python code blocks
+code_block_re = re.compile(r"^(```python\n)(.*?)(```\n)", re.DOTALL | re.MULTILINE)
+# Regular expression to match langchain import lines
 _IMPORT_RE = re.compile(r"(from\s+(langchain\.\w+(\.\w+)*?)\s+import\s+)(\w+)")
 
 _CODEBLOCK_PATH = "src/theme/CodeBlock/"
@@ -58,35 +60,60 @@ def main():
 
 
 def replace_imports(file):
-    """Replace imports in a markdown file with links to their documentation and return the import info"""
-    imports = []
-
+    """Replace imports in each Python code block with links to their documentation and append the import info in a comment"""
+    all_imports = []
     with open(file, "r") as f:
         data = f.read()
 
-    for match in _IMPORT_RE.finditer(data):
-        class_name = match.group(4)
-        try:
-            module_path = get_full_module_name(match.group(2), class_name)
-        except AttributeError as e:
-            logger.error(f"Could not find module for {class_name}", e)
-            continue
-        module_path_parts = module_path.replace(".", "/")
+    def replacer(match):
+        # Extract the code block content
+        code = match.group(2)
+        # Replace if any import comment exists
+        existing_comment_re = re.compile(r"^<!--IMPORTS:.*?-->\n", re.MULTILINE)
+        code = existing_comment_re.sub("", code)
 
-        url = (
-            _BASE_URL
-            + module_path_parts
-            + "/"
-            + module_path_parts
-            + "."
-            + class_name
-            + ".html"
-        )
+        # Process imports in the code block
+        imports = []
+        for import_match in _IMPORT_RE.finditer(code):
+            class_name = import_match.group(4)
+            try:
+                module_path = get_full_module_name(import_match.group(2), class_name)
+            except AttributeError as e:
+                logger.error(f"Could not find module for {class_name}", e)
+                continue
 
-        # Add the import information to our list
-        imports.append({"imported": class_name, "source": match.group(2), "docs": url})
+            url = (
+                _BASE_URL
+                + "/"
+                + module_path.split(".")[1]
+                + "/"
+                + module_path
+                + "."
+                + class_name
+                + ".html"
+            )
 
-    return imports
+            # Add the import information to our list
+            imports.append(
+                {"imported": class_name, "source": import_match.group(2), "docs": url}
+            )
+
+        if imports:
+            all_imports.extend(imports)
+            # Create a unique comment containing the import information
+            import_comment = f"<!--IMPORTS:{json.dumps(imports)}-->"
+            # Inject the import comment at the start of the code block
+            return match.group(1) + import_comment + "\n" + code + match.group(3)
+        else:
+            # If there are no imports, return the original match
+            return match.group(0)
+
+    # Use re.sub to replace each Python code block
+    data = code_block_re.sub(replacer, data)
+
+    with open(file, "w") as f:
+        f.write(data)
+    return all_imports
 
 
 if __name__ == "__main__":

docs/docs_skeleton/src/theme/CodeBlock/index.js

@@ -34,13 +34,25 @@ function Imports({ imports }) {
 }
 
 export default function CodeBlockWrapper({ children, ...props }) {
+  // Initialize imports as an empty array
+  let imports = [];
+
+  // Check if children is a string
   if (typeof children === "string") {
-    return <CodeBlock {...props}>{children}</CodeBlock>;
+    // Search for an IMPORTS comment in the code
+    const match = /<!--IMPORTS:(.*?)-->\n/.exec(children);
+    if (match) {
+      imports = JSON.parse(match[1]);
+      children = children.replace(match[0], "");
+    }
+  } else if (children.imports) {
+    imports = children.imports;
   }
+
   return (
     <>
-      <CodeBlock {...props}>{children.content}</CodeBlock>
-      <Imports imports={children.imports} />
+      <CodeBlock {...props}>{children}</CodeBlock>
+      {imports.length > 0 && <Imports imports={imports} />}
     </>
   );
-}
+}