Added answer options

app.py

@@ -6,26 +6,28 @@ load_dotenv()
 client = AcademicClient()
 
 
-def perform_qa(query):
-    return client.answer(query)
+def perform_qa(query: str, options: str) -> str:
+    return client.answer(query, options.split('\n'))
 
 
 css = """
 body {
-    align-items: center;
+    image-align: center;
     display:block;
 }
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown('Wisdom.AI'),
+    gr.Markdown('# Wisdom.AI'),
     gr.Image('misc/wisdom.jpg', height=600, width=400)
     with gr.Row():
-        inp = gr.Textbox('Що б ви хотіли дізнатися у мудрого?', label='Питання')
-        out = gr.Textbox('Мудрий каже...', label='Відповідь')
+        inp = gr.Textbox('Що б ви хотіли дізнатися у мудрого?', label='Питання', min_width=400)
+        out = gr.Textbox('Мудрий каже...', label='Відповідь', min_width=400)
+
+    options = gr.Textbox(label='Варіанти відповіді:', min_width=800)
 
     btn = gr.Button('Спитати')
-    btn.click(fn=perform_qa, inputs=inp, outputs=out)
+    btn.click(fn=perform_qa, inputs=[inp, options], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()

src/chains.py

@@ -1,31 +1,71 @@
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough, RunnableLambda
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableParallel
 
 from langchain_openai import ChatOpenAI
 from langchain.prompts import PromptTemplate
 from langchain_community.utilities import GoogleSerperAPIWrapper
+from typing import List
+
+from operator import itemgetter
 
 CUSTOM_RAG_PROMPT = """
-Використай наступні **надійні** елементи, для того, щоб відповісти на питання в кінці. 
-Якщо вони не містять відповіді, зверни увагу на відповідь з інтернету, хоча вона може бути не надійною. 
+Використай наступні **надійні** елементи, для того, щоб вибрати відповідь на питання з запропонованих.
+Якщо вони не містять відповіді, зверни увагу на інформацію з інтернету. 
 Якщо ти не знаєш відповіді, використаши всі свої джерела, то просто скажи про це, не потрібно вигадувати відповідь.
-Використовуй не більше трьох речень, та намагайся відповісти коротко та чітко.
+Напиши у відповіді номер правильного варіанту відповіді. Якщо серед варантів немає правильної відповіді, напиши коротко відповідь самостійно.
+
 
 {context}
 
 Відповідь з інтернету: {internet}
 
+Приклад: 
+
+Чия типологія поділяється на традиційні, харизматичні й раціональні системи?
+1) Вебер 
+2) Ленін
+3) Сталін
+4) Обама
+
+Правильна відповідь: 1 - Вебер. 
+
+
 Питання: {question}
 
-Корисна відповідь:"""
+Варіанти відповіді: 
+{options}
+
+Правильна відповідь:"""
 
 CUSTOM_RAG_PROMPT = PromptTemplate.from_template(CUSTOM_RAG_PROMPT)
 
+VERIFICATION_PROMPT = """
+Вам було задано наступне питання: 
+{question}
+З варіантами відповіді: 
+{options}
+На яку було запропоновано відповідь:
+{answer}
+
+Повторіть відповідь, якщо вона правильна. Інакше, скажіть "відповідь відсутня".
+
+Відповідь:
+"""
+VERIFICATION_PROMPT = PromptTemplate.from_template(VERIFICATION_PROMPT)
+
 
 def documents_parser(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
 
+def prepare_options(options):
+    return "\n".join([f"{i + 1}) {option}" for i, option in enumerate(options)])
+
+def flatten_input(d):
+    ret = d.pop('input')
+    ret.update(d)
+    return ret
+
 class PdfAndGoogleChain:
 
     def use_google_search(self, query):
@@ -34,17 +74,43 @@ class PdfAndGoogleChain:
         except Exception as ex:
             return 'NONE'
 
+    def retrieve_multiple(self, query_dict):
+        query = query_dict['query']
+        options = query_dict['options']
+        ret = self.retriever.get_relevant_documents(query)
+        for option in options:
+            ret.extend(self.retriever.get_relevant_documents(option)[:2])
+        return ret
+
     def __init__(self, retriever, llm_name: str = "gpt-3.5-turbo-0125"):
         self.search = GoogleSerperAPIWrapper()
+        self.retriever = retriever
         self.llm = ChatOpenAI(model=llm_name)
 
         self.rag_chain = (
-                {"context": retriever | documents_parser, "internet": RunnableLambda(self.use_google_search),
-                 "question": RunnablePassthrough()}
+                {"context": RunnableLambda(self.retrieve_multiple) | documents_parser,
+                 "internet": itemgetter("query") | RunnableLambda(self.use_google_search),
+                 "question": itemgetter("query") | RunnablePassthrough(),
+                 "options": itemgetter("options") | RunnableLambda(prepare_options)}
                 | CUSTOM_RAG_PROMPT
                 | self.llm
                 | StrOutputParser()
         )
 
-    def answer(self, query: str):
-        return self.rag_chain.invoke(query)
+        self.verification_chain = (
+                {"question": itemgetter("query") | RunnablePassthrough(),
+                 "options": itemgetter("options") | RunnableLambda(prepare_options),
+                 "answer": itemgetter("answer") | RunnablePassthrough()}
+                | VERIFICATION_PROMPT
+                | self.llm
+                | StrOutputParser()
+        )
+
+        self.global_chain = (RunnableParallel(input=RunnablePassthrough(), answer=self.rag_chain)
+                             | RunnableLambda(flatten_input)
+                             | self.verification_chain)
+
+    def answer(self, query: str, options: List[str]):
+        options = list(filter(lambda x: x is not None and len(x) > 0, options))
+        return self.global_chain.invoke({"query": query, "options": options})
+

src/clients.py

@@ -43,7 +43,8 @@ class AcademicClient:
 
     def __init__(self):
         self.create_vectordb()
-        self.chain = PdfAndGoogleChain(self.vectordb.as_retriever())
+        self.chain = PdfAndGoogleChain(
+            self.vectordb.as_retriever(search_type="mmr", search_kwargs={"fetch_k": 30, "k": 6}))
 
-    def answer(self, query):
-        return self.chain.answer(query)
+    def answer(self, query, options):
+        return self.chain.answer(query, options)

src/complex.ipynb

@@ -11,7 +11,7 @@
     "from langchain.vectorstores import Chroma\n",
     "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
     "from langchain_core.output_parsers import StrOutputParser\n",
-    "from langchain_core.runnables import RunnablePassthrough, RunnableLambda\n",
+    "from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableParallel\n",
     "from langchain.document_loaders import PyPDFLoader\n",
     "from langchain_openai import ChatOpenAI\n",
     "from dotenv import load_dotenv\n",
@@ -149,7 +149,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 26,
    "outputs": [],
    "source": [
     "retriever = vectordb.as_retriever()\n",
@@ -174,12 +174,37 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-04-09T13:38:45.242263100Z",
-     "start_time": "2024-04-09T13:38:45.221315700Z"
+     "end_time": "2024-04-09T18:59:04.620561600Z",
+     "start_time": "2024-04-09T18:59:04.602246900Z"
     }
    },
    "id": "64cb22281c854513"
   },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "langchain_core.runnables.base.RunnableSequence"
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "type(rag_chain)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-09T18:59:09.631586300Z",
+     "start_time": "2024-04-09T18:59:09.610952200Z"
+    }
+   },
+   "id": "c2fe5487662fc6f0"
+  },
   {
    "cell_type": "code",
    "execution_count": 19,
@@ -263,6 +288,67 @@
    },
    "id": "b0a54b5e476b46e0"
   },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "outputs": [],
+   "source": [
+    "from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableParallel\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-09T19:11:51.492629300Z",
+     "start_time": "2024-04-09T19:11:51.470103600Z"
+    }
+   },
+   "id": "5c977fcc519c1a6e"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "outputs": [],
+   "source": [
+    "def flatten_input(d):\n",
+    "    ret = d.pop('a')\n",
+    "    ret.update(d)\n",
+    "    return ret\n",
+    "a = RunnableParallel(a = RunnablePassthrough(), b = RunnableLambda(lambda x: \"abracadabra\")) | RunnableLambda(flatten_input)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-09T19:14:33.339534100Z",
+     "start_time": "2024-04-09T19:14:33.319287200Z"
+    }
+   },
+   "id": "720b1320bc0fb7d8"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "{'xx': 'yy', 'zz': 11, 'b': 'abracadabra'}"
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "a.invoke({\"xx\" : \"yy\", \"zz\" : 11})"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-04-09T19:14:33.976773500Z",
+     "start_time": "2024-04-09T19:14:33.918734500Z"
+    }
+   },
+   "id": "465e8521af889ff6"
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -271,7 +357,7 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "5c977fcc519c1a6e"
+   "id": "27c9bf7387f058cf"
   }
  ],
  "metadata": {