chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,243 @@
|
||||
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})
|
||||
Reference in New Issue
Block a user