app.py

import os
import gc
import tempfile
import uuid
import pandas as pd
import openai

from gitingest import ingest
from llama_index.core import Settings
from llama_index.llms.sambanovasystems import SambaNovaCloud
from llama_index.core import PromptTemplate
from llama_index.embeddings.mixedbreadai import MixedbreadAIEmbedding
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import MarkdownNodeParser

import streamlit as st

# Fetch API keys securely from Hugging Face secrets
SAMBANOVA_API_KEY = os.getenv("SAMBANOVA_API_KEY")
MXBAI_API_KEY = os.getenv("MXBAI_API_KEY")

# Ensure both API keys are available
if not SAMBANOVA_API_KEY:
    raise ValueError("SAMBANOVA_API_KEY is not set in the Hugging Face secrets.")
if not MXBAI_API_KEY:
    raise ValueError("MXBAI_API_KEY is not set in the Hugging Face secrets.")

if "id" not in st.session_state:
    st.session_state.id = uuid.uuid4()
    st.session_state.file_cache = {}

session_id = st.session_state.id
client = None

@st.cache_resource
def load_llm():
    # Instantiate the SambaNova model
    llm = SambaNovaCloud(
        model="Meta-Llama-3.1-405B-Instruct",  # Use the correct model name
        context_window=100000,
        max_tokens=1024,
        temperature=0.7,
        top_k=1,
        top_p=0.01,
    )
    return llm

def reset_chat():
    st.session_state.messages = []
    st.session_state.context = None
    gc.collect()

def process_with_gitingets(github_url):
    # or from URL
    summary, tree, content = ingest(github_url)
    return summary, tree, content


with st.sidebar:
    st.header(f"Add your GitHub repository!")
    
    github_url = st.text_input("Enter GitHub repository URL", placeholder="GitHub URL")
    load_repo = st.button("Load Repository")

    if github_url and load_repo:
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                st.write("Processing your repository...")
                repo_name = github_url.split('/')[-1]
                file_key = f"{session_id}-{repo_name}"
                
                if file_key not in st.session_state.get('file_cache', {}):
                    if os.path.exists(temp_dir):
                        summary, tree, content = process_with_gitingets(github_url)

                        # Write summary to a markdown file
                        with open("content.md", "w", encoding="utf-8") as f:
                            f.write(content)

                        # Write summary to a markdown file in temp directory
                        content_path = os.path.join(temp_dir, f"{repo_name}_content.md")
                        with open(content_path, "w", encoding="utf-8") as f:
                            f.write(content)
                        loader = SimpleDirectoryReader(
                            input_dir=temp_dir,
                        )
                    else:    
                        st.error('Could not find the file you uploaded, please check again...')
                        st.stop()
                    
                    docs = loader.load_data()

                    # setup llm & embedding model
                    llm = load_llm()

                    # Mixedbread AI embedding setup
                    embed_model = MixedbreadAIEmbedding(
                        api_key=MXBAI_API_KEY,  # Use the API key from Hugging Face secret
                        model_name="mixedbread-ai/mxbai-embed-large-v1",  # Specify the model name
                    )

                    # Creating an index over loaded data
                    Settings.embed_model = embed_model
                    node_parser = MarkdownNodeParser()
                    index = VectorStoreIndex.from_documents(documents=docs, transformations=[node_parser], show_progress=True)

                    # Create the query engine
                    Settings.llm = llm
                    query_engine = index.as_query_engine(streaming=True)

                    # ====== Customise prompt template ======
                    qa_prompt_tmpl_str = (
                        "Context information is below.\n"
                        "---------------------\n"
                        "{context_str}\n"
                        "---------------------\n"
                        "Given the context information above I want you to think step by step to answer the query in a highly precise and crisp manner focused on the final answer, in case you don't know the answer say 'I don't know!'.\n"
                        "Query: {query_str}\n"
                        "Answer: "
                    )
                    qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)

                    query_engine.update_prompts(
                        {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
                    )
                    
                    st.session_state.file_cache[file_key] = query_engine
                else:
                    query_engine = st.session_state.file_cache[file_key]

                # Inform the user that the file is processed and display the PDF uploaded
                st.success("Ready to Chat!")
        except Exception as e:
            st.error(f"An error occurred: {e}")
            st.stop()     


col1, col2 = st.columns([6, 1])

with col1:
    st.header(f"Chat with GitHub using RAG </>")

with col2:
    st.button("Clear ↺", on_click=reset_chat)

# Initialize chat history
if "messages" not in st.session_state:
    reset_chat()


# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


# Accept user input
if prompt := st.chat_input("What's up?"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)

    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        
        try:
            # Get the repo name from the GitHub URL
            repo_name = github_url.split('/')[-1]
            file_key = f"{session_id}-{repo_name}"
            
            # Get query engine from session state
            query_engine = st.session_state.file_cache.get(file_key)
            
            if query_engine is None:
                st.error("Please load a repository first!")
                st.stop()
                
            # Use the query engine
            response = query_engine.query(prompt)
            
            # Handle streaming response
            if hasattr(response, 'response_gen'):
                for chunk in response.response_gen:
                    if isinstance(chunk, str):  # Only process string chunks
                        full_response += chunk
                        message_placeholder.markdown(full_response + "▌")
            else:
                # Handle non-streaming response
                full_response = str(response)
                message_placeholder.markdown(full_response)

            message_placeholder.markdown(full_response)
        except Exception as e:
            st.error(f"An error occurred while processing your query: {str(e)}")
            full_response = "Sorry, I encountered an error while processing your request."
            message_placeholder.markdown(full_response)

    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": full_response})