Update app.py

app.py

@@ -7,46 +7,24 @@ from langchain_community.document_loaders import UnstructuredURLLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
-from langchain_core.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
+from langchain_google_genai import ChatGoogleGenerativeAI
 
-# Get HuggingFace API token from environment variables
-token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 
+# Get HuggingFace API token from environment variables
+token = os.environ.get("API_TOKEN")
 
 # ------------------------
-# LLM Model (LLaMA 2)
+# LLM 
 # ------------------------
-pipe = pipeline(
-    task="text-generation",
-    model="HuggingFaceTB/SmolLM2-360M",
-    temperature=0.7,
-    max_new_tokens=512,
-    token=token,
-    device_map="auto"
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.5-flash",
+    temperature=0.7,api_key = token
 )
 
-llm = HuggingFacePipeline(
-    pipeline=pipe,
-    model_kwargs={"temperature": 0.7}
-)
-
-# ------------------------
-# Prompt template
-# ------------------------
-prompt = PromptTemplate.from_template(
-    """Given the following extracted parts of a long document and a question, create a final answer with references.
-If you don't know the answer, just say that you don't know.
-Question: {question}"""
-)
 
 # Global variable to store the QA chain
-simple_chain = None
+chain  = None
 
-# ------------------------
-# Function to process URLs with real-time logging
-# ------------------------
-# ------------------------
 # Paths to save FAISS and URLs
 # ------------------------
 FAISS_FILE = "vectorstore.pkl"
@@ -56,7 +34,7 @@ URLS_FILE = "urls.pkl"
 # Function to process URLs with logging and FAISS management
 # ------------------------
 def process_urls_with_logs(url1, url2, url3):
-    global simple_chain
+    global chain 
     
     urls = [url1, url2, url3]
     urls = [u.strip() for u in urls if u.strip() != ""]
@@ -101,9 +79,9 @@ def process_urls_with_logs(url1, url2, url3):
             pickle.dump(urls, f)
 
         print("Initializing LLM chain...")
-        retriever = vectorstore.as_retriever()
-        from langchain_core.runnables import RunnableSequence
-        simple_chain = RunnableSequence(prompt, llm, StrOutputParser())
+        
+        chain = RetrievalQAWithSourcesChain.from_llm( llm=llm, retriever=vectorstore.as_retriever())
+        
         return "FAISS successfully created/recreated!"
     
     else:
@@ -112,23 +90,24 @@ def process_urls_with_logs(url1, url2, url3):
         with open(FAISS_FILE, "rb") as f:
             vectorstore = pickle.load(f)
 
-        retriever = vectorstore.as_retriever()
-        from langchain_core.runnables import RunnableSequence
-        simple_chain = RunnableSequence(prompt, llm, StrOutputParser())
-
+        chain = RetrievalQAWithSourcesChain.from_llm( llm=llm, retriever=vectorstore.as_retriever())
         return "Existing FAISS loaded."
 
 # ------------------------
 # Function to answer questions
 # ------------------------
 def ask_question(question):
-    global simple_chain
+    global chain 
 
-    if simple_chain is None:
+    if chain  is None:
         return "Please process URLs first."
     
-    result = simple_chain.invoke({"question": question})
-    return result
+    result = chain.invoke({'question': question})
+
+    answer = result.get("answer", "")
+    sources = result.get("sources", "")
+        
+    return answer, sources
 
 # ------------------------
 # Gradio Interface
@@ -158,6 +137,7 @@ with gr.Blocks() as app:
             
             ask_btn = gr.Button("Ask")
             answer_output = gr.Textbox(label="Answer", lines=8)
+            sources_output = gr.Textbox(label="Sources", lines=4)
     
     # Connect buttons to suas funções
     process_btn.click(
@@ -167,10 +147,10 @@ with gr.Blocks() as app:
     )
     
     ask_btn.click(
-        ask_question,
-        inputs=question_box,
-        outputs=answer_output
-    )
+    ask_question,
+    inputs=question_box,
+    outputs=[answer_output, sources_output]
+)
 
 # Launch the Gradio app
 app.launch()