From 98fa9fc0d5d946f8f4baf473d557d2d5440b320a Mon Sep 17 00:00:00 2001 From: remy Date: Wed, 30 Apr 2025 19:00:50 +1000 Subject: [PATCH] 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. --- .env.sample | 5 +- .gitignore | 9 ++- main.py | 10 ++- nodes.py | 36 ++++++---- requirements.txt | 2 +- utils/call_llm.py | 136 +++++++++++++++++++++++++++++-------- utils/crawl_local_files.py | 106 ++++++++++++++++++----------- utils/fix_yaml.py | 20 ++++++ 8 files changed, 241 insertions(+), 83 deletions(-) create mode 100644 utils/fix_yaml.py diff --git a/.env.sample b/.env.sample index c26cdce..51bbc47 100644 --- a/.env.sample +++ b/.env.sample @@ -1,2 +1,5 @@ GEMINI_PROJECT_ID= -GITHUB_TOKEN= \ No newline at end of file +GEMINI_API_KEY= +GITHUB_TOKEN= +OPENROUTER_API_KEY = +OPENROUTER_MODEL = \ No newline at end of file diff --git a/.gitignore b/.gitignore index 7546330..08505f8 100644 --- a/.gitignore +++ b/.gitignore @@ -99,4 +99,11 @@ coverage/ llm_cache.json # Output files -output/ \ No newline at end of file +output/ + +# uv manage +pyproject.toml +uv.lock + +docs/*.pdf +docs/design-cn.md diff --git a/main.py b/main.py index e0ccc82..5e42ad8 100644 --- a/main.py +++ b/main.py @@ -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() @@ -61,6 +65,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": [], @@ -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() diff --git a/nodes.py b/nodes.py index 67ab034..d2045f2 100644 --- a/nodes.py +++ b/nodes.py @@ -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,10 +445,12 @@ 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 # The 'previous_chapters_summary' will be built progressively in the exec context diff --git a/requirements.txt b/requirements.txt index 06253bc..c08935b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,4 @@ requests>=2.28.0 gitpython>=3.1.0 google-cloud-aiplatform>=1.25.0 google-genai>=1.9.0 -python-dotenv>=1.0.0 +python-dotenv>=1.0.0 \ No newline at end of file diff --git a/utils/call_llm.py b/utils/call_llm.py index 0d794b4..714d4dd 100644 --- a/utils/call_llm.py +++ b/utils/call_llm.py @@ -3,86 +3,93 @@ 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 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: + 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] - + # 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", ""), ) - response_text = response.text + # 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 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: + 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: + 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 @@ -101,7 +108,7 @@ def call_llm(prompt: str, use_cache: bool = True) -> str: # return response.content[1].text # # Use OpenAI o1 -# def call_llm(prompt, use_cache: bool = True): +# def call_llm(prompt, use_cache: bool = True): # from openai import OpenAI # client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "your-api-key")) # r = client.chat.completions.create( @@ -115,11 +122,86 @@ 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?" - + # First call - should hit the API print("Making call...") response1 = call_llm(test_prompt, use_cache=False) print(f"Response: {response1}") - diff --git a/utils/crawl_local_files.py b/utils/crawl_local_files.py index 837215c..50c470d 100644 --- a/utils/crawl_local_files.py +++ b/utils/crawl_local_files.py @@ -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,53 +23,78 @@ def crawl_local_files(directory, include_patterns=None, exclude_patterns=None, m files_dict = {} - for root, _, files in os.walk(directory): - for filename in files: - filepath = os.path.join(root, filename) - + 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 - if use_relative_paths: - relpath = os.path.relpath(filepath, directory) - else: - relpath = filepath + 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 file matches any include pattern - included = False - if include_patterns: - for pattern in include_patterns: - if fnmatch.fnmatch(relpath, pattern): - included = True - break - else: - included = True - - # Check if file matches any exclude pattern - excluded = False - if exclude_patterns: - for pattern in exclude_patterns: - if fnmatch.fnmatch(relpath, pattern): - excluded = True - break + # 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) + + # Get path relative to directory if requested + if use_relative_paths: + relpath = os.path.relpath(filepath, directory) + else: + relpath = filepath - if not included or excluded: - continue - - # Check file size - if max_file_size and os.path.getsize(filepath) > max_file_size: - continue - - try: - with open(filepath, 'r', encoding='utf-8') as f: - content = f.read() - files_dict[relpath] = content - except Exception as e: - print(f"Warning: Could not read file {filepath}: {e}") + # Check if file matches any include pattern + included = False + if include_patterns: + for pattern in include_patterns: + if fnmatch.fnmatch(relpath, pattern): + included = True + break + else: + included = True + + # Check if file matches any exclude pattern + excluded = False + if exclude_patterns: + for pattern in exclude_patterns: + if fnmatch.fnmatch(relpath, pattern) or fnmatch.fnmatch(relpath, pattern.replace("/", "\\")): + print(relpath, pattern) + excluded = True + break + + if not included or excluded: + continue + + # Check file size + if max_file_size and os.path.getsize(filepath) > max_file_size: + continue + + try: + with open(filepath, 'r', encoding='utf-8') as f: + content = f.read() + files_dict[relpath] = content + except Exception as e: + print(f"Warning: Could not read file {filepath}: {e}") return {"files": files_dict} 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}") \ No newline at end of file diff --git a/utils/fix_yaml.py b/utils/fix_yaml.py new file mode 100644 index 0000000..1e7f55d --- /dev/null +++ b/utils/fix_yaml.py @@ -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) \ No newline at end of file