app.py

import streamlit as st
import os
import pandas as pd
from command_center import CommandCenter
from process_documents import process_documents
from embed_documents import create_retriever
import json
from langchain.callbacks import get_openai_callback
from langchain_openai import ChatOpenAI
import base64
from chat_chains import rag_chain, parse_model_response
from langchain_core.messages import AIMessage, HumanMessage
from autoqa_chains import auto_qa_chain, followup_qa_chain, auto_qa_output_parser

st.set_page_config(layout="wide")
os.environ["OPENAI_API_KEY"] = "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"

format_citations = lambda citations: "\n\n".join(
    [f"{citation['quote']} ... [{citation['source_id']}]" for citation in citations]
)


def session_state_2_llm_chat_history(session_state):
    chat_history = []
    for ss in session_state:
        if not ss[0].startswith("/"):
            chat_history.append(HumanMessage(content=ss[0]))
            chat_history.append(AIMessage(content=ss[1]))
    return chat_history


ai_message_format = lambda message, references: (
    f"{message}\n\n---\n\n{format_citations(references)}"
    if references != ""
    else message
)

welcome_message = """
Hi I'm Agent Zeta, your AI assistant, dedicated to making your journey through machine learning research papers as insightful and interactive as possible. Whether you're diving into the latest studies or brushing up on foundational papers, I'm here to help navigate, discuss, and analyze content with you.

Here's a quick guide to getting started with me:

| Command | Description |
|---------|-------------|
| `/upload` <list of urls> | Upload and process documents for our conversation. |
| `/index` | View an index of processed documents to easily navigate your research. |
| `/cost` | Calculate the cost of our conversation, ensuring transparency in resource usage. |
| `/download` | Download conversation data for your records or further analysis. |
| `/auto` <document id> | Automatically generate questions and answers for a document. |

<br>

Feel free to use these commands to enhance your research experience. Let's embark on this exciting journey of discovery together!

Use `/man` at any point of time to view this guide again.
"""


def process_documents_wrapper(inputs):
    snippets, documents = process_documents(inputs)
    st.session_state.retriever = create_retriever(snippets)
    st.session_state.source_doc_urls = inputs
    st.session_state.index = [
        [snip.metadata["chunk_id"], snip.metadata["header"]] for snip in snippets
    ]
    response = f"Uploaded and processed documents {inputs}"
    st.session_state.messages.append((f"/upload {inputs}", response, ""))
    st.session_state.documents = documents
    return response


def index_documents_wrapper(inputs=None):
    response = pd.DataFrame(
        st.session_state.index, columns=["id", "reference"]
    ).to_markdown()
    st.session_state.messages.append(("/index", response, ""))
    return response


def calculate_cost_wrapper(inputs=None):
    try:
        stats_df = pd.DataFrame(st.session_state.costing)
        stats_df.loc["total"] = stats_df.sum()
        response = stats_df.to_markdown()
    except ValueError:
        response = "No cost incurred yet"
    st.session_state.messages.append(("/cost", response, ""))
    return response


def download_conversation_wrapper(inputs=None):
    conversation_data = json.dumps(
        {
            "document_urls": (
                st.session_state.source_doc_urls
                if "source_doc_urls" in st.session_state
                else []
            ),
            "document_snippets": (
                st.session_state.index if "index" in st.session_state else []
            ),
            "conversation": [
                {"human": message[0], "ai": message[1], "references": message[2]}
                for message in st.session_state.messages
            ],
            "costing": (
                st.session_state.costing if "costing" in st.session_state else []
            ),
            "total_cost": (
                {
                    k: sum(d[k] for d in st.session_state.costing)
                    for k in st.session_state.costing[0]
                }
                if "costing" in st.session_state and len(st.session_state.costing) > 0
                else {}
            ),
        }
    )
    conversation_data = base64.b64encode(conversation_data.encode()).decode()
    st.session_state.messages.append(("/download", "Conversation data downloaded", ""))
    return f'<a href="data:text/csv;base64,{conversation_data}" download="conversation_data.json">Download Conversation</a>'


def query_llm_wrapper(inputs):
    retriever = st.session_state.retriever
    qa_chain = rag_chain(
        retriever, ChatOpenAI(model="gpt-4-0125-preview", temperature=0)
    )
    relevant_docs = retriever.get_relevant_documents(inputs)
    with get_openai_callback() as cb:
        response = qa_chain.invoke(
            {
                "question": inputs,
                "chat_history": session_state_2_llm_chat_history(
                    st.session_state.messages
                ),
            }
        ).content
        stats = cb
    response = parse_model_response(response)
    answer = response["answer"]
    citations = response["citations"]
    citations.append(
        {
            "source_id": " ".join(
                [
                    f"[{ref}]"
                    for ref in sorted(
                        [ref.metadata["chunk_id"] for ref in relevant_docs],
                        key=lambda x: int(x.split("_")[1]),
                    )
                ]
            ),
            "quote": "other sources",
        }
    )

    st.session_state.messages.append((inputs, answer, citations))
    st.session_state.costing.append(
        {
            "prompt tokens": stats.prompt_tokens,
            "completion tokens": stats.completion_tokens,
            "cost": stats.total_cost,
        }
    )
    return answer, citations


def auto_qa_chain_wrapper(inputs):
    document = st.session_state.documents[inputs]
    llm = ChatOpenAI(model="gpt-4-turbo-preview", temperature=0)
    auto_qa_conversation = []
    with get_openai_callback() as cb:
        auto_qa_response = auto_qa_chain(llm).invoke({"paper": document})
        auto_qa_response_parsed = auto_qa_output_parser.invoke(auto_qa_response)[
            "questions"
        ]
        auto_qa_conversation = [
            (f'/auto {qa["question"]}', qa["answer"], "")
            for qa in auto_qa_response_parsed
        ]
        stats = cb
    st.session_state.messages.append(
        (f"/auto {inputs}", "Auto Convervation Generated", "")
    )
    for qa in auto_qa_conversation:
        st.session_state.messages.append((qa[0], qa[1], ""))

    st.session_state.costing.append(
        {
            "prompt tokens": stats.prompt_tokens,
            "completion tokens": stats.completion_tokens,
            "cost": stats.total_cost,
        }
    )
    return "\n\n".join(
        f"Q: {qa['question']}\n\nA: {qa['answer']}" for qa in auto_qa_response_parsed
    )


def boot(command_center):
    st.write("# Agent Zeta")
    if "costing" not in st.session_state:
        st.session_state.costing = []
    if "messages" not in st.session_state:
        st.session_state.messages = []
    st.chat_message("ai").write(welcome_message, unsafe_allow_html=True)
    for message in st.session_state.messages:
        st.chat_message("human").write(message[0])
        st.chat_message("ai").write(
            ai_message_format(message[1], message[2]), unsafe_allow_html=True
        )
    if query := st.chat_input():
        st.chat_message("human").write(query)
        response = command_center.execute_command(query)
        if response is None:
            pass
        elif type(response) == tuple:
            result, references = response
            st.chat_message("ai").write(
                ai_message_format(result, references), unsafe_allow_html=True
            )
        else:
            st.chat_message("ai").write(response, unsafe_allow_html=True)


if __name__ == "__main__":
    all_commands = [
        ("/upload", list, process_documents_wrapper),
        ("/index", None, index_documents_wrapper),
        ("/cost", None, calculate_cost_wrapper),
        ("/download", None, download_conversation_wrapper),
        ("/man", None, lambda x: welcome_message),
        ("/auto", int, auto_qa_chain_wrapper),
    ]
    command_center = CommandCenter(
        default_input_type=str,
        default_function=query_llm_wrapper,
        all_commands=all_commands,
    )
    boot(command_center)