| 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 chains import rag_chain, parse_model_response | |
| from langchain_core.messages import AIMessage, HumanMessage | |
| 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` | 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. | | |
| <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 = 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, "")) | |
| 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 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), | |
| ] | |
| command_center = CommandCenter( | |
| default_input_type=str, | |
| default_function=query_llm_wrapper, | |
| all_commands=all_commands, | |
| ) | |
| boot(command_center) | |