Update src/rag_engine.py

src/rag_engine.py

@@ -7,59 +7,87 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough, RunnableParallel
 
+
 class ProjectRAGEngine:
-    def __init__(self, api_key):
-        self.embeddings = OpenAIEmbeddings(openai_api_key=api_key)
-        self.llm = ChatOpenAI(model="openai/gpt-oss-120b:free", openai_api_key=api_key, temperature=0)
+    def __init__(self):
+        # ✅ OpenAI embeddings (OFFICIAL)
+        self.embeddings = OpenAIEmbeddings(
+            model="text-embedding-3-small"
+        )
+
+        # ✅ OpenRouter LLM
+        self.llm = ChatOpenAI(
+            model="openai/gpt-oss-120b:free",
+            temperature=0,
+            openai_api_key=os.getenv("OPENROUTER_API_KEY"),
+            openai_api_base="https://openrouter.ai/api/v1",
+            default_headers={
+                "HTTP-Referer": "http://localhost:8501",
+                "X-Title": "Project-RAG-App"
+            }
+        )
+
         self.vector_store = None
 
     def process_documents(self, pdf_paths):
         all_docs = []
+
         for path in pdf_paths:
-            try:
-                loader = PyPDFLoader(path)
-                docs = loader.load()
-                all_docs.extend(docs)
-            except Exception as e:
-                print(f"Error loading {path}: {e}")
-        
-        # Splitting logic to handle large reports [cite: 10]
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(all_docs)
-        self.vector_store = FAISS.from_documents(splits, self.embeddings)
+            loader = PyPDFLoader(path)
+            all_docs.extend(loader.load())
+
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200
+        )
+
+        splits = splitter.split_documents(all_docs)
+
+        self.vector_store = FAISS.from_documents(
+            splits, self.embeddings
+        )
 
     def _format_docs(self, docs):
-        return "\n\n".join(doc.page_content for doc in docs)
+        return "\n\n".join(d.page_content for d in docs)
 
     def get_answer(self, query):
         if not self.vector_store:
             return "Please upload documents first.", []
 
-        # System prompt ensuring grounded responses [cite: 18, 25]
         template = """
-        You are a professional Project Analyst. Answer strictly based on the provided context.
-        If the answer is not in the context, say you don't know. 
-        Cite document names and page numbers for every answer. Include direct quotes.
+        You are a professional Project Analyst.
+        Answer strictly using the context.
+        If unknown, say you don't know.
+        Cite document names and page numbers.
+        
+        Context:
+        {context}
         
-        Context: {context}
-        Question: {question}
+        Question:
+        {question}
         """
+
         prompt = ChatPromptTemplate.from_template(template)
         retriever = self.vector_store.as_retriever(search_kwargs={"k": 5})
 
-        # Pure LCEL Chain composition
-        rag_chain_from_docs = (
-            RunnablePassthrough.assign(context=(lambda x: self._format_docs(x["context"])))
+        rag_chain = (
+            RunnablePassthrough.assign(
+                context=lambda x: self._format_docs(x["context"])
+            )
             | prompt
             | self.llm
             | StrOutputParser()
         )
 
-        rag_chain_with_source = RunnableParallel(
+        chain = RunnableParallel(
             {"context": retriever, "question": RunnablePassthrough()}
-        ).assign(answer=rag_chain_from_docs)
+        ).assign(answer=rag_chain)
 
-        result = rag_chain_with_source.invoke(query)
-        
-        sources = [{"content": doc.page_content, "metadata": doc.metadata} for doc in result["context"]]
-        return result["answer"], sources
+        result = chain.invoke(query)
+
+        sources = [
+            {"content": d.page_content, "metadata": d.metadata}
+            for d in result["context"]
+        ]
+
+        return result["answer"], sources