File size: 7,110 Bytes
40fd220 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 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}) |