Rename app.py to rag.py

rag.py

@@ -0,0 +1,61 @@
+# rag.py
+
+from langchain.vectorstores import Chroma
+from langchain.chat_models import ChatOllama
+from langchain.embeddings import FastEmbedEmbeddings
+from langchain.schema.output_parser import StrOutputParser
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.schema.runnable import RunnablePassthrough
+from langchain.prompts import PromptTemplate
+from langchain.vectorstores.utils import filter_complex_metadata
+
+
+class ChatPDF:
+    vector_store = None
+    retriever = None
+    chain = None
+
+    def __init__(self):
+        self.model = ChatOllama(model="mistral")
+        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=100)
+        self.prompt = PromptTemplate.from_template(
+            """
+            <s> [INST] You are an assistant for question-answering tasks. Use the following pieces of retrieved context 
+            to answer the question. If you don't know the answer, just say that you don't know. Use three sentences
+             maximum and keep the answer concise. [/INST] </s> 
+            [INST] Question: {question} 
+            Context: {context} 
+            Answer: [/INST]
+            """
+        )
+
+    def ingest(self, pdf_file_path: str):
+        docs = PyPDFLoader(file_path=pdf_file_path).load()
+        chunks = self.text_splitter.split_documents(docs)
+        chunks = filter_complex_metadata(chunks)
+
+        vector_store = Chroma.from_documents(documents=chunks, embedding=FastEmbedEmbeddings())
+        self.retriever = vector_store.as_retriever(
+            search_type="similarity_score_threshold",
+            search_kwargs={
+                "k": 3,
+                "score_threshold": 0.5,
+            },
+        )
+
+        self.chain = ({"context": self.retriever, "question": RunnablePassthrough()}
+                      | self.prompt
+                      | self.model
+                      | StrOutputParser())
+
+    def ask(self, query: str):
+        if not self.chain:
+            return "Please, add a PDF document first."
+
+        return self.chain.invoke(query)
+
+    def clear(self):
+        self.vector_store = None
+        self.retriever = None
+        self.chain = None