app.py

import gradio as gr
import os
from queue import SimpleQueue
from langchain.callbacks.manager import CallbackManager
from langchain.chat_models import ChatOpenAI
from pydantic import BaseModel
import requests
import typing
from typing import TypeVar, Generic
import tqdm
from langchain.chains import ConversationalRetrievalChain
import os
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import DeepLake
import random
import time
import os
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
import math
import subprocess

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import LLMResult
from typing import Any, Union

job_done = object()

class StreamingGradioCallbackHandler(BaseCallbackHandler):
    def __init__(self, q: SimpleQueue):
        self.q = q

    def on_llm_start(
        self, serialized: typing.Dict[str, Any], prompts: typing.List[str], **kwargs: Any
    ) -> None:
        """Run when LLM starts running. Clean the queue."""
        while not self.q.empty():
            try:
                self.q.get(block=False)
            except Empty:
                continue

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        self.q.put(token)

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Run when LLM ends running."""
        self.q.put(job_done)

    def on_llm_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Run when LLM errors."""
        self.q.put(job_done)


class Response(BaseModel):
    result: typing.Any
    error: str
    stdout: str
    repo: str

class HumanPrompt(BaseModel):
    prompt: str

class GithubResponse(BaseModel):
    result: typing.Any
    error: str
    stdout: str
    repo: str


repo_name = gr.State()
git_tickets = gr.State()
git_titles = gr.State()
git_ticket_choices = gr.State()
vector_db_url = gr.State()
git_tickets.value = []
git_titles.value = []
git_ticket_choices.value = []

embeddings = OpenAIEmbeddings(disallowed_special=())

def git_clone(repo_url):
    subprocess.run(["git", "clone", repo_url])
    dirpath = repo_url.split('/')[-1]
    if dirpath.lower().endswith('.git'):
        dirpath = dirpath[:-4]
    return dirpath


def index_repo(textbox: str, dropdown: str) -> Response:

    mapping = {
        "Langchain" : "https://github.com/langchain-ai/langchain.git",
        "Weaviate": "https://github.com/weaviate/weaviate.git",
        "Llama2": "https://github.com/facebookresearch/llama.git",
        "OpenAssistant": "https://github.com/LAION-AI/Open-Assistant.git",
        "MemeAI": "https://github.com/aiswaryasankar/memeAI.git",
        "GenerativeAgents": "https://github.com/joonspk-research/generative_agents.git"
    }

    if textbox != "":
        repo = textbox
    else:
        repo = mapping[dropdown[0]]

    repo_name.value = repo
    pathName = git_clone(repo)
    root_dir = './' + pathName

    activeloop_username = "aiswaryas"
    dataset_path = f"hub://{activeloop_username}/" + pathName
    invalid_dataset_path = False

    try:
        try:
            db = DeepLake(dataset_path=dataset_path,
                    embedding_function=embeddings,
                    token=os.environ['ACTIVELOOP_TOKEN'],
                    read_only=True,
                    num_workers=12,
                    runtime = {"tensor_db": True}
                )
        except Exception as e:
            print("Failed to read: " + str(e))
            if "scheduled for deletion" in str(e):
                dataset_path = f"hub://{activeloop_username}/" + pathName + str(random.randint(1,100))
                invalid_dataset_path = True

        print(invalid_dataset_path)
        print(db)
        print(len(db.vectorstore.dataset))
        if invalid_dataset_path or db is None or len(db.vectorstore.dataset) == 0:
            print("Dataset doesn't exist, fetching data")
            try:
                docs = []
                for dirpath, dirnames, filenames in os.walk(root_dir):
                    for file in filenames:
                        print(file)
                        try:
                            loader = TextLoader(os.path.join(dirpath, file), encoding='utf-8')
                            docs.extend(loader.load_and_split())
                        except Exception as e:
                            print("Exception: " + str(e) + "| File: " + os.path.join(dirpath, file))
                            pass

                activeloop_username = "aiswaryas"
                text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
                texts = text_splitter.split_documents(docs)

                db = DeepLake(dataset_path=dataset_path,
                        embedding_function=embeddings,
                        token=os.environ['ACTIVELOOP_TOKEN'],
                        read_only=False,
                        num_workers=12,
                        runtime = {"tensor_db": True}
                )
                # Do this in chunks to avoid hitting the ratelimit immediately
                for i in range(0, len(texts), 500):
                    print("Adding documents " + str(i))
                    db.add_documents(texts[i:i+500])
                    time.sleep(.5)

            except Exception as e:
                return Response(
                    result= "Failed to index github repo",
                    repo="",
                    error=str(e),
                    stdout="",
                )

    except Exception as e:
        return Response(
            result= "Failed to index github repo",
            repo="",
            error=str(e),
            stdout="",
        )

    vector_db_url.value = dataset_path

    return {
        success_response: "SUCCESS",
        launch_product: gr.update(visible=True)
    }


def answer_questions(question: str, github: str, **kwargs) -> Response:

    repoName = repo_name.value
    github = repoName[:-4]
    print("REPO NAME: " + github)

    try:
        embeddings = OpenAIEmbeddings(disallowed_special=())
        pathName = github.split('/')[-1]
        dataset_path = vector_db_url.value

        print("before reading repo")
        db = DeepLake(dataset_path=dataset_path, read_only=True, embedding_function=embeddings)

        print("finished indexing repo")
        retriever = db.as_retriever()
        retriever.search_kwargs['distance_metric'] = 'cos'
        retriever.search_kwargs['fetch_k'] = 100
        retriever.search_kwargs['maximal_marginal_relevance'] = True
        retriever.search_kwargs['k'] = 20

        q = SimpleQueue()
        model = ChatOpenAI(
            model_name='gpt-3.5-turbo-16k',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever, max_tokens_limit=16000)
        chat_history = []

    except Exception as e:
        print("Exception: " + str(e))
        return Response(
            result="",
            repo="",
            error=str(e),
            stdout="",
        )

    return Response(
        result=qa({"question": question, "chat_history": chat_history}),
        repo="",
        error="",
        stdout="",
    )


def fetchGithubIssues(**kwargs) -> Response:
    """
        This endpoint should get a list of all the github issues that are open for this repository
    """
    repo = "/".join(repo_name.value[:-4].split("/")[-2:])
    print("REPO NAME IN FETCH GITHUB ISSUES: " + str(repo))

    batch = []
    all_issues = []
    per_page = 100  # Number of issues to return per page
    num_pages = math.ceil(20 / per_page)
    base_url = "https://api.github.com/repos"

    GITHUB_TOKEN = "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"  # Copy your GitHub token here
    headers = {"Authorization": f"token {GITHUB_TOKEN}"}

    issues_data = []

    for page in range(num_pages):
        # Query with state=all to get both open and closed issues
        query = f"issues?page={page}&per_page={per_page}&state=all"
        issues = requests.get(f"{base_url}/{repo}/{query}", headers=headers)
        print(f"{base_url}/{repo}/{query}")

        batch.extend(issues.json())
        for issue in issues.json():
            issues_data.append({
                "issue_url": issue["url"],
                "title": issue["title"],
                "body": issue["body"],
                "comments_url": issue["comments_url"],
            })

    # This should set the state variables for tickets
    git_tickets.value = issues_data
    git_ticket_choices.value = {ticket["title"]: ticket for ticket in issues_data}
    git_titles.value = [ticket["title"] for ticket in issues_data]
    return  issues_data


def generateFolderNamesForRepo(repo):
    """
        This endpoint will first take the repo structure and return the folder and subfolder names.
        From those names, it will then prompt the model to generate an architecture diagram of that folder.
        There will be three "modules" no input just output that take the autogenerated prompts based on the
        input data and generate the responses that are displayed in the UI.
    """
    pathName = git_clone(repo)
    root_dir = './' + pathName

    files, dirs, docs = [], [], []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for file in filenames:
            try:
                loader = TextLoader(os.path.join(dirpath, file), encoding='utf-8')
                docs.extend(loader.load_and_split())
                files.append(file)
                dirs.append(dirnames)
            except Exception as e:
                print("Exception: " + str(e) + "| File: " + os.path.join(dirpath, file))
                pass

    return dirs


def generateDocumentationPerFolder(dir, github):

    if dir == "overview":
        prompt= """
            Summarize the structure of the {} repository.  Make a list of all endpoints and their behavior.  Explain
            how this module is used in the scope of the larger project.  Format the response as code documentation with an
            Overview, Architecture and Implementation Details.  Within implementation details, list out each function and provide
            an overview of that function.
        """.format(github)
    else:
        prompt= """
            Summarize how {} is implemented in the {} repository.  Make a list of all functions and their behavior.  Explain
            how this module is used in the scope of the larger project.  Format the response as code documentation with an
            Overview, Architecture and Implementation Details.  Within implementation details, list out each function and provide
            an overview of that function.
        """.format(dir, github)

    try:
        embeddings = OpenAIEmbeddings(disallowed_special=())
        pathName = github.split('/')[-1]
        dataset_path = vector_db_url.value

        db = DeepLake(dataset_path=dataset_path, read_only=True, embedding_function=embeddings)

        retriever = db.as_retriever()
        retriever.search_kwargs['distance_metric'] = 'cos'
        retriever.search_kwargs['fetch_k'] = 100
        retriever.search_kwargs['maximal_marginal_relevance'] = True
        retriever.search_kwargs['k'] = 20

        # streaming_handler = kwargs.get('streaming_handler')
        model = ChatOpenAI(
            model_name='gpt-3.5-turbo-16k',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever, max_tokens_limit=16000)
        chat_history = []
        return qa({"question": prompt, "chat_history": chat_history})["answer"]

    except Exception as e:
        print (str(e))
        return "Failed to generate documentation"


def solveGithubIssue(ticket, history) -> Response:
    """
        This endpoint takes in a github issue and then queries the db for the question against the codebase.
    """
    repoName = repo_name.value
    github = repoName[:-4]

    repoFolder = github.split("/")[-1]
    body = git_ticket_choices.value[ticket]["body"]
    title = git_ticket_choices.value[ticket]["title"]
    question = """
        Given the code in the {} repo, propose a solution for this ticket {} that includes a
        high level implementation, narrowing down the root cause of the issue and psuedocode if
        applicable on how to resolve the issue. If multiple changes are required to address the
        problem, list out each of the steps and a brief explanation for each one.
        """.format(repoFolder, body)

    q_display = """
            Can you explain how to approach solving this ticket: {}.  Here is a summary of the issue: {}
        """.format(title, body)

    try:
        embeddings = OpenAIEmbeddings(disallowed_special=())
        pathName = github.split('/')[-1]
        dataset_path = vector_db_url.value

        db = DeepLake(dataset_path=dataset_path, read_only=True, embedding=embeddings)

        retriever = db.as_retriever()
        retriever.search_kwargs['distance_metric'] = 'cos'
        retriever.search_kwargs['fetch_k'] = 100
        retriever.search_kwargs['maximal_marginal_relevance'] = True
        retriever.search_kwargs['k'] = 20

        q = SimpleQueue()
        model = ChatOpenAI(
            model_name='gpt-3.5-turbo-16k',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever, max_tokens_limit=16000)

    except Exception as e:
        return [[str(e), None]]

    history = [[q_display, ""]]
    history[-1][1] = ""
    # Flatten the list of lists
    flat_list = [item for sublist in history for item in sublist]
    flat_list = [item for item in flat_list if item is not None]

    print(flat_list)
    for char in qa({"question": question, "chat_history": []})["answer"]:
        history[-1][1] += char
        yield history


def user(message, history):
    return "", history + [[message, None]]


def bot(history, **kwargs):

    user_message = history[-1][0]

    # global repoName
    repoName = repo_name.value
    print("STATE REPO NAME: " + repoName)
    github = repoName[:-4]
    try:
        embeddings = OpenAIEmbeddings(disallowed_special=())
        pathName = github.split('/')[-1]
        dataset_path = vector_db_url.value

        db = DeepLake(dataset_path=dataset_path, read_only=True, embedding_function=embeddings)
        retriever = db.as_retriever()
        retriever.search_kwargs['distance_metric'] = 'cos'
        retriever.search_kwargs['fetch_k'] = 100
        retriever.search_kwargs['maximal_marginal_relevance'] = True
        retriever.search_kwargs['k'] = 20

        q = SimpleQueue()
        model = ChatOpenAI(
            model_name='gpt-3.5-turbo-16k',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever, max_tokens_limit=16000, return_source_documents=True, get_chat_history=lambda h : h)
        chat_history = []

    except Exception as e:
        print("Exception: " + str(e))
        return str(e)

    history[-1][1] = ""

    for char in qa({"question": user_message, "chat_history": []})["answer"]:
        history[-1][1] += char
        yield history


with gr.Blocks() as demo:

    # repoName = gr.State(value="https://github.com/sourcegraph/cody.git")

    gr.Markdown("""
        <h1 align="center"> Entelligence AI </h1>
        <p style="text-align: center; font-size:36">Enabling your product team to ship product 10x faster.</p>
        """)

    repoTextBox = gr.Textbox(label="Github Repository")

    gr.Markdown("""Choose from any of the following repositories""")
    ingestedRepos = gr.CheckboxGroup(choices=['Langchain', 'Weaviate', 'OpenAssistant', 'GenerativeAgents','Llama2', "MemeAI"], label="Github Repository", value="Langchain")

    success_response = gr.Textbox(label="")
    ingest_btn = gr.Button("Index repo")

    with gr.Column(visible=False) as launch_product:

        # Toggle visibility of the chat, bugs, docs, model windows
        with gr.Tab("Code Chat"):
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
            clear = gr.Button("Clear")

            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)


        index = 0
        with gr.Tab("Bug Triage"):

            # Display the titles in the dropdown
            def create_ticket_dropdown():
                print(git_titles.value)
                return ticketDropdown.update(
                    choices=git_titles.value
                )

            ticketDropdown = gr.Dropdown(choices=[], title="Github Issues", interactive=True)
            ticketDropdown.focus(create_ticket_dropdown, outputs=ticketDropdown)
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
            clear = gr.Button("Clear")
            ticketDropdown.change(solveGithubIssue, inputs=[ticketDropdown, chatbot], outputs=[chatbot])
            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)


        # with gr.Tab("AI Code Documentation"):

        #     repoName = repo_name.value
        #     # First parse through the folder structure and store that as a list of clickable buttons
        #     gr.Markdown("""
        #         ## AI Generated Code Documentation
        #         Code documentation comes in 3 flavors - internal engineering, external API documentation and product documentation.  Each offers different layers of abstraction over the code base.
        #     """)

        #     #   docs = generateDocumentationPerFolder("overview", repo_name)

        #     # For now let's just display all of the docs in one big file
        #     allDocs = ""
        #     dirNames = generateFolderNamesForRepo(repoName[:-4])
        #     for dir in dirNames:
        #         if dir[0] != ".":
        #             allDocs += generateDocumentationPerFolder(dir, repoName[:-4]) + '\n\n'

        #     gr.Markdown(allDocs)

        #     def button_click_callback(markdown):
        #         docs = generateDocumentationPerFolder("overview", repoName[:-4])
        #         markdown.update(docs)

        #     markdown = gr.Markdown()
        #     # Generate the left column buttons and their names and wrap each one in a function
        #     with gr.Row():
        #         with gr.Column(scale=.5, min_width=300):
        #             dirNames = generateFolderNamesForRepo(repoName[:-4])
        #             buttons = [gr.Button(folder_name) for folder_name in dirNames]
        #             for btn, folder_name in zip(buttons, dirNames):
        #                 btn.click(button_click_callback, [markdown], [markdown] )


        #         # Generate the overall documentation for the main bubble at the same time
        #         with gr.Column(scale=2, min_width=300):
        #             docs = generateDocumentationPerFolder("overview", repoName[:-4])
        #             markdown.update(docs)
        #             # markdown.render()


        with gr.Tab("Custom Model Finetuning"):
            # First provide a summary of offering
            gr.Markdown("""
                # Enterprise Custom Model Finetuning
                Finetuning code generation models directly on your enterprise code base has shown up to 10% increase in model suggestion acceptance rate.
                """)

            # Choose base model - radio with model size
            gr.Radio(choices=["Santacoder (1.1B parameter model)", "Incoder (6B parameter model)", "Codegen (16B parameter model)", "Starcoder (15.5B parameter model)"] , value="Starcoder (15.5B parameter model)")

            # Choose existing code base or input a new code base for finetuning -
            with gr.Row():
                gr.Markdown("""
                    If you'd like to use the current code base, click this toggle otherwise input the entire code base below.
                """)
                existing_repo = gr.Checkbox(value=True, label="Use existing repository")
            gr.Textbox(label="Input repository", visible=False)

            # Allow option to remove generated files etc etc
            gr.Markdown("""
                Finetuned model performance is highly dependent on training data quality. We have currently found that excluding the following file types improves performance. If you'd like to include them, please toggle them.
            """)
            file_types = gr.CheckboxGroup(choices=['.bin', '.gen', '.git', '.gz','.jpg', '.lz', '.midi', '.mpq','.png', '.tz'], label="Removed file types")

            # Based on data above, we should show a field for estimated fine tuning cost
            # Then we should show the chart for loss
            def wandb_report(url):
                iframe = f'<iframe src={url} style="border:none;height:1024px;width:100%">'
                return gr.HTML(iframe)

            submit_btn = gr.Button("Start Training")
            with gr.Column(visible=False) as start_training:
                # Include the epoch loss table
                epoch_loss = gr.Dataframe(
                    headers=["Step", "Training Loss", "Validation Loss"],
                    datatype=["number", "number", "number"],
                    row_count=5,
                    col_count=(3, "fixed"),
                    value=[[500, 1.868200, 1.548535], [1000, 1.450100, 1.518277], [1500, 1.659000, 1.486497],
                    [2000, 1.364900, 1.452842], [2500, 1.406300, 1.405151], [3000, 1.276000, 1.346159]]
                )

                # After you start training you should see the Wandb report
                report_url = 'https://wandb.ai/aiswaryasankar/aiswarya-santacoder-finetuning/reports/Aiswarya-Santacoder-Finetuning--Vmlldzo0ODM3MDA4'
                report = wandb_report(report_url)

                # Include a playground to compare different models on given tasks
                # Link to the generated huggingface spaces model if you opt into it
                # Toggle to select model for the remaining functionality

            def startTraining(existing_repo, file_types):
                return {
                    start_training: gr.update(visible=True),
                }

            submit_btn.click(
                startTraining,
                inputs=[existing_repo, file_types],
                outputs=[start_training], # report, epoch_loss,
            )


    ingest_btn.click(fn=index_repo, inputs=[repoTextBox, ingestedRepos], outputs=[success_response, launch_product], api_name="index_repo").then(fn=fetchGithubIssues)



demo.queue()
demo.launch(debug=True)