Skip to content

Commit 51b64f8

Browse files
committed
feat: add --max-abstractions flag to control number of identified abstractions
1 parent ba7d863 commit 51b64f8

File tree

3 files changed

+14
-4
lines changed

3 files changed

+14
-4
lines changed

README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,8 @@ This is a tutorial project of [Pocket Flow](https://github.com/The-Pocket/Pocket
112112
- `-e, --exclude` - Files to exclude (e.g., "tests/*" "docs/*")
113113
- `-s, --max-size` - Maximum file size in bytes (default: 100KB)
114114
- `--language` - Language for the generated tutorial (default: "english")
115+
- `--max-abstractions` - Maximum number of abstractions to identify (default: 10)
116+
- `--no-cache` - Disable LLM response caching (default: caching enabled)
115117

116118
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).
117119

main.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ def main():
4040
parser.add_argument("--language", default="english", help="Language for the generated tutorial (default: english)")
4141
# Add use_cache parameter to control LLM caching
4242
parser.add_argument("--no-cache", action="store_true", help="Disable LLM response caching (default: caching enabled)")
43+
# Add max_abstraction_num parameter to control the number of abstractions
44+
parser.add_argument("--max-abstractions", type=int, default=10, help="Maximum number of abstractions to identify (default: 20)")
4345

4446
args = parser.parse_args()
4547

@@ -68,6 +70,9 @@ def main():
6870

6971
# Add use_cache flag (inverse of no-cache flag)
7072
"use_cache": not args.no_cache,
73+
74+
# Add max_abstraction_num parameter
75+
"max_abstraction_num": args.max_abstractions,
7176

7277
# Outputs will be populated by the nodes
7378
"files": [],

nodes.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ def prep(self, shared):
8686
project_name = shared["project_name"] # Get project name
8787
language = shared.get("language", "english") # Get language
8888
use_cache = shared.get("use_cache", True) # Get use_cache flag, default to True
89+
max_abstraction_num = shared.get("max_abstraction_num", 10) # Get max_abstraction_num, default to 20
8990

9091
# Helper to create context from files, respecting limits (basic example)
9192
def create_llm_context(files_data):
@@ -110,7 +111,8 @@ def create_llm_context(files_data):
110111
project_name,
111112
language,
112113
use_cache,
113-
) # Return use_cache
114+
max_abstraction_num,
115+
) # Return all parameters
114116

115117
def exec(self, prep_res):
116118
(
@@ -120,7 +122,8 @@ def exec(self, prep_res):
120122
project_name,
121123
language,
122124
use_cache,
123-
) = prep_res # Unpack use_cache
125+
max_abstraction_num,
126+
) = prep_res # Unpack all parameters
124127
print(f"Identifying abstractions using LLM...")
125128

126129
# Add language instruction and hints only if not English
@@ -140,7 +143,7 @@ def exec(self, prep_res):
140143
{context}
141144
142145
{language_instruction}Analyze the codebase context.
143-
Identify the top 5-20 core most important abstractions to help those new to the codebase.
146+
Identify the top 5-{max_abstraction_num} core most important abstractions to help those new to the codebase.
144147
145148
For each abstraction, provide:
146149
1. A concise `name`{name_lang_hint}.
@@ -167,7 +170,7 @@ def exec(self, prep_res):
167170
Another core concept, similar to a blueprint for objects.{desc_lang_hint}
168171
file_indices:
169172
- 5 # path/to/another.js
170-
# ... up to 20 abstractions
173+
# ... up to {max_abstraction_num} abstractions
171174
```"""
172175
response = call_llm(prompt, use_cache=use_cache) # Pass use_cache parameter
173176

0 commit comments

Comments
 (0)