citations

app.py

@@ -6,22 +6,34 @@ from process_documents import process_documents
 from embed_documents import create_retriever
 import json
 from langchain.callbacks import get_openai_callback
-from langchain.chains import ConversationalRetrievalChain
 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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"
 
-get_references = lambda relevant_docs: " ".join(
-    [f"[{ref}]" for ref in sorted([ref.metadata["chunk_id"] for ref in relevant_docs])]
+format_citations = lambda citations: "\n\n".join(
+    [f"{citation['quote']} ... [{citation['source_id']}]" for citation in citations]
 )
-session_state_2_llm_chat_history = lambda session_state: [
-    ss[:2] for ss in session_state if not ss[0].startswith("/")
-]
+
+
+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{references}" if references != "" else message
+    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.
 
@@ -108,26 +120,39 @@ def download_conversation_wrapper(inputs=None):
 
 def query_llm_wrapper(inputs):
     retriever = st.session_state.retriever
-    qa_chain = ConversationalRetrievalChain.from_llm(
-        llm=ChatOpenAI(model="gpt-4-0125-preview", temperature=0),
-        retriever=retriever,
-        return_source_documents=True,
-        chain_type="stuff",
+    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:
-        result = qa_chain(
+        response = qa_chain.invoke(
             {
                 "question": inputs,
                 "chat_history": session_state_2_llm_chat_history(
                     st.session_state.messages
                 ),
             }
-        )
+        ).content
         stats = cb
-    result = result["answer"]
-    references = get_references(relevant_docs)
-    st.session_state.messages.append((inputs, result, references))
+    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,
@@ -135,7 +160,7 @@ def query_llm_wrapper(inputs):
             "cost": stats.total_cost,
         }
     )
-    return result, references
+    return answer, citations
 
 
 def boot(command_center):

chains.py

@@ -0,0 +1,107 @@
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+import xml.etree.ElementTree as ET
+import re
+
+contextualize_q_system_prompt = """Given a chat history and the latest user question \
+which might reference context in the chat history, formulate a standalone question \
+which can be understood without the chat history. Do NOT answer the question, \
+just reformulate it if needed and otherwise return it as is."""
+contextualize_q_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", contextualize_q_system_prompt),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{question}"),
+    ]
+)
+contextualize_q_chain = lambda llm: contextualize_q_prompt | llm | StrOutputParser()
+
+qa_system_prompt = """As Zeta, your mission is to assist users in navigating the vast sea of machine learning research with ease and insight. When responding to inquiries, adhere to the following guidelines to ensure the utmost accuracy and utility:
+
+Contextual Understanding: When presented with a question, apply your understanding of machine learning concepts to interpret the context provided accurately. Utilize this context to guide your search for answers within the specified research papers.
+
+Answer Provision: Always provide an answer that is directly supported by the research papers' content. If the information needed to answer the question is not available, clearly state, "I don't know."
+
+Citation Requirement: For every answer given, include multiple citations from the research papers. A citation must include a direct quote from the paper that supports your answer, along with the identification (ID) of the paper. This ensures that all provided information can be traced back to its source, maintaining a high level of credibility and transparency.
+
+Formatting Guidelines: Present your citations in the following structured format at the end of your answer to maintain clarity and consistency:
+
+
+<citations>
+    <citation><source_id>[Source ID]</source_id><quote>[Direct quote from the source]</quote></citation>
+    ...
+</citations>
+
+
+Conflict Resolution: In cases where multiple sources offer conflicting information, evaluate the context, relevance, and credibility of each source to determine the most accurate answer. Explain your reasoning within the citation section to provide insight into your decision-making process.
+
+User Engagement: Encourage user engagement by asking clarifying questions if the initial inquiry is ambiguous or lacks specific context. This helps in providing more targeted and relevant responses.
+
+Continual Learning: Although you are not expected to generate new text or insights beyond the provided papers, be open to learning from new information as it becomes available to you through user interactions and queries.
+
+By following these guidelines, you ensure that users receive valuable, accurate, and source-backed insights into their inquiries, making their exploration of machine learning research more productive and enlightening.
+
+{context}"""
+qa_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", qa_system_prompt),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{question}"),
+    ]
+)
+
+
+def format_docs(docs):
+    return "\n\n".join(
+        f"{doc.metadata['chunk_id']}: {doc.page_content}" for doc in docs
+    )
+
+
+def contextualized_question(input: dict):
+    if input.get("chat_history"):
+        return contextualize_q_chain
+    else:
+        return input["question"]
+
+
+rag_chain = lambda retriever, llm: (
+    RunnablePassthrough.assign(
+        context=contextualized_question | retriever | format_docs
+    )
+    | qa_prompt
+    | llm
+)
+
+
+def parse_model_response(input_string):
+    parsed_data = {"answer": "", "citations": []}
+    xml_matches = re.findall(r"<citations>.*?</citations>", input_string, re.DOTALL)
+    if not xml_matches:
+        parsed_data["answer"] = input_string
+        return parsed_data
+
+    outside_text_parts = []
+    last_end_pos = 0
+
+    for xml_string in xml_matches:
+        match = re.search(re.escape(xml_string), input_string[last_end_pos:], re.DOTALL)
+
+        if match:
+            outside_text_parts.append(
+                input_string[last_end_pos : match.start() + last_end_pos]
+            )
+            last_end_pos += match.end()
+
+        root = ET.fromstring(xml_string)
+
+        for citation in root.findall("citation"):
+            source_id = citation.find("source_id").text
+            quote = citation.find("quote").text
+            parsed_data["citations"].append({"source_id": source_id, "quote": quote})
+
+    outside_text_parts.append(input_string[last_end_pos:])
+
+    parsed_data["answer"] = "".join(outside_text_parts)
+
+    return parsed_data