app.py

import gradio as gr
import torch

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 math
import tqdm
from langchain.chains import ConversationalRetrievalChain
import os
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import DeepLake
import random

os.environ['OPENAI_API_KEY']='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'
os.environ['ACTIVELOOP_TOKEN']='eyJhbGciOiJIUzUxMiIsImlhdCI6MTY4MTU5NTgyOCwiZXhwIjoxNzEzMjE4MTU5fQ.eyJpZCI6ImFpc3dhcnlhcyJ9.eoiMFZsS20zzMXXupFbowUlLdgIgf_MA1ck_DByzREeoQvNm8GPhKEfqea2y1Qak-ud2jo9dhSTBTfRe1ztezw'


import os
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter

import subprocess
repo_name = "https://github.com/aiswaryasankar/memeAI.git"

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

global ticket_choices

job_done = object() # signals the processing is done

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


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(repo: str) -> Response:
    pathName = git_clone(repo)
    root_dir = './' + pathName

    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())
            except Exception as e:
                print("Exception: " + str(e) + "| File: " + os.path.join(dirpath, file))
                pass

    activeloop_username = "aiswaryas"
    dataset_path = f"hub://{activeloop_username}/" + pathName
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    texts = text_splitter.split_documents(docs)

    print(texts)
    for text in texts:
        print(text)

    try:
        db = DeepLake(dataset_path=dataset_path,
                        embedding_function=embeddings,
                        token=os.environ['ACTIVELOOP_TOKEN'], read_only=False)
        # NOTE: read_only=False because we want to ingest documents
        # NOTE: This will raise a `deeplake.util.exceptions.LockedException` if dataset is already locked
        # NOTE: change it to read_only=True when querying the dataset

        # Delete dataset if not empty:
        if len(db.ds) > 0:
            print("Dataset not empty. Deleting existing dataset...")
            db.ds.delete()
            print("Done.")
            # Reinitialize
            db = DeepLake(dataset_path=dataset_path,
                            embedding_function=embeddings,
                            token=os.environ['ACTIVELOOP_TOKEN'], read_only=False)

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

    try:
        db.add_documents(texts)

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

    finally:
        db.ds._unlock()

    return "SUCCESS"


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

    global repo_name
    github = repo_name[:-4]
    try:
        embeddings = OpenAIEmbeddings(openai_api_key="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")
        pathName = github.split('/')[-1]
        dataset_path = "hub://aiswaryas/" + pathName

        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-4',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
            openai_api_key="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",
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever)
        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(repo: str, num_issues:int, **kwargs) -> Response:
    """
        This endpoint should get a list of all the github issues that are open for this repository
    """

    batch = []
    all_issues = []
    per_page = 100  # Number of issues to return per page
    num_pages = math.ceil(num_issues / 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)

        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"],
          })

    print(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[0]

def generateDocumentationPerFolder(dir, github):

  if dir == "overview":
    prompt= """
        Summarize the structure of the memeAI 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(dir)
  else:
    prompt= """
        Summarize how {} is implemented in the memeAI 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)

  print(prompt)
  try:
    embeddings = OpenAIEmbeddings(openai_api_key="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")
    pathName = github.split('/')[-1]
    dataset_path = "hub://aiswaryas/" + pathName

    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

    # streaming_handler = kwargs.get('streaming_handler')
    model = ChatOpenAI(
        model_name='gpt-4',
        temperature=0.0,
        verbose=True,
        streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
        openai_api_key="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",
    )
    qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever)
    chat_history = []

  except Exception as e:
    return str(e)

#   history[-1][1] = ""
#   for char in qa({"question": prompt, "chat_history": chat_history}):
#     history[-1][1] += char
#     time.sleep(0.01)
#     yield history

  return qa({"question": prompt, "chat_history": chat_history})["answer"]
  return response["answer"]


def generateArchitectureDiagram(folder) -> Response:
    """
        This endpoint should generate a Mermaid diagram for the given input files.  It will return the
    """


def solveGithubIssue(ticket, history) -> Response:
    """
        This endpoint takes in a github issue and then queries the db for the question against the codebase.
    """
    print(history)
    global repo_name, ticket_choices
    github = repo_name[:-4]
    repoFolder = github.split("/")[-1]
    body = ticket_choices[ticket]["body"]
    title = ticket_choices[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 = """
        How would I approach solving this ticket: {}.  Here is a summary of the issue: {}
    """.format(title, body)

    print(question)

    try:
        embeddings = OpenAIEmbeddings(openai_api_key="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")
        pathName = github.split('/')[-1]
        dataset_path = "hub://aiswaryas/" + pathName

        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-4',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
            openai_api_key="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",
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever)

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

    history = [[q_display, ""]]
    history[-1][1] = ""
    for char in qa({"question": prompt, "chat_history": chat_history}):
        history[-1][1] += char
        time.sleep(0.01)
        yield history

    # return [[qa({"question": question, "chat_history": chat_history})["answer"], None]]


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


def bot(history, **kwargs):
    print(history)
    user_message = history[-1][0]
    global repo_name
    github = repo_name[:-4]
    try:
        embeddings = OpenAIEmbeddings(openai_api_key="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")
        pathName = github.split('/')[-1]
        dataset_path = "hub://aiswaryas/" + pathName

        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-4',
            temperature=0.0,
            verbose=True,
            streaming=True,  # Pass `streaming=True` to make sure the client receives the data.
            callback_manager=CallbackManager(
                [StreamingGradioCallbackHandler(q)]
            ),
            openai_api_key="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",
        )
        qa = ConversationalRetrievalChain.from_llm(model,retriever=retriever)
        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": chat_history})["answer"]:
      history[-1][1] += char
      yield history


with gr.Blocks() as demo:

    gr.Markdown("""
    # Entelligence AI

    Enabling your product team to ship product 10x faster.
    """)

    repoTextBox = gr.Textbox(label="Github Repository")
    repo_name = "https://github.com/aiswaryasankar/memeAI.git"
    # def update_state(value):
    #   repo_name.value = value
    #   return value

    # repoTextBox.change(update_state, repoTextBox)
    # print(repo_name.value)
    success_response = gr.Textbox(label="")
    ingest_btn = gr.Button("Index repo")
    ingest_btn.click(fn=index_repo, inputs=repoTextBox, outputs=success_response, api_name="index_repo")


    # 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(tickets):

        return gr.Dropdown.update(
            choices=titles, value=titles[0]
        ), gr.update(visible=True)

      # Here you want to first call the getGithubIssues function
      # repo = gr.Interface.get_session_state("repo")
      print(repo_name)
      repo = "/".join(repo_name[:-4].split("/")[-2:])
      tickets = fetchGithubIssues(repo, 10)

      # Create the dropdown
      global ticket_choices
      ticket_choices = {ticket["title"]: ticket for ticket in tickets}
      ticket_titles = [ticket["title"] for ticket in tickets]

      ticketDropdown = gr.Dropdown(choices=ticket_titles, title="Github Issues")

      # Extract the ticket title, body for the selected ticket
      chatbot = gr.Chatbot()
      msg = gr.Textbox()
      clear = gr.Button("Clear")

      if index == 0:
        msg.submit(solveGithubIssue, [ticketDropdown, chatbot], [msg, chatbot], queue=False).then(
            bot, chatbot, chatbot
        )
        ticketDropdown.change(solveGithubIssue, inputs=[ticketDropdown, chatbot], outputs=[chatbot])
        index += 1
      else:
        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
            bot, chatbot, chatbot
        )
        index += 1
      clear.click(lambda: None, None, chatbot, queue=False)


    with gr.Tab("AI Code Documentation"):
      # 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)
      #   markdown = gr.Markdown(value=docs)

      def button_click_callback(label):
        docs = generateDocumentationPerFolder(label, repo_name[:-4])
        markdown.update(docs)

      # 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(repo_name[:-4])
            print(dirNames)
            buttons = [gr.Button(folder_name, onclick=button_click_callback) for folder_name in dirNames]

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


      # For each folder, generate a diagram and 2-3 prompts that dive deeper into explaining content


      # Render all the content in the UI

      #

    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
      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
        start_training= gr.update(visible=True)
        # return {
        #     report: report,
        #     epoch_loss: epoch_loss,
        #     start_training: gr.update(visible=True),
        # }

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

demo.launch(debug=True, share=True)