fix: improve file filtering, add new utility,

- Improved the speed of file filtering in `crawl_local_files.py` with folder-level exclusion
- Added `fix_yaml.py` utility for YAML indentation fixes
- Updated `nodes.py` to support up to 20 core abstractions
- add option for no cache.
This commit is contained in:
remy
2025-04-30 19:00:50 +10:00
parent 4da93741fe
commit 98fa9fc0d5
8 changed files with 241 additions and 83 deletions
+3
View File
@@ -1,2 +1,5 @@
GEMINI_PROJECT_ID=<GEMINI_PROJECT_ID>
GEMINI_API_KEY=<GEMINI_API_KEY>
GITHUB_TOKEN=<GITHUB_TOKEN>
OPENROUTER_API_KEY = <OPENROUTER_API_KEY>
OPENROUTER_MODEL = <OPENROUTER_MODEL>
+7
View File
@@ -100,3 +100,10 @@ llm_cache.json
# Output files
output/
# uv manage
pyproject.toml
uv.lock
docs/*.pdf
docs/design-cn.md
+9 -1
View File
@@ -14,8 +14,10 @@ DEFAULT_INCLUDE_PATTERNS = {
}
DEFAULT_EXCLUDE_PATTERNS = {
"assets/*", "data/*", "examples/*", "images/*", "public/*", "static/*", "temp/*",
"docs/*",
"venv/*", ".venv/*", "*test*", "tests/*", "docs/*", "examples/*", "v1/*",
"dist/*", "build/*", "experimental/*", "deprecated/*",
"dist/*", "build/*", "experimental/*", "deprecated/*", "misc/*",
"legacy/*", ".git/*", ".github/*", ".next/*", ".vscode/*", "obj/*", "bin/*", "node_modules/*", "*.log"
}
@@ -36,6 +38,8 @@ def main():
parser.add_argument("-s", "--max-size", type=int, default=100000, help="Maximum file size in bytes (default: 100000, about 100KB).")
# Add language parameter for multi-language support
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)")
args = parser.parse_args()
@@ -62,6 +66,9 @@ def main():
# Add language for multi-language support
"language": args.language,
# Add use_cache flag (inverse of no-cache flag)
"use_cache": not args.no_cache,
# Outputs will be populated by the nodes
"files": [],
"abstractions": [],
@@ -73,6 +80,7 @@ def main():
# Display starting message with repository/directory and language
print(f"Starting tutorial generation for: {args.repo or args.dir} in {args.language.capitalize()} language")
print(f"LLM caching: {'Disabled' if args.no_cache else 'Enabled'}")
# Create the flow instance
tutorial_flow = create_tutorial_flow()
+22 -12
View File
@@ -1,9 +1,12 @@
import os
import re
import yaml
from pocketflow import Node, BatchNode
from utils.crawl_github_files import crawl_github_files
from utils.call_llm import call_llm
from utils.crawl_local_files import crawl_local_files
from utils.fix_yaml import add_indentation
# Helper to get content for specific file indices
def get_content_for_indices(files_data, indices):
@@ -79,6 +82,7 @@ class IdentifyAbstractions(Node):
files_data = shared["files"]
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
# Helper to create context from files, respecting limits (basic example)
def create_llm_context(files_data):
@@ -94,10 +98,10 @@ class IdentifyAbstractions(Node):
context, file_info = create_llm_context(files_data)
# Format file info for the prompt (comment is just a hint for LLM)
file_listing_for_prompt = "\n".join([f"- {idx} # {path}" for idx, path in file_info])
return context, file_listing_for_prompt, len(files_data), project_name, language # Return language
return context, file_listing_for_prompt, len(files_data), project_name, language, use_cache # Return use_cache
def exec(self, prep_res):
context, file_listing_for_prompt, file_count, project_name, language = prep_res # Unpack project name and language
context, file_listing_for_prompt, file_count, project_name, language, use_cache = prep_res # Unpack use_cache
print(f"Identifying abstractions using LLM...")
# Add language instruction and hints only if not English
@@ -117,7 +121,7 @@ Codebase Context:
{context}
{language_instruction}Analyze the codebase context.
Identify the top 5-10 core most important abstractions to help those new to the codebase.
Identify the top 5-20 core most important abstractions to help those new to the codebase.
For each abstraction, provide:
1. A concise `name`{name_lang_hint}.
@@ -144,12 +148,14 @@ 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 10 abstractions
# ... up to 20 abstractions
```"""
response = call_llm(prompt)
response = call_llm(prompt, use_cache=use_cache) # Pass use_cache parameter
# --- Validation ---
yaml_str = response.strip().split("```yaml")[1].split("```")[0].strip()
# add whitespace to fix llm generation error(except -)
yaml_str = add_indentation(yaml_str)
abstractions = yaml.safe_load(yaml_str)
if not isinstance(abstractions, list):
@@ -203,6 +209,7 @@ class AnalyzeRelationships(Node):
files_data = shared["files"]
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
# Create context with abstraction names, indices, descriptions, and relevant file snippets
context = "Identified Abstractions:\n"
@@ -230,10 +237,10 @@ class AnalyzeRelationships(Node):
)
context += file_context_str
return context, "\n".join(abstraction_info_for_prompt), project_name, language # Return language
return context, "\n".join(abstraction_info_for_prompt), project_name, language, use_cache # Return use_cache
def exec(self, prep_res):
context, abstraction_listing, project_name, language = prep_res # Unpack project name and language
context, abstraction_listing, project_name, language, use_cache = prep_res # Unpack use_cache
print(f"Analyzing relationships using LLM...")
# Add language instruction and hints only if not English
@@ -339,6 +346,7 @@ class OrderChapters(Node):
relationships = shared["relationships"] # Summary/label might be translated
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
# Prepare context for the LLM
abstraction_info_for_prompt = []
@@ -363,10 +371,10 @@ class OrderChapters(Node):
if language.lower() != "english":
list_lang_note = f" (Names might be in {language.capitalize()})"
return abstraction_listing, context, len(abstractions), project_name, list_lang_note
return abstraction_listing, context, len(abstractions), project_name, list_lang_note, use_cache # Return use_cache
def exec(self, prep_res):
abstraction_listing, context, num_abstractions, project_name, list_lang_note = prep_res
abstraction_listing, context, num_abstractions, project_name, list_lang_note, use_cache = prep_res # Unpack use_cache
print("Determining chapter order using LLM...")
# No language variation needed here in prompt instructions, just ordering based on structure
# The input names might be translated, hence the note.
@@ -437,9 +445,11 @@ Now, provide the YAML output:
class WriteChapters(BatchNode):
def prep(self, shared):
chapter_order = shared["chapter_order"] # List of indices
abstractions = shared["abstractions"] # List of dicts, name/desc potentially translated
files_data = shared["files"]
language = shared.get("language", "english") # Get language
abstractions = shared["abstractions"] # List of {"name": str, "description": str, "files": [int]}
files_data = shared["files"] # List of (path, content) tuples
project_name = shared["project_name"]
language = shared.get("language", "english")
use_cache = shared.get("use_cache", True) # Get use_cache flag, default to True
# Get already written chapters to provide context
# We store them temporarily during the batch run, not in shared memory yet
+100 -18
View File
@@ -3,23 +3,29 @@ import os
import logging
import json
from datetime import datetime
import requests
# Configure logging
log_directory = os.getenv("LOG_DIR", "logs")
os.makedirs(log_directory, exist_ok=True)
log_file = os.path.join(log_directory, f"llm_calls_{datetime.now().strftime('%Y%m%d')}.log")
log_file = os.path.join(
log_directory, f"llm_calls_{datetime.now().strftime('%Y%m%d')}.log"
)
# Set up logger
logger = logging.getLogger("llm_logger")
logger.setLevel(logging.INFO)
logger.propagate = False # Prevent propagation to root logger
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
file_handler.setFormatter(
logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
)
logger.addHandler(file_handler)
# Simple cache configuration
cache_file = "llm_cache.json"
# By default, we Google Gemini 2.5 pro, as it shows great performance for code understanding
def call_llm(prompt: str, use_cache: bool = True) -> str:
# Log the prompt
@@ -31,7 +37,7 @@ def call_llm(prompt: str, use_cache: bool = True) -> str:
cache = {}
if os.path.exists(cache_file):
try:
with open(cache_file, 'r') as f:
with open(cache_file, "r") as f:
cache = json.load(f)
except:
logger.warning(f"Failed to load cache, starting with empty cache")
@@ -42,21 +48,21 @@ def call_llm(prompt: str, use_cache: bool = True) -> str:
return cache[prompt]
# Call the LLM if not in cache or cache disabled
client = genai.Client(
vertexai=True,
# TODO: change to your own project id and location
project=os.getenv("GEMINI_PROJECT_ID", "your-project-id"),
location=os.getenv("GEMINI_LOCATION", "us-central1")
)
# You can comment the previous line and use the AI Studio key instead:
# client = genai.Client(
# api_key=os.getenv("GEMINI_API_KEY", "your-api_key"),
# vertexai=True,
# # TODO: change to your own project id and location
# project=os.getenv("GEMINI_PROJECT_ID", "your-project-id"),
# location=os.getenv("GEMINI_LOCATION", "us-central1")
# )
model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25")
response = client.models.generate_content(
model=model,
contents=[prompt]
# You can comment the previous line and use the AI Studio key instead:
client = genai.Client(
api_key=os.getenv("GEMINI_API_KEY", ""),
)
# model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25")
model = os.getenv("GEMINI_MODEL", "gemini-2.0-flash-exp")
response = client.models.generate_content(model=model, contents=[prompt])
response_text = response.text
# Log the response
@@ -68,7 +74,7 @@ def call_llm(prompt: str, use_cache: bool = True) -> str:
cache = {}
if os.path.exists(cache_file):
try:
with open(cache_file, 'r') as f:
with open(cache_file, "r") as f:
cache = json.load(f)
except:
pass
@@ -76,13 +82,14 @@ def call_llm(prompt: str, use_cache: bool = True) -> str:
# Add to cache and save
cache[prompt] = response_text
try:
with open(cache_file, 'w') as f:
with open(cache_file, "w") as f:
json.dump(cache, f)
except Exception as e:
logger.error(f"Failed to save cache: {e}")
return response_text
# # Use Anthropic Claude 3.7 Sonnet Extended Thinking
# def call_llm(prompt, use_cache: bool = True):
# from anthropic import Anthropic
@@ -115,6 +122,82 @@ def call_llm(prompt: str, use_cache: bool = True) -> str:
# )
# return r.choices[0].message.content
# Use OpenRouter API
# def call_llm(prompt: str, use_cache: bool = True) -> str:
# # Log the prompt
# logger.info(f"PROMPT: {prompt}")
# # Check cache if enabled
# if use_cache:
# # Load cache from disk
# cache = {}
# if os.path.exists(cache_file):
# try:
# with open(cache_file, "r") as f:
# cache = json.load(f)
# except:
# logger.warning(f"Failed to load cache, starting with empty cache")
# # Return from cache if exists
# if prompt in cache:
# logger.info(f"RESPONSE: {cache[prompt]}")
# return cache[prompt]
# # OpenRouter API configuration
# api_key = os.getenv("OPENROUTER_API_KEY", "")
# model = os.getenv("OPENROUTER_MODEL", "google/gemini-2.0-flash-exp:free")
# headers = {
# "Authorization": f"Bearer {api_key}",
# }
# data = {
# "model": model,
# "messages": [{"role": "user", "content": prompt}]
# }
# response = requests.post(
# "https://openrouter.ai/api/v1/chat/completions",
# headers=headers,
# json=data
# )
# if response.status_code != 200:
# error_msg = f"OpenRouter API call failed with status {response.status_code}: {response.text}"
# logger.error(error_msg)
# raise Exception(error_msg)
# try:
# response_text = response.json()["choices"][0]["message"]["content"]
# except Exception as e:
# error_msg = f"Failed to parse OpenRouter response: {e}; Response: {response.text}"
# logger.error(error_msg)
# raise Exception(error_msg)
# # Log the response
# logger.info(f"RESPONSE: {response_text}")
# # Update cache if enabled
# if use_cache:
# # Load cache again to avoid overwrites
# cache = {}
# if os.path.exists(cache_file):
# try:
# with open(cache_file, "r") as f:
# cache = json.load(f)
# except:
# pass
# # Add to cache and save
# cache[prompt] = response_text
# try:
# with open(cache_file, "w") as f:
# json.dump(cache, f)
# except Exception as e:
# logger.error(f"Failed to save cache: {e}")
# return response_text
if __name__ == "__main__":
test_prompt = "Hello, how are you?"
@@ -122,4 +205,3 @@ if __name__ == "__main__":
print("Making call...")
response1 = call_llm(test_prompt, use_cache=False)
print(f"Response: {response1}")
+31 -3
View File
@@ -1,9 +1,12 @@
import os
import fnmatch
from tqdm import tqdm
def crawl_local_files(directory, include_patterns=None, exclude_patterns=None, max_file_size=None, use_relative_paths=True):
"""
Crawl files in a local directory with similar interface as crawl_github_files.
Implements efficient folder-level filtering to skip entire directories that match exclude patterns,
significantly improving performance when excluding large directory trees.
Args:
directory (str): Path to local directory
@@ -20,7 +23,31 @@ def crawl_local_files(directory, include_patterns=None, exclude_patterns=None, m
files_dict = {}
for root, _, files in os.walk(directory):
for root, dirs, files in os.walk(directory):
print(f"root: {root}")
# Check if current directory should be excluded
if exclude_patterns:
# Get path relative to directory if requested
rel_root = os.path.relpath(root, directory) if use_relative_paths else root
# Handle the case where rel_root is the current directory
if rel_root == '.':
rel_root = ''
# Check if directory matches any exclude pattern
for pattern in exclude_patterns:
# Normalize pattern to handle both forward and backward slashes
norm_pattern = pattern.replace("/", os.path.sep)
# Check if the directory matches the pattern
if fnmatch.fnmatch(rel_root, norm_pattern) or \
fnmatch.fnmatch(os.path.join(rel_root, ''), norm_pattern + os.path.sep):
# Skip this directory and all subdirectories
dirs[:] = [] # Clear dirs list to prevent further traversal
print(f"Skipping directory: {rel_root} (matches pattern {pattern})")
break
else:
# print(root)
for filename in files:
filepath = os.path.join(root, filename)
@@ -44,7 +71,8 @@ def crawl_local_files(directory, include_patterns=None, exclude_patterns=None, m
excluded = False
if exclude_patterns:
for pattern in exclude_patterns:
if fnmatch.fnmatch(relpath, pattern):
if fnmatch.fnmatch(relpath, pattern) or fnmatch.fnmatch(relpath, pattern.replace("/", "\\")):
print(relpath, pattern)
excluded = True
break
@@ -66,7 +94,7 @@ def crawl_local_files(directory, include_patterns=None, exclude_patterns=None, m
if __name__ == "__main__":
print("--- Crawling parent directory ('..') ---")
files_data = crawl_local_files("..", exclude_patterns={"*.pyc", "__pycache__/*", ".git/*", "output/*"})
files_data = crawl_local_files("..", exclude_patterns={"*.pyc", "__pycache__/*",".venv/*", ".git/*","docs/*", "output/*"})
print(f"Found {len(files_data['files'])} files:")
for path in files_data["files"]:
print(f" {path}")
+20
View File
@@ -0,0 +1,20 @@
import re
def add_indentation(text):
# This pattern matches lines that don't start with a hyphen or whitespace
pattern = r'^(?![-\s])(.*)$'
# Replace with 4 spaces followed by the captured content
result = re.sub(pattern, r' \1', text, flags=re.MULTILINE)
return result
if __name__ == "__main__":
# Example usage
text = """This line will be indented
- This line won't be indented
This line won't be indented either
Another line that will be indented"""
indented_text = add_indentation(text)