From 51b64f8d282fe73be4fadc6e5eef8094d051fd19 Mon Sep 17 00:00:00 2001
From: remy <redreamality@gmail.com>
Date: Thu, 1 May 2025 15:43:26 +1000
Subject: [PATCH] feat: add --max-abstractions flag to control number of
 identified abstractions

---
 README.md |  2 ++
 main.py   |  5 +++++
 nodes.py  | 11 +++++++----
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index c6d0ef1..4ef2e5d 100644
--- a/README.md
+++ b/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).
 
diff --git a/main.py b/main.py
index 5e42ad8..3a094d3 100644
--- a/main.py
+++ b/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": [],
diff --git a/nodes.py b/nodes.py
index f6be367..05e0fca 100644
--- a/nodes.py
+++ b/nodes.py
@@ -86,6 +86,7 @@ class IdentifyAbstractions(Node):
         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 @@ class IdentifyAbstractions(Node):
             project_name,
             language,
             use_cache,
-        )  # Return use_cache
+            max_abstraction_num,
+        )  # Return all parameters
 
     def exec(self, prep_res):
         (
@@ -120,7 +122,8 @@ class IdentifyAbstractions(Node):
             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 @@ Codebase Context:
 {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 @@ Format the output as a YAML list of dictionaries:
     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