docs.py

import os
import logging
import requests
import gradio as gr
from groq import Groq
from urllib.parse import urlparse

# Configure logging
logging.basicConfig(level=logging.INFO)


def extract_repo_id(input_str: str) -> str:
    """
    Extract the repository ID in the format 'user/repo' from a URL or plain string.
    """
    input_str = input_str.strip()

    # Remove trailing '.git' if present
    if input_str.endswith(".git"):
        input_str = input_str[:-4]

    # If input is a URL, parse it to extract the repo path.
    if input_str.startswith("http"):
        try:
            parsed = urlparse(input_str)
            path_parts = parsed.path.strip("/").split("/")
            if len(path_parts) >= 2:
                repo_id = "/".join(path_parts[:2])
                logging.info("Extracted repo_id: %s", repo_id)
                return repo_id
            else:
                raise ValueError("Invalid GitHub URL format.")
        except Exception as e:
            raise ValueError(f"Error parsing URL: {e}")
    else:
        # Assume it's already in the correct format.
        return input_str


def fetch_repo_details(repo_id: str) -> dict:
    """
    Fetch repository details from GitHub using its API.
    """
    api_url = f"https://api.github.com/repos/{repo_id}"
    response = requests.get(api_url)
    if response.status_code != 200:
        raise ValueError(
            f"Error: Unable to fetch repository details for '{repo_id}'. "
            f"Status code: {response.status_code}"
        )
    return response.json()


def fetch_readme(repo_id: str) -> str:
    """
    Attempt to fetch the README.md from the repository's main branch.
    """
    readme_url = f"https://raw.githubusercontent.com/{repo_id}/main/README.md"
    response = requests.get(readme_url)
    if response.status_code == 200:
        return response.text
    else:
        return "No README found."


def generate_docs(repo_input: str) -> str:
    """
    Fetch repository details from GitHub and generate documentation using the Groq API.
    """
    try:
        # Extract a clean repository ID.
        repo_id = extract_repo_id(repo_input)
    except ValueError as e:
        return str(e)

    try:
        repo_data = fetch_repo_details(repo_id)
    except ValueError as e:
        return str(e)

    # Extract repository information.
    repo_name = repo_data.get("name", repo_id)
    repo_description = repo_data.get("description", "No description provided.")
    readme_content = fetch_readme(repo_id)

    # Construct the prompt for the Groq API.
    prompt = (
        f"Generate comprehensive documentation for the GitHub repository '{repo_id}'. "
        f"The repository is named '{repo_name}' and its description is: {repo_description}. "
        f"Use the following README content as reference:\n\n{readme_content}\n\n"
        "Please include sections such as an introduction, installation instructions, usage examples, "
        "and any relevant details that would help a new user understand and work with this repository."
    )

    # Retrieve the Groq API key from environment variables.
    groq_api_key = os.environ.get("GROQ_API_KEY")
    if not groq_api_key:
        return "Error: GROQ_API_KEY environment variable is not set."

    # Initialize the Groq client.
    client = Groq(api_key=groq_api_key)

    try:
        # Call the Groq API with the generated prompt.
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="deepseek-r1-distill-llama-70b",
            stream=False,
        )
    except Exception as e:
        return f"Error calling Groq API: {e}"

    # Extract and return the generated documentation.
    documentation = chat_completion.choices[0].message.content
    return documentation