langchain-streamlit demo with streaming llm, memory, and langsmith feedback
from datetime import datetime
from typing import Tuple, List, Dict, Any, Union, Optional

import anthropic
import langsmith.utils
import openai
import streamlit as st
from langchain.callbacks.tracers.langchain import LangChainTracer, wait_for_all_tracers
from langchain.callbacks.tracers.run_collector import RunCollectorCallbackHandler
from langchain.memory import ConversationBufferMemory, StreamlitChatMessageHistory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema.document import Document
from langchain.schema.retriever import BaseRetriever
from langsmith.client import Client
from streamlit_feedback import streamlit_feedback

from defaults import default_values

from llm_resources import get_runnable, get_llm, get_texts_and_retriever, StreamHandler

__version__ = "0.0.14"

# --- Initialization ---
st.set_page_config(
    page_title=f"langchain-streamlit-demo v{__version__}",
    page_icon="🦜",
)


def st_init_null(*variable_names) -> None:
    for variable_name in variable_names:
        if variable_name not in st.session_state:
            st.session_state[variable_name] = None


st_init_null(
    "chain",
    "client",
    "doc_chain",
    "document_chat_chain_type",
    "llm",
    "ls_tracer",
    "provider",
    "retriever",
    "run",
    "run_id",
    "trace_link",
    "LANGSMITH_API_KEY",
    "LANGSMITH_PROJECT",
    "AZURE_OPENAI_BASE_URL",
    "AZURE_OPENAI_API_VERSION",
    "AZURE_OPENAI_DEPLOYMENT_NAME",
    "AZURE_OPENAI_EMB_DEPLOYMENT_NAME",
    "AZURE_OPENAI_API_KEY",
    "AZURE_OPENAI_MODEL_VERSION",
    "AZURE_AVAILABLE",
)

# --- LLM globals ---
STMEMORY = StreamlitChatMessageHistory(key="langchain_messages")
MEMORY = ConversationBufferMemory(
    chat_memory=STMEMORY,
    return_messages=True,
    memory_key="chat_history",
)

RUN_COLLECTOR = RunCollectorCallbackHandler()

st.session_state.LANGSMITH_API_KEY = (
    st.session_state.LANGSMITH_API_KEY
    or default_values.PROVIDER_KEY_DICT.get("LANGSMITH")
)

st.session_state.LANGSMITH_PROJECT = st.session_state.LANGSMITH_PROJECT or (
    default_values.DEFAULT_LANGSMITH_PROJECT or "langchain-streamlit-demo"
)


def azure_state_or_default(*args):
    st.session_state.update(
        {
            arg: st.session_state.get(arg) or default_values.AZURE_DICT.get(arg)
            for arg in args
        },
    )


azure_state_or_default(
    "AZURE_OPENAI_BASE_URL",
    "AZURE_OPENAI_API_VERSION",
    "AZURE_OPENAI_DEPLOYMENT_NAME",
    "AZURE_OPENAI_EMB_DEPLOYMENT_NAME",
    "AZURE_OPENAI_API_KEY",
    "AZURE_OPENAI_MODEL_VERSION",
)

st.session_state.AZURE_AVAILABLE = all(
    [
        st.session_state.AZURE_OPENAI_BASE_URL,
        st.session_state.AZURE_OPENAI_API_VERSION,
        st.session_state.AZURE_OPENAI_DEPLOYMENT_NAME,
        st.session_state.AZURE_OPENAI_API_KEY,
        st.session_state.AZURE_OPENAI_MODEL_VERSION,
    ],
)

st.session_state.AZURE_EMB_AVAILABLE = (
    st.session_state.AZURE_AVAILABLE
    and st.session_state.AZURE_OPENAI_EMB_DEPLOYMENT_NAME
)

AZURE_KWARGS = (
    None
    if not st.session_state.AZURE_EMB_AVAILABLE
    else {
        "openai_api_base": st.session_state.AZURE_OPENAI_BASE_URL,
        "openai_api_version": st.session_state.AZURE_OPENAI_API_VERSION,
        "deployment": st.session_state.AZURE_OPENAI_EMB_DEPLOYMENT_NAME,
        "openai_api_key": st.session_state.AZURE_OPENAI_API_KEY,
        "openai_api_type": "azure",
    }
)


@st.cache_data
def get_texts_and_retriever_cacheable_wrapper(
    uploaded_file_bytes: bytes,
    openai_api_key: str,
    chunk_size: int = default_values.DEFAULT_CHUNK_SIZE,
    chunk_overlap: int = default_values.DEFAULT_CHUNK_OVERLAP,
    k: int = default_values.DEFAULT_RETRIEVER_K,
    azure_kwargs: Optional[Dict[str, str]] = None,
    use_azure: bool = False,
) -> Tuple[List[Document], BaseRetriever]:
    return get_texts_and_retriever(
        uploaded_file_bytes=uploaded_file_bytes,
        openai_api_key=openai_api_key,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        k=k,
        azure_kwargs=azure_kwargs,
        use_azure=use_azure,
    )


# --- Sidebar ---
sidebar = st.sidebar
with sidebar:
    st.markdown("# Menu")

    model = st.selectbox(
        label="Chat Model",
        options=default_values.SUPPORTED_MODELS,
        index=default_values.SUPPORTED_MODELS.index(default_values.DEFAULT_MODEL),
    )

    st.session_state.provider = default_values.MODEL_DICT[model]

    provider_api_key = (
        default_values.PROVIDER_KEY_DICT.get(
            st.session_state.provider,
        )
        or st.text_input(
            f"{st.session_state.provider} API key",
            type="password",
        )
        if st.session_state.provider != "Azure OpenAI"
        else ""
    )

    if st.button("Clear message history"):
        STMEMORY.clear()
        st.session_state.trace_link = None
        st.session_state.run_id = None

    # --- Document Chat Options ---
    with st.expander("Document Chat", expanded=False):
        uploaded_file = st.file_uploader("Upload a PDF", type="pdf")

        openai_api_key = (
            provider_api_key
            if st.session_state.provider == "OpenAI"
            else default_values.OPENAI_API_KEY
            or st.sidebar.text_input("OpenAI API Key: ", type="password")
        )

        document_chat = st.checkbox(
            "Document Chat",
            value=False,
            help="Uploaded document will provide context for the chat.",
        )

        k = st.slider(
            label="Number of Chunks",
            help="How many document chunks will be used for context?",
            value=default_values.DEFAULT_RETRIEVER_K,
            min_value=1,
            max_value=10,
        )

        chunk_size = st.slider(
            label="Number of Tokens per Chunk",
            help="Size of each chunk of text",
            min_value=default_values.MIN_CHUNK_SIZE,
            max_value=default_values.MAX_CHUNK_SIZE,
            value=default_values.DEFAULT_CHUNK_SIZE,
        )

        chunk_overlap = st.slider(
            label="Chunk Overlap",
            help="Number of characters to overlap between chunks",
            min_value=default_values.MIN_CHUNK_OVERLAP,
            max_value=default_values.MAX_CHUNK_OVERLAP,
            value=default_values.DEFAULT_CHUNK_OVERLAP,
        )

        chain_type_help_root = (
            "https://python.langchain.com/docs/modules/chains/document/"
        )

        chain_type_help = "\n".join(
            f"- [{chain_type_name}]({chain_type_help_root}/{chain_type_name})"
            for chain_type_name in (
                "stuff",
                "refine",
                "map_reduce",
                "map_rerank",
            )
        )

        document_chat_chain_type = st.selectbox(
            label="Document Chat Chain Type",
            options=[
                "stuff",
                "refine",
                "map_reduce",
                "map_rerank",
                "Q&A Generation",
                "Summarization",
            ],
            index=0,
            help=chain_type_help,
            disabled=not document_chat,
        )
        use_azure = st.toggle(
            label="Use Azure OpenAI",
            value=st.session_state.AZURE_EMB_AVAILABLE,
            help="Use Azure for embeddings instead of using OpenAI directly.",
        )

        if uploaded_file:
            if st.session_state.AZURE_EMB_AVAILABLE or openai_api_key:
                (
                    st.session_state.texts,
                    st.session_state.retriever,
                ) = get_texts_and_retriever_cacheable_wrapper(
                    uploaded_file_bytes=uploaded_file.getvalue(),
                    openai_api_key=openai_api_key,
                    chunk_size=chunk_size,
                    chunk_overlap=chunk_overlap,
                    k=k,
                    azure_kwargs=AZURE_KWARGS,
                    use_azure=use_azure,
                )
            else:
                st.error("Please enter a valid OpenAI API key.", icon="❌")

    # --- Advanced Settings ---
    with st.expander("Advanced Settings", expanded=False):
        st.markdown("## Feedback Scale")
        use_faces = st.toggle(label="`Thumbs` ⇄ `Faces`", value=False)
        feedback_option = "faces" if use_faces else "thumbs"

        system_prompt = (
            st.text_area(
                "Custom Instructions",
                default_values.DEFAULT_SYSTEM_PROMPT,
                help="Custom instructions to provide the language model to determine style, personality, etc.",
            )
            .strip()
            .replace("{", "{{")
            .replace("}", "}}")
        )

        temperature = st.slider(
            "Temperature",
            min_value=default_values.MIN_TEMP,
            max_value=default_values.MAX_TEMP,
            value=default_values.DEFAULT_TEMP,
            help="Higher values give more random results.",
        )

        max_tokens = st.slider(
            "Max Tokens",
            min_value=default_values.MIN_MAX_TOKENS,
            max_value=default_values.MAX_MAX_TOKENS,
            value=default_values.DEFAULT_MAX_TOKENS,
            help="Higher values give longer results.",
        )

    # --- LangSmith Options ---
    if default_values.SHOW_LANGSMITH_OPTIONS:
        with st.expander("LangSmith Options", expanded=False):
            st.session_state.LANGSMITH_API_KEY = st.text_input(
                "LangSmith API Key (optional)",
                value=st.session_state.LANGSMITH_API_KEY,
                type="password",
            )

            st.session_state.LANGSMITH_PROJECT = st.text_input(
                "LangSmith Project Name",
                value=st.session_state.LANGSMITH_PROJECT,
            )

    if st.session_state.client is None and st.session_state.LANGSMITH_API_KEY:
        st.session_state.client = Client(
            api_url="https://api.smith.langchain.com",
            api_key=st.session_state.LANGSMITH_API_KEY,
        )
        st.session_state.ls_tracer = LangChainTracer(
            project_name=st.session_state.LANGSMITH_PROJECT,
            client=st.session_state.client,
        )

    # --- Azure Options ---
    if default_values.SHOW_AZURE_OPTIONS:
        with st.expander("Azure Options", expanded=False):
            st.session_state.AZURE_OPENAI_BASE_URL = st.text_input(
                "AZURE_OPENAI_BASE_URL",
                value=st.session_state.AZURE_OPENAI_BASE_URL,
            )

            st.session_state.AZURE_OPENAI_API_VERSION = st.text_input(
                "AZURE_OPENAI_API_VERSION",
                value=st.session_state.AZURE_OPENAI_API_VERSION,
            )

            st.session_state.AZURE_OPENAI_DEPLOYMENT_NAME = st.text_input(
                "AZURE_OPENAI_DEPLOYMENT_NAME",
                value=st.session_state.AZURE_OPENAI_DEPLOYMENT_NAME,
            )

            st.session_state.AZURE_OPENAI_EMB_DEPLOYMENT_NAME = st.text_input(
                "AZURE_OPENAI_EMB_DEPLOYMENT_NAME",
                value=st.session_state.AZURE_OPENAI_EMB_DEPLOYMENT_NAME,
            )

            st.session_state.AZURE_OPENAI_API_KEY = st.text_input(
                "AZURE_OPENAI_API_KEY",
                value=st.session_state.AZURE_OPENAI_API_KEY,
                type="password",
            )

            st.session_state.AZURE_OPENAI_MODEL_VERSION = st.text_input(
                "AZURE_OPENAI_MODEL_VERSION",
                value=st.session_state.AZURE_OPENAI_MODEL_VERSION,
            )


# --- LLM Instantiation ---
st.session_state.llm = get_llm(
    provider=st.session_state.provider,
    model=model,
    provider_api_key=provider_api_key,
    temperature=temperature,
    max_tokens=max_tokens,
    azure_available=st.session_state.AZURE_AVAILABLE,
    azure_dict={
        "AZURE_OPENAI_BASE_URL": st.session_state.AZURE_OPENAI_BASE_URL,
        "AZURE_OPENAI_API_VERSION": st.session_state.AZURE_OPENAI_API_VERSION,
        "AZURE_OPENAI_DEPLOYMENT_NAME": st.session_state.AZURE_OPENAI_DEPLOYMENT_NAME,
        "AZURE_OPENAI_API_KEY": st.session_state.AZURE_OPENAI_API_KEY,
        "AZURE_OPENAI_MODEL_VERSION": st.session_state.AZURE_OPENAI_MODEL_VERSION,
    },
)

# --- Chat History ---
if len(STMEMORY.messages) == 0:
    STMEMORY.add_ai_message("Hello! I'm a helpful AI chatbot. Ask me a question!")

for msg in STMEMORY.messages:
    st.chat_message(
        msg.type,
        avatar="🦜" if msg.type in ("ai", "assistant") else None,
    ).write(msg.content)


# --- Current Chat ---
if st.session_state.llm:
    # --- Regular Chat ---
    chat_prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                system_prompt + "\nIt's currently {time}.",
            ),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{query}"),
        ],
    ).partial(time=lambda: str(datetime.now()))

    # --- Chat Input ---
    prompt = st.chat_input(placeholder="Ask me a question!")
    if prompt:
        st.chat_message("user").write(prompt)
        feedback_update = None
        feedback = None

        # --- Chat Output ---
        with st.chat_message("assistant", avatar="🦜"):
            callbacks = [RUN_COLLECTOR]

            if st.session_state.ls_tracer:
                callbacks.append(st.session_state.ls_tracer)

            config: Dict[str, Any] = dict(
                callbacks=callbacks,
                tags=["Streamlit Chat"],
            )
            if st.session_state.provider == "Anthropic":
                config["max_concurrency"] = 5

            use_document_chat = all(
                [
                    document_chat,
                    st.session_state.retriever,
                ],
            )

            full_response: Union[str, None] = None

            message_placeholder = st.empty()
            stream_handler = StreamHandler(message_placeholder)
            callbacks.append(stream_handler)

            st.session_state.chain = get_runnable(
                use_document_chat,
                document_chat_chain_type,
                st.session_state.llm,
                st.session_state.retriever,
                MEMORY,
                chat_prompt,
                prompt,
            )

            # --- LLM call ---
            try:
                full_response = st.session_state.chain.invoke(prompt, config)

            except (openai.error.AuthenticationError, anthropic.AuthenticationError):
                st.error(
                    f"Please enter a valid {st.session_state.provider} API key.",
                    icon="❌",
                )

            # --- Display output ---
            if full_response is not None:
                message_placeholder.markdown(full_response)

                # --- Tracing ---
                if st.session_state.client:
                    st.session_state.run = RUN_COLLECTOR.traced_runs[0]
                    st.session_state.run_id = st.session_state.run.id
                    RUN_COLLECTOR.traced_runs = []
                    wait_for_all_tracers()
                    try:
                        st.session_state.trace_link = st.session_state.client.read_run(
                            st.session_state.run_id,
                        ).url
                    except langsmith.utils.LangSmithError:
                        st.session_state.trace_link = None

    # --- LangSmith Trace Link ---
    if st.session_state.trace_link:
        with sidebar:
            st.markdown(
                f'<a href="{st.session_state.trace_link}" target="_blank"><button>Latest Trace: 🛠️</button></a>',
                unsafe_allow_html=True,
            )

    # --- Feedback ---
    if st.session_state.client and st.session_state.run_id:
        feedback = streamlit_feedback(
            feedback_type=feedback_option,
            optional_text_label="[Optional] Please provide an explanation",
            key=f"feedback_{st.session_state.run_id}",
        )

        # Define score mappings for both "thumbs" and "faces" feedback systems
        score_mappings: dict[str, dict[str, Union[int, float]]] = {
            "thumbs": {"👍": 1, "👎": 0},
            "faces": {"😀": 1, "🙂": 0.75, "😐": 0.5, "🙁": 0.25, "😞": 0},
        }

        # Get the score mapping based on the selected feedback option
        scores = score_mappings[feedback_option]

        if feedback:
            # Get the score from the selected feedback option's score mapping
            score = scores.get(
                feedback["score"],
            )

            if score is not None:
                # Formulate feedback type string incorporating the feedback option
                # and score value
                feedback_type_str = f"{feedback_option} {feedback['score']}"

                # Record the feedback with the formulated feedback type string
                # and optional comment
                feedback_record = st.session_state.client.create_feedback(
                    st.session_state.run_id,
                    feedback_type_str,
                    score=score,
                    comment=feedback.get("text"),
                )
                st.toast("Feedback recorded!", icon="📝")
            else:
                st.warning("Invalid feedback score.")

else:
    st.error(f"Please enter a valid {st.session_state.provider} API key.", icon="❌")

langchain-streamlit-demo/app.py
joshuasundance-swca/langchain-streamlit-demo
import os
from collections import namedtuple


MODEL_DICT = {
    "gpt-3.5-turbo": "OpenAI",
    "gpt-4": "OpenAI",
    "claude-instant-v1": "Anthropic",
    "claude-2": "Anthropic",
    "meta-llama/Llama-2-7b-chat-hf": "Anyscale Endpoints",
    "meta-llama/Llama-2-13b-chat-hf": "Anyscale Endpoints",
    "meta-llama/Llama-2-70b-chat-hf": "Anyscale Endpoints",
    "codellama/CodeLlama-34b-Instruct-hf": "Anyscale Endpoints",
    "Azure OpenAI": "Azure OpenAI",
}

SUPPORTED_MODELS = list(MODEL_DICT.keys())

DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "gpt-3.5-turbo")

DEFAULT_SYSTEM_PROMPT = os.environ.get(
    "DEFAULT_SYSTEM_PROMPT",
    "You are a helpful chatbot.",
)

MIN_TEMP = float(os.environ.get("MIN_TEMPERATURE", 0.0))
MAX_TEMP = float(os.environ.get("MAX_TEMPERATURE", 1.0))
DEFAULT_TEMP = float(os.environ.get("DEFAULT_TEMPERATURE", 0.7))

MIN_MAX_TOKENS = int(os.environ.get("MIN_MAX_TOKENS", 1))
MAX_MAX_TOKENS = int(os.environ.get("MAX_MAX_TOKENS", 100000))
DEFAULT_MAX_TOKENS = int(os.environ.get("DEFAULT_MAX_TOKENS", 1000))

DEFAULT_LANGSMITH_PROJECT = os.environ.get("LANGCHAIN_PROJECT")

AZURE_VARS = [
    "AZURE_OPENAI_BASE_URL",
    "AZURE_OPENAI_API_VERSION",
    "AZURE_OPENAI_DEPLOYMENT_NAME",
    "AZURE_OPENAI_EMB_DEPLOYMENT_NAME",
    "AZURE_OPENAI_API_KEY",
    "AZURE_OPENAI_MODEL_VERSION",
]

AZURE_DICT = {v: os.environ.get(v, "") for v in AZURE_VARS}


SHOW_LANGSMITH_OPTIONS = (
    os.environ.get("SHOW_LANGSMITH_OPTIONS", "true").lower() == "true"
)
SHOW_AZURE_OPTIONS = os.environ.get("SHOW_AZURE_OPTIONS", "true").lower() == "true"

PROVIDER_KEY_DICT = {
    "OpenAI": os.environ.get("OPENAI_API_KEY", ""),
    "Anthropic": os.environ.get("ANTHROPIC_API_KEY", ""),
    "Anyscale Endpoints": os.environ.get("ANYSCALE_API_KEY", ""),
    "LANGSMITH": os.environ.get("LANGCHAIN_API_KEY", ""),
}

OPENAI_API_KEY = PROVIDER_KEY_DICT["OpenAI"]


MIN_CHUNK_SIZE = 1
MAX_CHUNK_SIZE = 10000
DEFAULT_CHUNK_SIZE = 1000

MIN_CHUNK_OVERLAP = 0
MAX_CHUNK_OVERLAP = 10000
DEFAULT_CHUNK_OVERLAP = 0

DEFAULT_RETRIEVER_K = 4

DEFAULT_VALUES = namedtuple(
    "DEFAULT_VALUES",
    [
        "MODEL_DICT",
        "SUPPORTED_MODELS",
        "DEFAULT_MODEL",
        "DEFAULT_SYSTEM_PROMPT",
        "MIN_TEMP",
        "MAX_TEMP",
        "DEFAULT_TEMP",
        "MIN_MAX_TOKENS",
        "MAX_MAX_TOKENS",
        "DEFAULT_MAX_TOKENS",
        "DEFAULT_LANGSMITH_PROJECT",
        "AZURE_VARS",
        "AZURE_DICT",
        "PROVIDER_KEY_DICT",
        "OPENAI_API_KEY",
        "MIN_CHUNK_SIZE",
        "MAX_CHUNK_SIZE",
        "DEFAULT_CHUNK_SIZE",
        "MIN_CHUNK_OVERLAP",
        "MAX_CHUNK_OVERLAP",
        "DEFAULT_CHUNK_OVERLAP",
        "DEFAULT_RETRIEVER_K",
        "SHOW_LANGSMITH_OPTIONS",
        "SHOW_AZURE_OPTIONS",
    ],
)


default_values = DEFAULT_VALUES(
    MODEL_DICT,
    SUPPORTED_MODELS,
    DEFAULT_MODEL,
    DEFAULT_SYSTEM_PROMPT,
    MIN_TEMP,
    MAX_TEMP,
    DEFAULT_TEMP,
    MIN_MAX_TOKENS,
    MAX_MAX_TOKENS,
    DEFAULT_MAX_TOKENS,
    DEFAULT_LANGSMITH_PROJECT,
    AZURE_VARS,
    AZURE_DICT,
    PROVIDER_KEY_DICT,
    OPENAI_API_KEY,
    MIN_CHUNK_SIZE,
    MAX_CHUNK_SIZE,
    DEFAULT_CHUNK_SIZE,
    MIN_CHUNK_OVERLAP,
    MAX_CHUNK_OVERLAP,
    DEFAULT_CHUNK_OVERLAP,
    DEFAULT_RETRIEVER_K,
    SHOW_LANGSMITH_OPTIONS,
    SHOW_AZURE_OPTIONS,
)

langchain-streamlit-demo/defaults.py
joshuasundance-swca/langchain-streamlit-demo
from tempfile import NamedTemporaryFile
from typing import Tuple, List, Optional, Dict

from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import RetrievalQA, LLMChain
from langchain.chat_models import (
    AzureChatOpenAI,
    ChatOpenAI,
    ChatAnthropic,
    ChatAnyscale,
)
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.schema import Document, BaseRetriever
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS

from defaults import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, DEFAULT_RETRIEVER_K
from qagen import get_rag_qa_gen_chain
from summarize import get_rag_summarization_chain


def get_runnable(
    use_document_chat: bool,
    document_chat_chain_type: str,
    llm,
    retriever,
    memory,
    chat_prompt,
    summarization_prompt,
):
    if not use_document_chat:
        return LLMChain(
            prompt=chat_prompt,
            llm=llm,
            memory=memory,
        ) | (lambda output: output["text"])

    if document_chat_chain_type == "Q&A Generation":
        return get_rag_qa_gen_chain(
            retriever,
            llm,
        )
    elif document_chat_chain_type == "Summarization":
        return get_rag_summarization_chain(
            summarization_prompt,
            retriever,
            llm,
        )
    else:
        return RetrievalQA.from_chain_type(
            llm=llm,
            chain_type=document_chat_chain_type,
            retriever=retriever,
            memory=memory,
            output_key="output_text",
        ) | (lambda output: output["output_text"])


def get_llm(
    provider: str,
    model: str,
    provider_api_key: str,
    temperature: float,
    max_tokens: int,
    azure_available: bool,
    azure_dict: dict[str, str],
):
    if azure_available and provider == "Azure OpenAI":
        return AzureChatOpenAI(
            openai_api_base=azure_dict["AZURE_OPENAI_BASE_URL"],
            openai_api_version=azure_dict["AZURE_OPENAI_API_VERSION"],
            deployment_name=azure_dict["AZURE_OPENAI_DEPLOYMENT_NAME"],
            openai_api_key=azure_dict["AZURE_OPENAI_API_KEY"],
            openai_api_type="azure",
            model_version=azure_dict["AZURE_OPENAI_MODEL_VERSION"],
            temperature=temperature,
            streaming=True,
            max_tokens=max_tokens,
        )

    elif provider_api_key:
        if provider == "OpenAI":
            return ChatOpenAI(
                model_name=model,
                openai_api_key=provider_api_key,
                temperature=temperature,
                streaming=True,
                max_tokens=max_tokens,
            )

        elif provider == "Anthropic":
            return ChatAnthropic(
                model=model,
                anthropic_api_key=provider_api_key,
                temperature=temperature,
                streaming=True,
                max_tokens_to_sample=max_tokens,
            )

        elif provider == "Anyscale Endpoints":
            return ChatAnyscale(
                model_name=model,
                anyscale_api_key=provider_api_key,
                temperature=temperature,
                streaming=True,
                max_tokens=max_tokens,
            )

    return None


def get_texts_and_retriever(
    uploaded_file_bytes: bytes,
    openai_api_key: str,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    k: int = DEFAULT_RETRIEVER_K,
    azure_kwargs: Optional[Dict[str, str]] = None,
    use_azure: bool = False,
) -> Tuple[List[Document], BaseRetriever]:
    with NamedTemporaryFile() as temp_file:
        temp_file.write(uploaded_file_bytes)
        temp_file.seek(0)

        loader = PyPDFLoader(temp_file.name)
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )
        texts = text_splitter.split_documents(documents)
        embeddings_kwargs = {"openai_api_key": openai_api_key}
        if use_azure and azure_kwargs:
            embeddings_kwargs.update(azure_kwargs)
        embeddings = OpenAIEmbeddings(**embeddings_kwargs)

        bm25_retriever = BM25Retriever.from_documents(texts)
        bm25_retriever.k = k

        faiss_vectorstore = FAISS.from_documents(texts, embeddings)
        faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": k})

        ensemble_retriever = EnsembleRetriever(
            retrievers=[bm25_retriever, faiss_retriever],
            weights=[0.5, 0.5],
        )

        return texts, ensemble_retriever


class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, initial_text=""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)

langchain-streamlit-demo/llm_resources.py
joshuasundance-swca/langchain-streamlit-demo
from typing import List

from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.prompts.chat import (
    ChatPromptTemplate,
)
from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.retriever import BaseRetriever
from langchain.schema.runnable import RunnablePassthrough, RunnableSequence
from pydantic import BaseModel, Field


class QuestionAnswerPair(BaseModel):
    question: str = Field(..., description="The question that will be answered.")
    answer: str = Field(..., description="The answer to the question that was asked.")

    def to_str(self, idx: int) -> str:
        question_piece = f"{idx}. **Q:** {self.question}"
        whitespace = " " * (len(str(idx)) + 2)
        answer_piece = f"{whitespace}**A:** {self.answer}"
        return f"{question_piece}\n\n{answer_piece}"


class QuestionAnswerPairList(BaseModel):
    QuestionAnswerPairs: List[QuestionAnswerPair]

    def to_str(self) -> str:
        return "\n\n".join(
            [
                qap.to_str(idx)
                for idx, qap in enumerate(self.QuestionAnswerPairs, start=1)
            ],
        )


PYDANTIC_PARSER: PydanticOutputParser = PydanticOutputParser(
    pydantic_object=QuestionAnswerPairList,
)


templ1 = """You are a smart assistant designed to help college professors come up with reading comprehension questions.
Given a piece of text, you must come up with question and answer pairs that can be used to test a student's reading comprehension abilities.
Generate as many question/answer pairs as you can.
When coming up with the question/answer pairs, you must respond in the following format:
{format_instructions}

Do not provide additional commentary and do not wrap your response in Markdown formatting. Return RAW, VALID JSON.
"""
templ2 = """{prompt}
Please create question/answer pairs, in the specified JSON format, for the following text:
----------------
{context}"""
CHAT_PROMPT = ChatPromptTemplate.from_messages(
    [
        ("system", templ1),
        ("human", templ2),
    ],
).partial(format_instructions=PYDANTIC_PARSER.get_format_instructions)


def get_rag_qa_gen_chain(
    retriever: BaseRetriever,
    llm: BaseLanguageModel,
    input_key: str = "prompt",
) -> RunnableSequence:
    return (
        {"context": retriever, input_key: RunnablePassthrough()}
        | CHAT_PROMPT
        | llm
        | OutputFixingParser.from_llm(llm=llm, parser=PYDANTIC_PARSER)
        | (lambda parsed_output: parsed_output.to_str())
    )

langchain-streamlit-demo/qagen.py
joshuasundance-swca/langchain-streamlit-demo
from langchain.chains.base import Chain
from langchain.chains.summarize import load_summarize_chain
from langchain.prompts import PromptTemplate
from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.retriever import BaseRetriever
from langchain.schema.runnable import RunnableSequence, RunnablePassthrough

prompt_template = """Write a concise summary of the following text, based on the user input.
User input: {query}
Text:
```
{text}
```
CONCISE SUMMARY:"""

refine_template = (
    "You are iteratively crafting a summary of the text below based on the user input\n"
    "User input: {query}\n"
    "We have provided an existing summary up to a certain point: {existing_answer}\n"
    "We have the opportunity to refine the existing summary"
    "(only if needed) with some more context below.\n"
    "------------\n"
    "{text}\n"
    "------------\n"
    "Given the new context, refine the original summary.\n"
    "If the context isn't useful, return the original summary.\n"
    "If the context is useful, refine the summary to include the new context.\n"
    "Your contribution is helping to build a comprehensive summary of a large body of knowledge.\n"
    "You do not have the complete context, so do not discard pieces of the original summary."
)


def get_summarization_chain(
    llm: BaseLanguageModel,
    prompt: str,
) -> Chain:
    _prompt = PromptTemplate.from_template(
        prompt_template,
        partial_variables={"query": prompt},
    )
    refine_prompt = PromptTemplate.from_template(
        refine_template,
        partial_variables={"query": prompt},
    )
    return load_summarize_chain(
        llm=llm,
        chain_type="refine",
        question_prompt=_prompt,
        refine_prompt=refine_prompt,
        return_intermediate_steps=False,
        input_key="input_documents",
        output_key="output_text",
    )


def get_rag_summarization_chain(
    prompt: str,
    retriever: BaseRetriever,
    llm: BaseLanguageModel,
    input_key: str = "prompt",
) -> RunnableSequence:
    return (
        {"input_documents": retriever, input_key: RunnablePassthrough()}
        | get_summarization_chain(llm, prompt)
        | (lambda output: output["output_text"])
    )

langchain-streamlit-demo/summarize.py
joshuasundance-swca/langchain-streamlit-demo