From 8211705debf9d1335223c606275f46c43c78d8a2 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Mon, 8 Apr 2024 00:23:52 +0200
Subject: Updates

---
 rag/ui.py | 39 +++++++++++++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 4 deletions(-)

(limited to 'rag/ui.py')

diff --git a/rag/ui.py b/rag/ui.py
index 1e4dd64..277c084 100644
--- a/rag/ui.py
+++ b/rag/ui.py
@@ -1,8 +1,14 @@
 import streamlit as st
 
+from langchain_community.document_loaders.blob_loaders import Blob
+
+try:
+    from rag.rag import RAG
+except ModuleNotFoundError:
+    from rag import RAG
+
+rag = RAG()
 
-# from loguru import logger as log
-# from rag.rag import RAG
 
 def upload_pdfs():
     files = st.file_uploader(
@@ -11,10 +17,35 @@ def upload_pdfs():
         accept_multiple_files=True,
     )
     for file in files:
-        bytes = file.read()
-        st.write(bytes)
+        blob = Blob.from_data(file.read())
+        rag.add_pdf_from_blob(blob)
 
 
 if __name__ == "__main__":
     st.header("RAG-UI")
+
     upload_pdfs()
+    query = st.text_area(
+        "query",
+        key="query",
+        height=100,
+        placeholder="Enter query here",
+        help="",
+        label_visibility="collapsed",
+        disabled=False,
+    )
+
+    (result_column, context_column) = st.columns(2)
+
+    if query:
+        response = rag.retrive(query)
+
+        with result_column:
+            st.markdown("### Answer")
+            st.markdown(response.answer)
+
+        with context_column:
+            st.markdown("### Context")
+            for c in response.context:
+                st.markdown(c)
+                st.markdown("---")
-- 
cgit v1.2.3-70-g09d2