From ca6905b4e8aa7d20219d6e1f4687dc96be57400a Mon Sep 17 00:00:00 2001 From: shyam0508060409087 Date: Tue, 7 Jul 2026 02:04:37 +0530 Subject: [PATCH] feat: add configurable LLM providers via factory abstraction --- repo2readme/cli/main.py | 110 +++++++++++++++++++------ repo2readme/config.py | 65 +++++++++++---- repo2readme/llm/factory.py | 107 ++++++++++++++++++++++++ repo2readme/llm/init.py | 9 ++ repo2readme/readme/readme_generator.py | 14 +--- repo2readme/readme/reviewer_agent.py | 7 +- repo2readme/summarize/summary.py | 19 ++--- 7 files changed, 265 insertions(+), 66 deletions(-) create mode 100644 repo2readme/llm/factory.py create mode 100644 repo2readme/llm/init.py diff --git a/repo2readme/cli/main.py b/repo2readme/cli/main.py index 8780274..a4daf40 100644 --- a/repo2readme/cli/main.py +++ b/repo2readme/cli/main.py @@ -1,10 +1,13 @@ +import os + import click from rich import print as rprint from rich.progress import Progress + from repo2readme.config import get_api_keys, reset_api_keys -import os -from repo2readme.utils.tree import generate_tree +from repo2readme.llm.factory import configure_llm from repo2readme.utils.detect_language import detect_lang +from repo2readme.utils.tree import generate_tree @click.group() @@ -46,21 +49,53 @@ def main(): is_flag=True, help="Preview the analysis without making any API calls.", ) -def run(url, local, output, force, include_patterns, exclude_patterns, max_file_size_kb, dry_run): - """ Use --url for GitHub repo url and --local for local repo - """ +@click.option( + "--provider", + default=None, + help="LLM provider (groq, google, openai, anthropic, etc.)", +) +@click.option( + "--model", + default=None, + help="Model name", +) +@click.option( + "--base-url", + default=None, + help="Base URL for OpenAI-compatible APIs", +) +def run( + url, + local, + output, + force, + include_patterns, + exclude_patterns, + max_file_size_kb, + dry_run, + provider, + model, + base_url, +): + """Use --url for GitHub repo url and --local for local repo""" + if not url and not local: rprint("[red]Provide either --url or --local[/red]") return source = url if url else local - + from repo2readme.loaders.repo_loader import RepoLoader with Progress() as progress: task = progress.add_task("[cyan]Loading repository...", total=1) try: - loader = RepoLoader(source, include_patterns=include_patterns, exclude_patterns=exclude_patterns, max_file_size_kb=max_file_size_kb) + loader = RepoLoader( + source, + include_patterns=include_patterns, + exclude_patterns=exclude_patterns, + max_file_size_kb=max_file_size_kb, + ) files, root_path, loader_obj = loader.load() except Exception as e: rprint(f"[red]Failed to load repository: {e}[/red]") @@ -69,10 +104,13 @@ def run(url, local, output, force, include_patterns, exclude_patterns, max_file_ documents = [] for f in files: - documents.append({ - "content": f.page_content, - "metadata": f.metadata - }) + documents.append( + { + "content": f.page_content, + "metadata": f.metadata, + } + ) + tree = generate_tree(root_path) # Estimate token count (roughly 3 characters per token) @@ -91,21 +129,24 @@ def format_size(size_bytes): if dry_run: rprint("\n[bold]Repository Tree[/bold]\n") rprint(tree) + rprint("\n[bold]Files to be processed[/bold]\n") for doc in documents: rel_path = doc["metadata"].get("relative_path", "") rprint(f"✓ [green]{rel_path}[/green]") + rprint("\n[bold]Repository Analysis[/bold]\n") rprint(f"Files selected : {total_documents}") rprint(f"Estimated tokens : ~{estimated_tokens:,}") rprint(f"Request size : ~{format_size(total_size_bytes)}") + rprint("\n[green]Dry run complete.[/green]") rprint("[yellow]No API requests were made.[/yellow]") + if hasattr(loader_obj, "cleanup"): loader_obj.cleanup() return - # Normal execution: print estimation first rprint("\n[bold]Repository Analysis[/bold]\n") rprint(f"Files to summarize : {total_documents}") rprint(f"Estimated tokens : ~{estimated_tokens:,}") @@ -119,33 +160,49 @@ def format_size(size_bytes): return try: - groq_key, gemini_key = get_api_keys() - os.environ["GROQ_API_KEY"] = groq_key - os.environ["GOOGLE_API_KEY"] = gemini_key + # Backward-compatible configuration loading + get_api_keys() + + # Configure the selected LLM provider/model for the process. + configure_llm( + provider=provider, + model=model, + base_url=base_url, + ) + except Exception as e: - rprint(f"[red]Failed to configure API keys: {e}[/red]") + rprint(f"[red]Failed to configure LLM: {e}[/red]") return - from repo2readme.summarize.summary import summarize_file from repo2readme.readme.agent_workflow import workflow + from repo2readme.summarize.summary import summarize_file summaries = [] errors = [] + with Progress() as progress: - task = progress.add_task("[cyan]Generating summaries...[/cyan]", total=total_documents) + task = progress.add_task( + "[cyan]Generating summaries...[/cyan]", + total=total_documents, + ) + for doc in documents: meta = doc["metadata"] + try: lang = detect_lang(meta.get("file_type", "text")) + summary = summarize_file( file_path=meta["file_path"], language=lang, - content=doc["content"] + content=doc["content"], ) - + summaries.append(summary) + except Exception as e: errors.append(f"Error processing {meta.get('file_path')}: {e}") + progress.update(task, advance=1) rprint("[cyan]Generating README...[/cyan]") @@ -156,12 +213,12 @@ def format_size(size_bytes): "iteration_no": 0, "max_iterations": 3, "latest_readme": "", - 'best_score': 0.0, - "best_readme": "" + "best_score": 0.0, + "best_readme": "", } final_state = workflow.invoke(initial_state) - readme = final_state['best_readme'] + readme = final_state["best_readme"] if output is None: rprint("\n[green]Generated README:[/green]\n") @@ -182,6 +239,11 @@ def format_size(size_bytes): rprint(f"[green]Saved to {output}[/green]") + if errors: + rprint("\n[yellow]Some files could not be processed:[/yellow]") + for err in errors: + rprint(f"[yellow]- {err}[/yellow]") + finally: if hasattr(loader_obj, "cleanup"): loader_obj.cleanup() @@ -199,4 +261,4 @@ def reset(): if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/repo2readme/config.py b/repo2readme/config.py index ed0351b..596d23f 100644 --- a/repo2readme/config.py +++ b/repo2readme/config.py @@ -5,9 +5,21 @@ ENV_PATH = os.path.join(os.path.expanduser("~"), ".repo2readme_env.json") +# ------------------------------------------------------- +# Supported API Keys +# Extend this dictionary when adding new providers. +# ------------------------------------------------------- + +SUPPORTED_API_KEYS = { + "GROQ_API_KEY": "Groq", + "GOOGLE_API_KEY": "Google Gemini", +} + + def load_env(): if not os.path.exists(ENV_PATH): return {} + with open(ENV_PATH, "r") as f: return json.load(f) @@ -18,33 +30,54 @@ def save_env(data): def get_api_keys(): - env = load_env() + """ + Loads all configured API keys. - groq = env.get("GROQ_API_KEY") - gemini = env.get("GOOGLE_API_KEY") + If any required API key is missing, + the user is prompted only for the missing key. - if groq and gemini: - return groq, gemini + Returns + ------- + dict - rprint("[yellow]API keys are missing! Let's add them.[/yellow]\n") + Example: - if not groq: - groq = input("Enter your Groq API key: ").strip() - if not gemini: - gemini = input("Enter your Google Gemini API key: ").strip() + { + "GROQ_API_KEY": "...", + "GOOGLE_API_KEY": "..." + } + """ - env["GROQ_API_KEY"] = groq - env["GOOGLE_API_KEY"] = gemini + env = load_env() + + missing = [ + key + for key in SUPPORTED_API_KEYS + if not env.get(key) + ] + + if missing: + rprint("[yellow]Some API keys are missing.[/yellow]\n") - save_env(env) + for key in missing: + provider_name = SUPPORTED_API_KEYS[key] + env[key] = input( + f"Enter your {provider_name} API key: " + ).strip() - rprint("[green]API keys saved successfully![/green]") + save_env(env) - return groq, gemini + rprint("[green]API keys saved successfully![/green]") + + return { + key: env.get(key) + for key in SUPPORTED_API_KEYS + } def reset_api_keys(): if os.path.exists(ENV_PATH): os.remove(ENV_PATH) return True - return False + + return False \ No newline at end of file diff --git a/repo2readme/llm/factory.py b/repo2readme/llm/factory.py new file mode 100644 index 0000000..b76dc24 --- /dev/null +++ b/repo2readme/llm/factory.py @@ -0,0 +1,107 @@ +import os + +from langchain_groq import ChatGroq +from langchain_google_genai import ChatGoogleGenerativeAI + +# ------------------------------------------------------------ +# Default configuration +# (Preserves existing behaviour) +# ------------------------------------------------------------ + +DEFAULT_CONFIG = { + "summarizer": { + "provider": "groq", + "model": "openai/gpt-oss-120b", + "temperature": 0.2, + }, + "generator": { + "provider": "google", + "model": "gemini-2.5-flash", + "temperature": 0.2, + }, + "reviewer": { + "provider": "google", + "model": "gemini-2.5-flash", + "temperature": 0.2, + }, +} + +# ------------------------------------------------------------ +# Runtime overrides from CLI +# ------------------------------------------------------------ + +_RUNTIME_CONFIG = { + "provider": None, + "model": None, + "base_url": None, +} + + +def configure_llm(provider=None, model=None, base_url=None): + """ + Configure runtime overrides supplied from the CLI. + """ + + _RUNTIME_CONFIG["provider"] = provider + _RUNTIME_CONFIG["model"] = model + _RUNTIME_CONFIG["base_url"] = base_url + + +def create_llm( + stage: str, + provider: str | None = None, + model: str | None = None, + temperature: float | None = None, +): + """ + Creates a chat model for the requested pipeline stage. + """ + + if stage not in DEFAULT_CONFIG: + raise ValueError(f"Unknown pipeline stage: {stage}") + + config = DEFAULT_CONFIG[stage] + + provider = ( + provider + or _RUNTIME_CONFIG["provider"] + or config["provider"] + ) + + model = ( + model + or _RUNTIME_CONFIG["model"] + or config["model"] + ) + + temperature = ( + temperature + if temperature is not None + else config["temperature"] + ) + + provider = provider.lower() + + if provider == "groq": + return ChatGroq( + model=model, + api_key=os.getenv("GROQ_API_KEY"), + temperature=temperature, + ) + + if provider in ("google", "gemini"): + return ChatGoogleGenerativeAI( + model=model, + google_api_key=os.getenv("GOOGLE_API_KEY"), + temperature=temperature, + ) + + raise ValueError(f"Unsupported provider: {provider}") + + +def get_llm(stage: str): + """ + Returns the configured LLM for a pipeline stage. + """ + + return create_llm(stage) \ No newline at end of file diff --git a/repo2readme/llm/init.py b/repo2readme/llm/init.py new file mode 100644 index 0000000..38e5082 --- /dev/null +++ b/repo2readme/llm/init.py @@ -0,0 +1,9 @@ +from .factory import ( + create_llm, + DEFAULT_CONFIG, +) + +__all__ = [ + "create_llm", + "DEFAULT_CONFIG", +] \ No newline at end of file diff --git a/repo2readme/readme/readme_generator.py b/repo2readme/readme/readme_generator.py index f6b70e3..b165959 100644 --- a/repo2readme/readme/readme_generator.py +++ b/repo2readme/readme/readme_generator.py @@ -1,14 +1,13 @@ -from langchain_google_genai import ChatGoogleGenerativeAI -from langchain_groq import ChatGroq +from repo2readme.llm.factory import get_llm import os from dotenv import load_dotenv load_dotenv() from langchain_core.prompts import PromptTemplate from langchain_core.output_parsers import StrOutputParser from typing import List + def generate_readme(summaries:List[str],tree_structure:str,feedback:List[str],latest_readme:str): - api_key = os.getenv("GOOGLE_API_KEY") - model=ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key) + model = get_llm("generator") prompt = PromptTemplate( template=""" @@ -93,9 +92,4 @@ def generate_readme(summaries:List[str],tree_structure:str,feedback:List[str],la "latest_readme":latest_readme, "feedback": feedback }) - return response - - - - - + return response \ No newline at end of file diff --git a/repo2readme/readme/reviewer_agent.py b/repo2readme/readme/reviewer_agent.py index 0087a57..d7c3bc9 100644 --- a/repo2readme/readme/reviewer_agent.py +++ b/repo2readme/readme/reviewer_agent.py @@ -1,7 +1,7 @@ from langchain_core.output_parsers import PydanticOutputParser from pydantic import BaseModel,Field from langchain_core.prompts import PromptTemplate -from langchain_google_genai import ChatGoogleGenerativeAI +from repo2readme.llm.factory import get_llm import os from dotenv import load_dotenv @@ -11,8 +11,7 @@ class ReviewSchema(BaseModel): feedback:str=Field(description="Actionable comment how to improve the README") def readme_reviewer(readme:str): - api_key = os.getenv("GOOGLE_API_KEY") - model=ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key) + model = get_llm("reviewer") parser=PydanticOutputParser(pydantic_object=ReviewSchema) review_prompt=PromptTemplate( template=""" @@ -43,4 +42,4 @@ def readme_reviewer(readme:str): response=chain.invoke({ 'readme':readme }) - return response + return response \ No newline at end of file diff --git a/repo2readme/summarize/summary.py b/repo2readme/summarize/summary.py index 433273a..fcd5c3f 100644 --- a/repo2readme/summarize/summary.py +++ b/repo2readme/summarize/summary.py @@ -1,19 +1,14 @@ -from langchain_groq import ChatGroq +from repo2readme.llm.factory import get_llm from langchain_core.prompts import PromptTemplate from dotenv import load_dotenv -import os from langchain_core.output_parsers import JsonOutputParser load_dotenv() def create_summarizer(file_path: str, language: str, content: str): - model = ChatGroq( - model="openai/gpt-oss-120b", - api_key=os.getenv("GROQ_API_KEY"), - temperature=0.2, - ) - parser=JsonOutputParser() + model = get_llm("summarizer") + parser = JsonOutputParser() prompt = PromptTemplate( template=""" You are an expert code analyst. @@ -72,9 +67,9 @@ def create_summarizer(file_path: str, language: str, content: str): Return ONLY JSON. {format_instructions} """, - input_variables=["file_path", "language", "content"], - partial_variables={"format_instructions": parser.get_format_instructions()} -) + input_variables=["file_path", "language", "content"], + partial_variables={"format_instructions": parser.get_format_instructions()} + ) chain = prompt | model | parser return chain @@ -91,4 +86,4 @@ def summarize_file(file_path: str, language: str, content: str): except Exception as e: print(file_path) print("SUMMARY ERROR:", e) - return {"file_path": file_path, "error": str(e)} + return {"file_path": file_path, "error": str(e)} \ No newline at end of file