Create app.py

app.py

@@ -0,0 +1,181 @@
+import os
+import gc
+import tempfile
+import uuid
+import pandas as pd
+import openai  # Import openai for Sambanova API
+
+from gitingest import ingest
+from llama_index.core import Settings
+from llama_index.core import PromptTemplate
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
+from llama_index.core.node_parser import MarkdownNodeParser
+
+import streamlit as st
+
+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
+
+# Update the load_llm function to use Sambanova's API
+@st.cache_resource
+def load_llm():
+    # Initialize the Sambanova OpenAI client
+    client = openai.OpenAI(
+        api_key=os.environ.get("SAMBANOVA_API_KEY"),
+        base_url="https://api.sambanova.ai/v1",
+    )
+    return client
+
+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()  # Load the Sambanova LLM client
+                    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5", trust_remote_code=True)
+                    # 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, where we use a cohere reranker on the fetched nodes
+                    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, incase 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 to get the context for the query
+            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})