app.py

import os
import gc
import re
import uuid
import subprocess
import requests
from dotenv import load_dotenv

os.environ["HF_HOME"] = "weights"
os.environ["TORCH_HOME"] = "weights"
import streamlit as st

from llama_index.core import Settings
from llama_index.llms.ollama import Ollama
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 langchain_community.embeddings import HuggingFaceEmbeddings
from llama_index.embeddings.langchain import LangchainEmbedding

from rag_101.retriever import (
    load_embedding_model,
    load_reranker_model
)

# setup the llm
ollama_url = 'http://localhost:11434/api/chat'

llm = Ollama(model="mistral:instruct", url=ollama_url ,request_timeout=1000.0)

# TODO: setup the embedding model
lc_embedding_model = load_embedding_model()
embed_model = LangchainEmbedding(lc_embedding_model)


# 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):
    try:
        result = subprocess.run(["git", "clone", repo_url], check=True, text=True, capture_output=True)
        print(result.stdout)
        return result
    except subprocess.CalledProcessError as e:
        print(f"Error occurred: {e.stderr}")
        raise e


def validate_owner_repo(owner, repo):
    return bool(owner) and bool(repo)


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:

    # 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)):
                        clone_repo(github_url)

                    if os.path.exists(input_dir_path):
                        loader = SimpleDirectoryReader(
                            input_dir=input_dir_path,
                            required_exts=[".py", ".ipynb", ".js", ".ts", ".md"],
                            recursive=True
                        )
                    else:
                        st.error('Error occurred while cloning the repo, carefully check the URL')
                        st.stop()

                    docs = loader.load_data()

                    # TODO: ====== Create vector store and upload data ======
                    Settings.embed_model = embed_model
                    index = VectorStoreIndex.from_documents(docs)

                    # setup a query engine
                    Settings.llm = llm
                    query_engine = index.as_query_engine(streaming=True, similarity_top_k=4)

                    # customize 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 crisp manner, in case you don't know the answer say 'I don't know!'.\n"
                        "Query: {query_str}\n"
                        "Answer: "
                    )
                    qa_prompt_tmpl_str = PromptTemplate(qa_prompt_tmpl_str)

                    query_engine.update_prompts(
                        {"response_synthesizer:text_qa_template": qa_prompt_tmpl_str}
                    )

                    if docs:
                        message_container.success("Data loaded successfully!!")
                    else:
                        message_container.write(
                            "No Data found, check if 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 repo')
            st.stop()

col1, col2 = st.columns([6, 1])

with col1:
    st.header(f"GitChat🌐")
    st.header(f"Chat with your code! </>")

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
# TODO: old one
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})

# todo: new one
# prompt = st.chat_input("What's up?")
# if prompt:
#     # 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
#     query_engine = st.session_state.query_engine
#
#     # simulate stream of response with milliseconds delay
#     try:
#         # Construct the request payload
#         payload = {
#             "message": prompt,
#             "model": "mistral:instruct"
#         }
#
#         # Send the request
#         response = requests.post(ollama_url, json=payload)
#
#         # Check for HTTP errors
#         response.raise_for_status()
#
#         # Print the full response to debug
#         response_json = response.json()
#         print(response_json)
#
#         # Process the response
#         if "response_gen" in response_json:
#             for chunk in response_json["response_gen"]:
#                 full_response += chunk
#                 message_placeholder.markdown(full_response + "▌")
#             message_placeholder.markdown(full_response)
#
#             # add assistant response to chat history
#             st.session_state.messages.append({"role": "assistant", "content": full_response})
#         else:
#             st.error("Unexpected response format: 'response_gen' key not found")
#
#     except requests.exceptions.HTTPError as e:
#         st.error(f"HTTP error: {e}")