File size: 3,049 Bytes
1925b26
 
bd75a92
1925b26
 
 
 
bf9bcf9
 
9bacaee
186c2bc
d31a43f
1925b26
d31a43f
bf9bcf9
 
 
 
 
d31a43f
bf9bcf9
186c2bc
 
 
 
 
1925b26
 
 
 
d31a43f
1925b26
d31a43f
 
 
 
e3f3e21
9d57136
d31a43f
 
 
 
bf9bcf9
d31a43f
 
 
1925b26
 
d31a43f
1925b26
 
 
 
 
 
d31a43f
 
 
 
 
 
 
 
1925b26
d31a43f
1925b26
e3f3e21
1925b26
d31a43f
 
 
 
1925b26
 
 
 
 
d31a43f
1925b26
d31a43f
1925b26
d31a43f
 
 
 
 
 
 
186c2bc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableParallel
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.chat_models import ChatLiteLLM
from langchain_core.messages import HumanMessage, AIMessage

class ProjectRAGEngine:
    def __init__(self):
        # ✅ Hugging Face Embeddings (LOCAL / FREE)
        self.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={"device": "cpu"},   # change to "cuda" if GPU available
            encode_kwargs={"normalize_embeddings": True}
        )
        # ✅ OpenRouter LLM (Chat only)
        self.llm = ChatOpenAI(
                model="openai/gpt-oss-120b:free",
                base_url="https://openrouter.ai/api/v1",
                api_key=os.getenv("OPENROUTER_API_KEY"),
                extra_body={"reasoning": {"enabled": True}})
        self.vector_store = None

    def process_documents(self, pdf_paths):
        all_docs = []

        for path in pdf_paths:
            loader = PyPDFLoader(path)
            all_docs.extend(loader.load())

        splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50
        )

        splits = splitter.split_documents(all_docs)

        # ✅ FAISS with HuggingFace embeddings
        self.vector_store = FAISS.from_documents(
            splits, self.embeddings
        )

    def _format_docs(self, 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.", []

        template = """
        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}
        Question:
        {question}
        """

        prompt = ChatPromptTemplate.from_template(template)
        retriever = self.vector_store.as_retriever(search_type="mmr", search_kwargs={"k": 5, "lambda_mult":0.25})

        rag_chain = (
            RunnablePassthrough.assign(
                context=lambda x: self._format_docs(x["context"])
            )
            | prompt
            | self.llm
            | StrOutputParser()
        )

        chain = RunnableParallel(
            {"context": retriever, "question": RunnablePassthrough()}
        ).assign(answer=rag_chain)

        result = chain.invoke(query)

        sources = [
            {"content": d.page_content, "metadata": d.metadata}
            for d in result["context"]
        ]

        return result["answer"], sources