feat: add --max-abstractions flag to control number of identified abstractions

Author: redreamality

README.md

@@ -112,6 +112,8 @@ This is a tutorial project of [Pocket Flow](https://github.com/The-Pocket/Pocket
     - `-e, --exclude` - Files to exclude (e.g., "tests/*" "docs/*")
     - `-s, --max-size` - Maximum file size in bytes (default: 100KB)
     - `--language` - Language for the generated tutorial (default: "english")
+    - `--max-abstractions` - Maximum number of abstractions to identify (default: 10)
+    - `--no-cache` - Disable LLM response caching (default: caching enabled)
 
 The application will crawl the repository, analyze the codebase structure, generate tutorial content in the specified language, and save the output in the specified directory (default: ./output).
 

main.py

@@ -40,6 +40,8 @@ def main():
     parser.add_argument("--language", default="english", help="Language for the generated tutorial (default: english)")
     # Add use_cache parameter to control LLM caching
     parser.add_argument("--no-cache", action="store_true", help="Disable LLM response caching (default: caching enabled)")
+    # Add max_abstraction_num parameter to control the number of abstractions
+    parser.add_argument("--max-abstractions", type=int, default=10, help="Maximum number of abstractions to identify (default: 20)")
 
     args = parser.parse_args()
 
@@ -68,6 +70,9 @@ def main():
 
         # Add use_cache flag (inverse of no-cache flag)
         "use_cache": not args.no_cache,
+        
+        # Add max_abstraction_num parameter
+        "max_abstraction_num": args.max_abstractions,
 
         # Outputs will be populated by the nodes
         "files": [],

nodes.py

@@ -86,6 +86,7 @@ def prep(self, shared):
         project_name = shared["project_name"]  # Get project name
         language = shared.get("language", "english")  # Get language
         use_cache = shared.get("use_cache", True)  # Get use_cache flag, default to True
+        max_abstraction_num = shared.get("max_abstraction_num", 10)  # Get max_abstraction_num, default to 20
 
         # Helper to create context from files, respecting limits (basic example)
         def create_llm_context(files_data):
@@ -110,7 +111,8 @@ def create_llm_context(files_data):
             project_name,
             language,
             use_cache,
-        )  # Return use_cache
+            max_abstraction_num,
+        )  # Return all parameters
 
     def exec(self, prep_res):
         (
@@ -120,7 +122,8 @@ def exec(self, prep_res):
             project_name,
             language,
             use_cache,
-        ) = prep_res  # Unpack use_cache
+            max_abstraction_num,
+        ) = prep_res  # Unpack all parameters
         print(f"Identifying abstractions using LLM...")
 
         # Add language instruction and hints only if not English
@@ -140,7 +143,7 @@ def exec(self, prep_res):
 {context}
 
 {language_instruction}Analyze the codebase context.
-Identify the top 5-20 core most important abstractions to help those new to the codebase.
+Identify the top 5-{max_abstraction_num} core most important abstractions to help those new to the codebase.
 
 For each abstraction, provide:
 1. A concise `name`{name_lang_hint}.
@@ -167,7 +170,7 @@ def exec(self, prep_res):
     Another core concept, similar to a blueprint for objects.{desc_lang_hint}
   file_indices:
     - 5 # path/to/another.js
-# ... up to 20 abstractions
+# ... up to {max_abstraction_num} abstractions
 ```"""
         response = call_llm(prompt, use_cache=use_cache)  # Pass use_cache parameter