rag/generate.py

"""
Batch-generates DSA knowledge base .txt files using the Groq API.

Setup:
    pip install groq python-dotenv
    Create a .env file with: GROQ_API_KEY=your_key_here

Run:
    python generate_knowledge_base.py

Output:
    Files are written to ./knowledge_base/<topic_snake_case>.txt
"""

import os
import re
import time
from dotenv import load_dotenv
from groq import Groq

load_dotenv()

client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
MODEL = "llama-3.3-70b-versatile"   # swap for another Groq-hosted model if you prefer
OUTPUT_DIR = "knowledge_base"
DELAY_BETWEEN_CALLS = 1.5  # seconds, be nice to rate limits

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── PROMPT TEMPLATE ─────────────────────────────────────────────────────────
PROMPT_TEMPLATE = """You are writing a knowledge base entry for a DSA (Data Structures & Algorithms)
tutoring RAG chatbot. Write in your own original words — do NOT copy verbatim from
GeeksforGeeks, textbooks, or any other source. Paraphrase and synthesize instead.
Keep explanations beginner-friendly but technically accurate.

Generate the entry for this topic: {topic}

Output ONLY in the exact following plain text structure, no markdown formatting,
no extra commentary before or after:

TOPIC: {topic}

DEFINITION:
<2-3 sentence clear definition of what it is and what problem it solves>

TIME_COMPLEXITY:
<Best/Average/Worst case if they differ, otherwise a single Big-O with brief justification>

SPACE_COMPLEXITY:
<Big-O with brief note on what's using the space>

USE_WHEN:
<1-2 sentences: scenarios/conditions where this is the right tool>

AVOID_WHEN:
<1-2 sentences: scenarios where this is a poor choice, and what to use instead>

EXAMPLE:
<A small, concrete worked example (3-6 elements/nodes max), shown as a step-by-step
walkthrough with intermediate states, ending in a clear result marked with a checkmark.
Use simple ASCII/text formatting (arrows, brackets, indentation) — no images.>

REAL_WORLD_ANALOGY:
<1-2 sentence relatable comparison to something non-technical>

SOURCE_NOTE:
<Optional: e.g. "Concepts referenced from CLRS Introduction to Algorithms" or
"GeeksforGeeks - {topic} (paraphrased, no verbatim text used)". Leave blank if not applicable.>
"""

# ── TOPIC LIST ───────────────────────────────────────────────────────────────
TOPICS = [
    # Arrays & Searching
    "Linear Search", "Binary Search", "Two Pointer Technique", "Sliding Window Technique",
    # Sorting
    "Bubble Sort", "Selection Sort", "Insertion Sort", "Merge Sort", "Quick Sort",
    "Heap Sort", "Counting Sort","radix sort","bucket sort",
    # Linked Lists
    "Singly Linked List", "Doubly Linked List", "Circular Linked List",
    # Stacks & Queues
    "Stack", "Queue", "Circular Queue", "Deque", "Monotonic Stack",
    # Trees
    "Binary Tree", "Binary Search Tree", "AVL Tree", "Heap (Min/Max)", "Trie",
    "Segment Tree", "Fenwick Tree (Binary Indexed Tree)",
    # Graphs
    "BFS", "DFS", "Dijkstra's Algorithm", "Bellman-Ford Algorithm",
    "Floyd-Warshall Algorithm", "Kruskal's Algorithm (MST)", "Prim's Algorithm (MST)",
    "Topological Sort", "Union-Find (Disjoint Set)",
    # Dynamic Programming
    "Dynamic Programming Overview", "0/1 Knapsack", "Longest Common Subsequence",
    "Longest Increasing Subsequence", "Coin Change Problem", "Edit Distance",
    # Greedy
    "Greedy Algorithms Overview", "Activity Selection Problem", "Fractional Knapsack",
    # Hashing
    "Hash Table", "Hash Set vs Hash Map",
    # Recursion & Backtracking
    "Recursion Overview", "Backtracking Overview", "N-Queens Problem",
    "Subset Sum / Power Set Generation",
    # Bit Manipulation
    "Bit Manipulation Basics", "XOR Tricks",
]


def to_snake_case(topic: str) -> str:
    """Converts a topic name into a safe snake_case filename."""
    cleaned = re.sub(r"[^\w\s-]", "", topic)          # strip punctuation like ()/,'
    cleaned = re.sub(r"\s+", "_", cleaned.strip())     # spaces -> underscores
    return cleaned.lower()


def generate_entry(topic: str) -> str:
    prompt = PROMPT_TEMPLATE.format(topic=topic)
    response = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.4,
        max_tokens=1024,
    )
    return response.choices[0].message.content.strip()


def main():
    print(f"Generating {len(TOPICS)} knowledge base files into ./{OUTPUT_DIR}/\n")

    for i, topic in enumerate(TOPICS, start=1):
        filename = to_snake_case(topic) + ".txt"
        filepath = os.path.join(OUTPUT_DIR, filename)

        if os.path.exists(filepath):
            print(f"[{i}/{len(TOPICS)}] SKIP (already exists): {filename}")
            continue

        print(f"[{i}/{len(TOPICS)}] Generating: {topic} -> {filename}")
        try:
            content = generate_entry(topic)
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(content)
        except Exception as e:
            print(f"    ERROR generating '{topic}': {e}")

        time.sleep(DELAY_BETWEEN_CALLS)

    print("\nDone. Review files in ./knowledge_base/ before indexing into ChromaDB.")


if __name__ == "__main__":
    main()