code refactoring

app.py

@@ -10,9 +10,10 @@ from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from pypdf import PdfReader, PdfWriter
 from pathlib import Path
+from typing import List
 
 
-def build_rag_chain(pdf_paths):
+def build_rag_chain(pdf_paths: List[str], chunk_size: int =1000, chunk_overlap: int =200, model_name: str ='gpt-4-0125-preview'):
     loaders = [PyPDFLoader(path) for path in pdf_paths]
 
     docs = []
@@ -21,8 +22,6 @@ def build_rag_chain(pdf_paths):
             loader.load()[0:] # skip first page
         )
 
-    chunk_size = 1000
-    chunk_overlap = 200
 
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, 
                                                 chunk_overlap=chunk_overlap)
@@ -35,7 +34,7 @@ def build_rag_chain(pdf_paths):
 
     # model_name = 'gpt-3.5-turbo-0125'
     # model_name = 'gpt-4-1106-preview'
-    model_name = 'gpt-4-0125-preview'
+    # model_name = 'gpt-4-0125-preview'
     llm = ChatOpenAI(model_name=model_name, temperature=0)
 
     def format_docs(docs):
@@ -51,9 +50,10 @@ def build_rag_chain(pdf_paths):
     return rag_chain
 
 
-def predict(query, pdf_file):
-    print(type(pdf_file))
-    if pdf_file:
+def predict(query: str, pdf_id: str =None, user_id: str = None):
+    print(type(pdf_id))
+    # print(user_id)
+    if pdf_id:
         # pdf_path = Path(pdf_file)
         # pdf_reader = PdfReader(pdf_path)
         # pdf_writer = PdfWriter()
@@ -72,7 +72,7 @@ def predict(query, pdf_file):
         # os.system("ls data/pdf")
             
     # pdf_paths = load_pdf_paths(data_root)
-        rag_chain = build_rag_chain([pdf_file])
+        rag_chain = build_rag_chain([pdf_id])
         return rag_chain.invoke(query)
     return "Please upload PDF file"