wehub-resource-sync
243 lines
8.3 KiB
Python
243 lines
8.3 KiB
Python
import os
|
|
|
|
import gc
|
|
import re
|
|
import glob
|
|
import uuid
|
|
import textwrap
|
|
import subprocess
|
|
import nest_asyncio
|
|
from dotenv import load_dotenv
|
|
|
|
import streamlit as st
|
|
|
|
from llama_index.core import Settings
|
|
from llama_index.core import PromptTemplate
|
|
from llama_index.core import SimpleDirectoryReader
|
|
from llama_index.core import VectorStoreIndex
|
|
from llama_index.core.storage.storage_context import StorageContext
|
|
from llama_index.core.node_parser import CodeSplitter, MarkdownNodeParser
|
|
from llama_index.llms.openai import OpenAI
|
|
from llama_index.llms.anthropic import Anthropic
|
|
from llama_index.core.indices.vector_store.base import VectorStoreIndex
|
|
from llama_index.vector_stores.qdrant import QdrantVectorStore
|
|
from llama_index.embeddings.fastembed import FastEmbedEmbedding
|
|
import qdrant_client
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
# setting up the llm
|
|
@st.cache_resource
|
|
def load_llm(model_name, provider="openai"):
|
|
if provider == "anthropic":
|
|
return Anthropic(model=model_name)
|
|
elif provider == "openai":
|
|
return OpenAI(model=model_name)
|
|
else:
|
|
raise ValueError(f"Unsupported provider: {provider}")
|
|
|
|
|
|
# utility functions
|
|
def parse_github_url(url):
|
|
pattern = r"https://github\.com/([^/]+)/([^/]+)"
|
|
match = re.match(pattern, url)
|
|
return match.groups() if match else (None, None)
|
|
|
|
def clone_repo(repo_url):
|
|
return subprocess.run(["git", "clone", repo_url], check=True, text=True, capture_output=True)
|
|
|
|
|
|
def validate_owner_repo(owner, repo):
|
|
return bool(owner) and bool(repo)
|
|
|
|
def parse_docs_by_file_types(ext, language, input_dir_path):
|
|
"""Parse documents based on file extension"""
|
|
files = glob.glob(f"{input_dir_path}/**/*{ext}", recursive=True)
|
|
|
|
if len(files) > 0:
|
|
print(f"Found {len(files)} files with extension {ext}")
|
|
loader = SimpleDirectoryReader(
|
|
input_dir=input_dir_path, required_exts=[ext], recursive=True
|
|
)
|
|
docs = loader.load_data()
|
|
parser = (
|
|
MarkdownNodeParser()
|
|
if ext == ".md"
|
|
else CodeSplitter.from_defaults(language=language)
|
|
)
|
|
nodes = parser.get_nodes_from_documents(docs)
|
|
print(f"Processed {len(nodes)} nodes from {ext} files")
|
|
return nodes
|
|
return []
|
|
|
|
|
|
# create an qdrant collection and return an index
|
|
def create_index(nodes, client):
|
|
unique_collection_id = uuid.uuid4()
|
|
collection_name = f"chat_with_docs_{unique_collection_id}"
|
|
vector_store = QdrantVectorStore(client=client, collection_name=collection_name)
|
|
storage_context = StorageContext.from_defaults(vector_store=vector_store)
|
|
index = VectorStoreIndex(
|
|
nodes,
|
|
storage_context=storage_context,
|
|
)
|
|
return index
|
|
|
|
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
|
|
|
|
|
|
def reset_chat():
|
|
st.session_state.messages = []
|
|
st.session_state.context = None
|
|
gc.collect()
|
|
|
|
|
|
with st.sidebar:
|
|
|
|
# Model selection
|
|
model_options = {
|
|
"OpenAI o3-mini": {"provider": "openai", "model": "o3-mini"},
|
|
"Claude 3.7 Sonnet": {"provider": "anthropic", "model": "claude-3-7-sonnet-20250219"}
|
|
}
|
|
selected_model = st.selectbox(
|
|
"Select Model",
|
|
options=list(model_options.keys()),
|
|
index=0
|
|
)
|
|
|
|
# Input for GitHub URL
|
|
github_url = st.text_input("GitHub Repository URL")
|
|
|
|
# Button to load and process the GitHub repository
|
|
process_button = st.button("Load")
|
|
|
|
message_container = st.empty() # Placeholder for dynamic messages
|
|
|
|
if process_button and github_url:
|
|
owner, repo = parse_github_url(github_url)
|
|
if validate_owner_repo(owner, repo):
|
|
with st.spinner(f"Loading {repo} repository by {owner}..."):
|
|
try:
|
|
input_dir_path = f"./{repo}"
|
|
|
|
if not os.path.exists(input_dir_path):
|
|
subprocess.run(["git", "clone", github_url], check=True, text=True, capture_output=True)
|
|
|
|
if os.path.exists(input_dir_path):
|
|
file_types = {
|
|
".md": "markdown",
|
|
".py": "python",
|
|
".ipynb": "python",
|
|
".js": "javascript",
|
|
".ts": "typescript"
|
|
}
|
|
|
|
nodes = []
|
|
for ext, language in file_types.items():
|
|
nodes += parse_docs_by_file_types(ext, language, input_dir_path)
|
|
else:
|
|
st.error('Error occurred while cloning the repository, carefully check the url')
|
|
st.stop()
|
|
|
|
# setting up the embedding model
|
|
Settings.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-base-en-v1.5")
|
|
try:
|
|
index = create_index(nodes)
|
|
except:
|
|
index = VectorStoreIndex(nodes=nodes)
|
|
|
|
# ====== Setup a query engine ======
|
|
model_info = model_options[selected_model]
|
|
Settings.llm = load_llm(model_name=model_info["model"], provider=model_info["provider"])
|
|
query_engine = index.as_query_engine(streaming=True, similarity_top_k=4)
|
|
|
|
# ====== Customise prompt template ======
|
|
qa_prompt_tmpl_str = (
|
|
"Context information is below.\n"
|
|
"---------------------\n"
|
|
"{context_str}\n"
|
|
"---------------------\n"
|
|
"Given the context information and your knowledge, I want you to think step by step to answer the query in a crisp manner, 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}
|
|
)
|
|
|
|
if nodes:
|
|
message_container.success("Data loaded successfully!!")
|
|
else:
|
|
message_container.write(
|
|
"No data found, check if the repository is not empty!"
|
|
)
|
|
st.session_state.query_engine = query_engine
|
|
|
|
except Exception as e:
|
|
st.error(f"An error occurred: {e}")
|
|
st.stop()
|
|
|
|
st.success("Ready to Chat!")
|
|
else:
|
|
st.error('Invalid owner or repository')
|
|
st.stop()
|
|
|
|
col1, col2 = st.columns([6, 1])
|
|
|
|
with col1:
|
|
st.header(f"Claude 3.7 Sonnet vs OpenAI o3! </>")
|
|
|
|
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 = ""
|
|
|
|
# context = st.session_state.context
|
|
query_engine = st.session_state.query_engine
|
|
|
|
# Simulate stream of response with milliseconds delay
|
|
streaming_response = query_engine.query(prompt)
|
|
|
|
for chunk in streaming_response.response_gen:
|
|
full_response += chunk
|
|
message_placeholder.markdown(full_response + "▌")
|
|
|
|
# full_response = query_engine.query(prompt)
|
|
|
|
message_placeholder.markdown(full_response)
|
|
# st.session_state.context = ctx
|
|
|
|
# Add assistant response to chat history
|
|
st.session_state.messages.append({"role": "assistant", "content": full_response}) |