Update app.py

app.py

@@ -19,13 +19,17 @@ from langchain.schema import BaseOutputParser, OutputParserException
 from htmlTemplates import css, bot_template, user_template
 from langchain.llms import HuggingFaceHub
 
-class ReferenceOutputParser(BaseOutputParser):
-    def parse(self, text: str) -> dict:
+class ReferenceOutputParser(BaseOutputParser[ChatGeneration]):
+    def parse(self, text: str) -> ChatGeneration:
         try:
             result, references = text.split("References:")
-            return {"result": result.strip(), "references": [ref.strip() for ref in references.split("\n") if ref.strip()]}
+            return ChatGeneration(
+                result=result.strip(),
+                extra_info={"references": [ref.strip() for ref in references.split("\n") if ref.strip()]}
+            )
         except ValueError:
             raise OutputParserException(f"Could not parse output: {text}")
+    
 
 def get_pdf_text(pdf_docs):
     text = ""
@@ -91,8 +95,8 @@ def handle_userinput(user_question):
         response = st.session_state.conversation({"question": user_question})
         st.session_state.chat_history = response["chat_history"]
 
-        result = response["result"]
-        references = response["references"]
+        result = response.result
+        references = response.extra_info["references"]
 
         st.write("//_^ User: " + user_question)
         st.write("🤖 ChatBot: " + result)