Update app.py

app.py

@@ -1,5 +1,3 @@
-from dataclasses import dataclass
-from operator import itemgetter
 from pathlib import Path
 from typing import List, Optional, Dict, Any
 import logging
@@ -8,17 +6,13 @@ from enum import Enum
 import gradio as gr
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
-from langchain.schema import BaseRetriever
 from langchain.embeddings.base import Embeddings
-from langchain.llms.base import BaseLanguageModel
 import PyPDF2
 from huggingface_hub import InferenceClient
-# Install required packages
-
-
-# Initialize models
 import torch
 from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+# Install required packages
+
 embed_model = HuggingFaceBgeEmbeddings(
     model_name="all-MiniLM-L6-v2",#"dunzhang/stella_en_1.5B_v5",
     model_kwargs={'device': 'cpu'},
@@ -39,6 +33,7 @@ class DocumentFormat(Enum):
     PDF = ".pdf"
     # Can be extended for other document types
 
+
 @dataclass
 class RAGConfig:
     """Configuration for RAG system parameters"""
@@ -47,15 +42,14 @@ class RAGConfig:
     retriever_k: int = 3
     persist_directory: str = "./chroma_db"
 
+
 class AdvancedRAGSystem:
     """Advanced RAG System with improved error handling and type safety"""
-    
-
     def __init__(
         self,
-        embed_model: Embeddings,
-        llm: BaseLanguageModel,
-        config: Optional[RAGConfig] = None
+        embed_model,
+        llm,
+        config = None
     ):
         """Initialize the RAG system with required models and optional configuration"""
         self.embed_model = embed_model
@@ -166,19 +160,12 @@ Context:
 	}
 ]
             
-            response_text = ""
             return self.llm.chat.completions.create(
                 model=model_name,
                 messages=messages,
                 max_tokens=500,
                 # stream=True
             ).choices[0].message.content
-                # return stream.choices[0].message.content
-                # if hasattr(chunk.choices[0].delta, 'content'):
-                #     content = chunk.choices[0].delta.content
-                #     if content is not None:
-                #         response_text += content
-                #         yield response_text
 
         except Exception as e:
             error_msg = f"Error during query processing: {str(e)}"
@@ -186,7 +173,9 @@ Context:
             return error_msg
 
 
-            
+
+
+
 def create_gradio_interface(rag_system: AdvancedRAGSystem) :
     """Create an improved Gradio interface for the RAG system"""
     
@@ -207,8 +196,6 @@ def create_gradio_interface(rag_system: AdvancedRAGSystem) :
     def query_streaming(question: str) :
         try:
             return rag_system.query(question)
-            # for response in rag_system.query(question):
-            #     yield response
         except Exception as e:
             return f"Error: {str(e)}"