wehub-resource-sync
160 lines
5.3 KiB
Python
160 lines
5.3 KiB
Python
import streamlit as st
|
|
import os
|
|
from llama_index.core import Settings, VectorStoreIndex, PromptTemplate
|
|
from llama_index.embeddings.nebius import NebiusEmbedding
|
|
from llama_index.llms.nebius import NebiusLLM
|
|
from llama_index.readers.github import GithubRepositoryReader, GithubClient
|
|
import re
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
def parse_github_url(url):
|
|
pattern = r"https?://github\.com/([^/]+)/([^/]+)(?:/tree/([^/]+))?"
|
|
match = re.match(pattern, url)
|
|
if not match:
|
|
raise ValueError("Invalid GitHub repository URL")
|
|
owner, repo, branch = match.groups()
|
|
return owner, repo, branch if branch else "main"
|
|
|
|
@st.cache_resource
|
|
def load_github_data(github_token, owner, repo, branch="main"):
|
|
github_client = GithubClient(github_token)
|
|
loader = GithubRepositoryReader(
|
|
github_client,
|
|
owner=owner,
|
|
repo=repo,
|
|
filter_file_extensions=(
|
|
[".py", ".ipynb", ".js", ".ts", ".md"],
|
|
GithubRepositoryReader.FilterType.INCLUDE
|
|
),
|
|
verbose=False,
|
|
concurrent_requests=5,
|
|
)
|
|
return loader.load_data(branch=branch)
|
|
|
|
def run_rag_completion(query_text: str, docs) -> str:
|
|
llm = NebiusLLM(
|
|
model="deepseek-ai/DeepSeek-V3",
|
|
api_key=os.getenv("NEBIUS_API_KEY")
|
|
)
|
|
|
|
embed_model = NebiusEmbedding(
|
|
model_name="BAAI/bge-en-icl",
|
|
api_key=os.getenv("NEBIUS_API_KEY")
|
|
)
|
|
|
|
Settings.llm = llm
|
|
Settings.embed_model = embed_model
|
|
|
|
index = VectorStoreIndex.from_documents(docs)
|
|
query_engine = index.as_query_engine(similarity_top_k=5, streaming=True)
|
|
|
|
qa_prompt_tmpl = PromptTemplate(
|
|
"Context information is below.\n"
|
|
"---------------------\n"
|
|
"{context_str}\n"
|
|
"---------------------\n"
|
|
"Given the context information, please answer the query.\n"
|
|
"Query: {query_str}\n"
|
|
"Answer: "
|
|
)
|
|
|
|
query_engine.update_prompts({"response_synthesizer:text_qa_template": qa_prompt_tmpl})
|
|
response = query_engine.query(query_text)
|
|
return str(response)
|
|
|
|
def main():
|
|
st.set_page_config(page_title="Chat with Code", layout="wide")
|
|
|
|
@st.fragment
|
|
def download_response(response:str) :
|
|
st.download_button(
|
|
label="Download message",
|
|
type="secondary",
|
|
data=response,
|
|
file_name="chatbot_response.md",
|
|
mime="text/plain",
|
|
icon=":material/download:",
|
|
)
|
|
|
|
# Initialize session states
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = []
|
|
if "docs" not in st.session_state:
|
|
st.session_state.docs = None
|
|
|
|
# Header with title and buttons
|
|
col1, col2, col5, col3, col4 = st.columns([3, 1, 1, 1, 1])
|
|
with col1:
|
|
st.title("🤖 Chat with Code ")
|
|
with col3:
|
|
st.link_button("⭐ Star Repo", "https://github.com/Arindam200/nebius-cookbook")
|
|
with col4:
|
|
if st.button("🗑️ Clear Chat"):
|
|
st.session_state.messages = []
|
|
st.rerun()
|
|
|
|
st.caption("Powered by Nebius AI (DeepSeek-V3) and LlamaIndex")
|
|
|
|
# Sidebar
|
|
with st.sidebar:
|
|
# st.title("Select Model")
|
|
# model = st.selectbox(
|
|
# "",
|
|
# ["DeepSeek-V3"],
|
|
# index=0
|
|
# )
|
|
# st.divider()
|
|
st.subheader("GitHub Repository URL")
|
|
repo_url = st.text_input("", placeholder="Enter repository URL")
|
|
|
|
if st.button("Load Repository"):
|
|
if repo_url:
|
|
try:
|
|
github_token = os.getenv("GITHUB_TOKEN")
|
|
nebius_api_key = os.getenv("NEBIUS_API_KEY")
|
|
|
|
if not github_token or not nebius_api_key:
|
|
st.error("Missing API keys")
|
|
st.stop()
|
|
|
|
owner, repo, branch = parse_github_url(repo_url)
|
|
with st.spinner("Loading repository..."):
|
|
st.session_state.docs = load_github_data(github_token, owner, repo, branch)
|
|
st.success("✓ Repository loaded successfully")
|
|
except Exception as e:
|
|
st.error(f"Error: {str(e)}")
|
|
|
|
# Display chat messages
|
|
for message in st.session_state.messages:
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
# Chat input
|
|
if prompt := st.chat_input("Ask about the repository..."):
|
|
if not st.session_state.docs:
|
|
st.error("Please load a repository first")
|
|
st.stop()
|
|
|
|
# Add user message
|
|
st.session_state.messages.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.markdown(prompt)
|
|
|
|
# Generate response
|
|
with st.chat_message("assistant"):
|
|
with st.spinner("Thinking..."):
|
|
try:
|
|
response = run_rag_completion(prompt, st.session_state.docs)
|
|
st.markdown(response)
|
|
st.session_state.messages.append({"role": "assistant", "content": response})
|
|
download_response(response)
|
|
except Exception as e:
|
|
st.error(f"Error: {str(e)}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|