LangChain: Code Explainer

# Install: pip install langchain langchain-openai langchain-community chromadb

from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

# --- Config ---
CODE_DIR = "./src"
CHROMA_DIR = "./chroma_code"

# --- Load source code ---
loader = DirectoryLoader(
    CODE_DIR,
    glob="**/*.{py,ts,js,go,rs}",
    loader_cls=TextLoader,
    loader_kwargs={"autodetect_encoding": True},
)
docs = loader.load()

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1500,
    chunk_overlap=200,
    separators=["\nclass ", "\ndef ", "\nfunction ", "\n\n", "\n"],
)
chunks = splitter.split_documents(docs)

# --- Build vector store ---
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_DIR)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# --- Code-aware prompt ---
prompt = PromptTemplate(
    template="""You are a senior developer explaining code to a teammate.
Use the code context below to answer the question.
Explain the logic, point out important patterns, and mention any gotchas.
If the answer is not in the context, say so.

Code context:
{context}

Question: {question}
Explanation:""",
    input_variables=["context", "question"],
)

llm = ChatOpenAI(model="gpt-4o", temperature=0)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type_kwargs={"prompt": prompt},
    return_source_documents=True,
)

if __name__ == "__main__":
    print("Code Explainer ready. Ask about any part of the codebase.")
    while True:
        query = input("\nQuestion: ")
        if query.lower() == "quit":
            break
        result = qa_chain.invoke({"query": query})
        print(f"\n{result['result']}")
        files = set(d.metadata["source"] for d in result["source_documents"])
        print(f"\nRelevant files: {', '.join(files)}")