diff --git a/README.md b/README.md index 93b2429..8e37002 100644 --- a/README.md +++ b/README.md @@ -82,6 +82,8 @@ Repo2Readme/ ├── pyproject.toml ├── repo2readme/ ├── config.py + ├── llm/ + ├── factory.py ├── cli/ ├── main.py ├── loaders/ diff --git a/pyproject.toml b/pyproject.toml index 67bc2fc..f0cbb5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,8 @@ dependencies = [ "langchain-community>=0.4.1", "langchain-groq>=1.1.1", "langchain-google-genai>=4.1.2", + "langchain-openai>=1.0.0", + "langchain-anthropic>=1.0.0", "groq>=0.37.1", "google-genai>=1.56.0" diff --git a/repo2readme/cli/main.py b/repo2readme/cli/main.py index 1d51535..63e77f9 100644 --- a/repo2readme/cli/main.py +++ b/repo2readme/cli/main.py @@ -1,7 +1,11 @@ import click from rich import print as rprint from rich.progress import Progress -from repo2readme.config import get_api_keys, reset_api_keys +from repo2readme.config import ( + get_api_keys, + get_api_key, + reset_api_keys, +) import os from repo2readme.utils.tree import generate_tree from repo2readme.utils.detect_language import detect_lang @@ -48,7 +52,22 @@ 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): +@click.option( + "--provider", + default=None, + help="LLM provider (groq, google, openai, anthropic, openrouter, ollama, etc.)", +) +@click.option( + "--model", + default=None, + help="LLM model name", +) +@click.option( + "--base-url", + default=None, + help="Base URL for OpenAI-compatible providers", +) +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: @@ -121,9 +140,26 @@ 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 + if provider: + api_key = get_api_key(provider) + + provider_env = { + "groq": "GROQ_API_KEY", + "google": "GOOGLE_API_KEY", + "gemini": "GOOGLE_API_KEY", + "openai": "OPENAI_API_KEY", + "anthropic": "ANTHROPIC_API_KEY", + "openrouter": "OPENROUTER_API_KEY", + "together": "TOGETHER_API_KEY", + } + + os.environ[provider_env[provider.lower()]] = api_key + + else: + groq_key, gemini_key = get_api_keys() + os.environ["GROQ_API_KEY"] = groq_key + os.environ["GOOGLE_API_KEY"] = gemini_key + except Exception as e: rprint(f"[red]Failed to configure API keys: {e}[/red]") return @@ -143,11 +179,21 @@ def process_document(doc): 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"] - ) + if provider or model or base_url: + summary = summarize_file( + file_path=meta["file_path"], + language=lang, + content=doc["content"], + provider=provider, + model_name=model, + base_url=base_url, + ) + else: + summary = summarize_file( + file_path=meta["file_path"], + language=lang, + content=doc["content"], + ) with summaries_lock: summaries.append(summary) except Exception as e: @@ -175,7 +221,10 @@ def process_document(doc): "max_iterations": 3, "latest_readme": "", 'best_score': 0.0, - "best_readme": "" + "best_readme": "", + "provider": provider, + "model": model, + "base_url": base_url, } final_state = workflow.invoke(initial_state) diff --git a/repo2readme/config.py b/repo2readme/config.py index ed0351b..1ebd79f 100644 --- a/repo2readme/config.py +++ b/repo2readme/config.py @@ -16,31 +16,54 @@ def save_env(data): with open(ENV_PATH, "w") as f: json.dump(data, f, indent=4) - def get_api_keys(): env = load_env() groq = env.get("GROQ_API_KEY") gemini = env.get("GOOGLE_API_KEY") - if groq and gemini: - return groq, gemini - - rprint("[yellow]API keys are missing! Let's add them.[/yellow]\n") - if not groq: - groq = input("Enter your Groq API key: ").strip() + groq = get_api_key("groq") + if not gemini: - gemini = input("Enter your Google Gemini API key: ").strip() + gemini = get_api_key("google") + + return groq, gemini + +def get_api_key(provider: str): + env = load_env() - env["GROQ_API_KEY"] = groq - env["GOOGLE_API_KEY"] = gemini + provider_map = { + "groq": "GROQ_API_KEY", + "google": "GOOGLE_API_KEY", + "gemini": "GOOGLE_API_KEY", + "openai": "OPENAI_API_KEY", + "anthropic": "ANTHROPIC_API_KEY", + "openrouter": "OPENROUTER_API_KEY", + "together": "TOGETHER_API_KEY", + } + provider = provider.lower() + + if provider not in provider_map: + raise ValueError(f"Unsupported provider: {provider}") + + env_var = provider_map[provider] + api_key = env.get(env_var) + + if api_key: + return api_key + + rprint(f"[yellow]{provider} API key is missing![/yellow]\n") + + api_key = input(f"Enter your {provider} API key: ").strip() + + env[env_var] = api_key save_env(env) - rprint("[green]API keys saved successfully![/green]") + rprint("[green]API key saved successfully![/green]") - return groq, gemini + return api_key def reset_api_keys(): diff --git a/repo2readme/llm/__init__.py b/repo2readme/llm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/repo2readme/llm/factory.py b/repo2readme/llm/factory.py new file mode 100644 index 0000000..13dc39a --- /dev/null +++ b/repo2readme/llm/factory.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import os + +from langchain_core.language_models.chat_models import BaseChatModel + +def create_llm( + provider: str, + model: str, + api_key: str | None = None, + base_url: str | None = None, + **kwargs, +) -> BaseChatModel: + """ + Factory function to create a LangChain chat model. + + Parameters + ---------- + provider : str + LLM provider name. + model : str + Model name. + api_key : str | None + Optional API key. Falls back to environment variable. + base_url : str | None + Optional base URL for OpenAI-compatible APIs. + """ + + provider = provider.lower() + + if provider == "groq": + from langchain_groq import ChatGroq + return ChatGroq( + model=model, + api_key=api_key or os.getenv("GROQ_API_KEY"), + **kwargs, + ) + + elif provider in ("google", "gemini"): + from langchain_google_genai import ChatGoogleGenerativeAI + return ChatGoogleGenerativeAI( + model=model, + google_api_key=api_key or os.getenv("GOOGLE_API_KEY"), + **kwargs, + ) + + elif provider == "openai": + from langchain_openai import ChatOpenAI + return ChatOpenAI( + model=model, + api_key=api_key or os.getenv("OPENAI_API_KEY"), + base_url=base_url, + **kwargs, + ) + + elif provider == "anthropic": + from langchain_anthropic import ChatAnthropic + return ChatAnthropic( + model=model, + api_key=api_key or os.getenv("ANTHROPIC_API_KEY"), + **kwargs, + ) + + elif provider == "openrouter": + from langchain_openai import ChatOpenAI + return ChatOpenAI( + model=model, + api_key=api_key or os.getenv("OPENROUTER_API_KEY"), + base_url=base_url or "https://openrouter.ai/api/v1", + **kwargs, + ) + + else: + raise ValueError( + f"Unsupported provider '{provider}'." + ) \ No newline at end of file diff --git a/repo2readme/readme/agent_workflow.py b/repo2readme/readme/agent_workflow.py index 8ba4afe..1fea4a4 100644 --- a/repo2readme/readme/agent_workflow.py +++ b/repo2readme/readme/agent_workflow.py @@ -13,6 +13,9 @@ class ReadmeState(TypedDict): best_score:float iteration_no:int max_iterations:int + provider: str | None + model: str | None + base_url: str | None def generate_readme_node(state:ReadmeState): latest_readme = ( @@ -22,14 +25,17 @@ def generate_readme_node(state:ReadmeState): summaries=state['summaries'], tree_structure=state['tree_structure'], feedback=state['feedback'], - latest_readme=latest_readme + latest_readme=latest_readme, + provider=state["provider"], + model_name=state["model"], + base_url=state["base_url"] ) return { 'readme':[readme] } def readme_reviewer_node(state:ReadmeState): latest_readme=state['readme'][-1] - review=readme_reviewer(latest_readme) + review=readme_reviewer(latest_readme,provider=state["provider"], model_name=state["model"], base_url=state["base_url"]) best_score=state['best_score'] best_readme=state['best_readme'] if review.score>best_score: diff --git a/repo2readme/readme/readme_generator.py b/repo2readme/readme/readme_generator.py index f6b70e3..b200cf6 100644 --- a/repo2readme/readme/readme_generator.py +++ b/repo2readme/readme/readme_generator.py @@ -1,14 +1,15 @@ -from langchain_google_genai import ChatGoogleGenerativeAI -from langchain_groq import ChatGroq -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) +from repo2readme.llm.factory import create_llm +def generate_readme(summaries:List[str],tree_structure:str,feedback:List[str],latest_readme:str,provider:str, model_name:str, base_url:str): + model = create_llm( + provider= provider or "groq", + model= model_name or "openai/gpt-oss-120b", + base_url=base_url +) prompt = PromptTemplate( template=""" diff --git a/repo2readme/readme/reviewer_agent.py b/repo2readme/readme/reviewer_agent.py index 0087a57..0f9c449 100644 --- a/repo2readme/readme/reviewer_agent.py +++ b/repo2readme/readme/reviewer_agent.py @@ -1,18 +1,26 @@ from langchain_core.output_parsers import PydanticOutputParser from pydantic import BaseModel,Field from langchain_core.prompts import PromptTemplate -from langchain_google_genai import ChatGoogleGenerativeAI import os from dotenv import load_dotenv +from repo2readme.llm.factory import create_llm load_dotenv() class ReviewSchema(BaseModel): score:float=Field(description="A score between 1 and 10 based on the quality of the README") 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) +def readme_reviewer( + readme: str, + provider: str | None = None, + model_name: str | None = None, + base_url: str | None = None, +): + model = create_llm( + provider=provider or "google", + model=model_name or "gemini-2.5-flash", + base_url=base_url, +) parser=PydanticOutputParser(pydantic_object=ReviewSchema) review_prompt=PromptTemplate( template=""" diff --git a/repo2readme/summarize/summary.py b/repo2readme/summarize/summary.py index 433273a..301955d 100644 --- a/repo2readme/summarize/summary.py +++ b/repo2readme/summarize/summary.py @@ -1,17 +1,18 @@ -from langchain_groq import ChatGroq from langchain_core.prompts import PromptTemplate from dotenv import load_dotenv import os from langchain_core.output_parsers import JsonOutputParser +from repo2readme.llm.factory import create_llm 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, +def create_summarizer(file_path, language, + content,provider=None, model_name=None, base_url=None,): + model = create_llm( + provider=provider or "groq", + model=model_name or "openai/gpt-oss-120b", + base_url=base_url, ) parser=JsonOutputParser() prompt = PromptTemplate( @@ -80,9 +81,16 @@ def create_summarizer(file_path: str, language: str, content: str): return chain -def summarize_file(file_path: str, language: str, content: str): +def summarize_file( + file_path, + language, + content, + provider=None, + model_name=None, + base_url=None, +): try: - chain = create_summarizer(file_path, language, content) + chain = create_summarizer(file_path, language, content, provider, model_name, base_url,) return chain.invoke({ "file_path": file_path.replace("\\", "/"), "language": language,