Update app.py

app.py

@@ -17,6 +17,8 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain import hub
 from langgraph.graph import END, StateGraph, START
 from typing_extensions import TypedDict
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.callbacks import get_openai_callback
 
 # Load environment variables
 dotenv.load_dotenv()
@@ -48,15 +50,28 @@ def initialize_retriever():
         try:
             loader = PyPDFLoader(file_path)
             docs = loader.load()  # Each doc is a page
+
+            # Split each page into smaller chunks
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1024,  # or 512, adjust as needed
+                chunk_overlap=100
+            )
+            split_docs = []
             for doc in docs:
-                doc.metadata["source_file"] = os.path.basename(file_path)
-                doc.metadata["file_path"] = file_path
-                doc.metadata["chunk_size"] = len(doc.page_content)
-                doc.metadata["chunk_id"] = f"{os.path.basename(file_path)}-page-{doc.metadata.get('page', '0')}"
-                if "page" in doc.metadata:
-                    doc.metadata["page_num"] = doc.metadata["page"]
-            logger.info(f"Processed {file_path}: extracted {len(docs)} page-chunks")
-            return docs
+                for chunk in text_splitter.split_text(doc.page_content):
+                    new_doc = doc.__class__(
+                        page_content=chunk,
+                        metadata=doc.metadata.copy()
+                    )
+                    new_doc.metadata["source_file"] = os.path.basename(file_path)
+                    new_doc.metadata["file_path"] = file_path
+                    new_doc.metadata["chunk_size"] = len(chunk)
+                    new_doc.metadata["chunk_id"] = f"{os.path.basename(file_path)}-page-{doc.metadata.get('page', '0')}-chunk"
+                    if "page" in doc.metadata:
+                        new_doc.metadata["page_num"] = doc.metadata["page"]
+                    split_docs.append(new_doc)
+            logger.info(f"Processed {file_path}: extracted {len(split_docs)} chunks")
+            return split_docs
         except Exception as e:
             logger.error(f"Error processing {file_path}: {str(e)}")
             return []
@@ -99,7 +114,7 @@ def initialize_retriever():
     return retriever
 
 # Define graders and components
-def setup_components(retriever):
+def setup_components(retriever, model_choice):
     # Data models for grading
     class GradeDocuments(BaseModel):
         """Binary score for relevance check on retrieved documents."""
@@ -120,7 +135,7 @@ def setup_components(retriever):
         )
     
     # LLM models
-    llm = ChatOpenAI(model="gpt-4.1", temperature=0)
+    llm = ChatOpenAI(model=model_choice, temperature=0)
     doc_grader = llm.with_structured_output(GradeDocuments)
     hallucination_grader_llm = llm.with_structured_output(GradeHallucinations)
     answer_grader_llm = llm.with_structured_output(GradeAnswer)
@@ -304,62 +319,76 @@ def build_rag_graph(components):
     # Compile the graph
     return workflow.compile()
 
+# Initialize global variables
+retriever = None
+rag_app = None
+components = None
+current_model_choice = "gpt-4.1"  # Default
+
+# Run PDF processing and RAG setup ONCE at startup, with default model
+retriever = initialize_retriever()
+if retriever is not None:
+    components = setup_components(retriever, current_model_choice)
+    rag_app = build_rag_graph(components)
+else:
+    logger.error("No retriever could be initialized. Please add PDF files to the Data directory.")
+
 # Processing function for Gradio
-def process_query(question, display_logs=False):
+def process_query(question, display_logs=False, model_choice="gpt-4.1"):
     logs = []
     answer = ""
-    
+    token_usage = {}
     try:
-        # Initialize retriever if not done yet
-        global retriever, rag_app
-        if 'retriever' not in globals() or retriever is None:
-            logs.append("Initializing retriever...")
-            retriever = initialize_retriever()
-            if retriever is None:
-                return "Error: No PDF files found. Please add PDF files to the Data directory.", "\n".join(logs)
-            
-            logs.append("Setting up RAG components...")
-            components = setup_components(retriever)
-            
-            logs.append("Building RAG graph...")
-            rag_app = build_rag_graph(components)
+        global retriever, rag_app, components, current_model_choice
+        if retriever is None:
+            logs.append("Error: No PDF files found. Please add PDF files to the Data directory and restart the app.")
+            return "Error: No PDF files found. Please add PDF files to the Data directory.", "\n".join(logs), token_usage
         
+        # If model_choice changed, re-initialize components and rag_app
+        if model_choice != current_model_choice:
+            logs.append(f"Switching model to {model_choice} ...")
+            components = setup_components(retriever, model_choice)
+            rag_app = build_rag_graph(components)
+            current_model_choice = model_choice
+
         logs.append("Processing query: " + question)
-        
-        # Run the query through the RAG graph
+        logs.append(f"Using model: {model_choice}")
         logs.append("Starting RAG pipeline...")
         final_output = None
-        for i, output in enumerate(rag_app.stream({"question": question})):
-            step_info = f"Step {i+1}: "
+        
+        with get_openai_callback() as cb:
+            for i, output in enumerate(rag_app.stream({"question": question})):
+                step_info = f"Step {i+1}: "
+                if 'retrieve' in output:
+                    step_info += f"Retrieved {len(output['retrieve']['documents'])} documents"
+                elif 'grade_documents' in output:
+                    step_info += f"Graded documents, {len(output['grade_documents']['documents'])} deemed relevant"
+                elif 'transform_query' in output:
+                    step_info += f"Transformed query to: {output['transform_query']['question']}"
+                elif 'generate' in output:
+                    step_info += "Generated answer"
+                    final_output = output
+                logs.append(step_info)
             
-            if 'retrieve' in output:
-                step_info += f"Retrieved {len(output['retrieve']['documents'])} documents"
-            elif 'grade_documents' in output:
-                step_info += f"Graded documents, {len(output['grade_documents']['documents'])} deemed relevant"
-            elif 'transform_query' in output:
-                step_info += f"Transformed query to: {output['transform_query']['question']}"
-            elif 'generate' in output:
-                step_info += "Generated answer"
-                final_output = output
+            # Store token usage information
+            token_usage = {
+                "total_tokens": cb.total_tokens,
+                "prompt_tokens": cb.prompt_tokens,
+                "completion_tokens": cb.completion_tokens,
+                "total_cost": cb.total_cost
+            }
+            logs.append(f"Token usage: {token_usage}")
             
-            logs.append(step_info)
-        
         if final_output and 'generate' in final_output:
             answer = final_output['generate']['generation']
             logs.append("Final answer generated successfully")
         else:
             answer = "No answer could be generated. Please try rephrasing your question."
             logs.append("Failed to generate answer")
-        
     except Exception as e:
         logs.append(f"Error: {str(e)}")
         answer = f"An error occurred: {str(e)}"
-    
-    return answer, "\n".join(logs) if display_logs else ""
-
-# Initialize global variables
-retriever = None
-rag_app = None
+    return answer, "\n".join(logs) if display_logs else "", token_usage
 
 # Create Gradio interface
 with gr.Blocks(title="Self-RAG Document Assistant", theme=gr.themes.Base()) as demo:
@@ -383,6 +412,11 @@ with gr.Blocks(title="Self-RAG Document Assistant", theme=gr.themes.Base()) as d
             )
         
         with gr.Column(scale=1):
+            model_choice_input = gr.Dropdown(
+                label="Model",
+                choices=["gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano"],
+                value="gpt-4.1"
+            )
             show_logs = gr.Checkbox(label="Show Debugging Logs", value=False)
             submit_btn = gr.Button("Submit", variant="primary")
     
@@ -400,12 +434,18 @@ with gr.Blocks(title="Self-RAG Document Assistant", theme=gr.themes.Base()) as d
             lines=15,
             visible=False
         )
+        
+    with gr.Row():
+        token_usage_output = gr.JSON(
+            label="Token Usage Statistics",
+            visible=True
+        )
     
     # Event handlers
     submit_btn.click(
         fn=process_query,
-        inputs=[query_input, show_logs],
-        outputs=[answer_output, logs_output]
+        inputs=[query_input, show_logs, model_choice_input],
+        outputs=[answer_output, logs_output, token_usage_output]
     )
     
     show_logs.change(