LangChain GitLoader example

rag.py

import getpass
import os
import bs4
from langchain import hub
# from langchain_community.document_loaders import WebBaseLoader
from langchain_chroma import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

from langchain_community.document_loaders import GitLoader

os.environ["OPENAI_API_KEY"] = getpass.getpass()

llm = ChatOpenAI(model="gpt-3.5-turbo-0125")

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

prompt = hub.pull("rlm/rag-prompt")

class Bot():
    def __init__(self):
        self.vectorstore = None
        self.rag_chain = None

    def load_git(self, repo_path, branch):
        loader = GitLoader(repo_path=repo_path, branch=branch)
        docs = loader.load()

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(docs)
        self.vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())

        retriever = self.vectorstore.as_retriever()

        self.rag_chain = (
            {"context": retriever | format_docs, "question": RunnablePassthrough()}
            | prompt
            | llm
            | StrOutputParser()
        )

    def ask(q):
        return self.rag_chain.invoke(q)

    def cleanup():
        self.vectorstore.delete_collection()

    def start():
        print("*** START OF CHAT ***")
        n = 1
        while True:
            try:
                inp = input(f"User[{n}]: ")
                ans = self.ask(inp)
                print()
                print(f"AI[{n}]: " + str(ans))
                print()
            except EOFError:
                print()
                print("*** END OF CHAT ***")
                return
            n += 1

bot = Bot()