Spaces:
Running
Running
| import os | |
| from dotenv import load_dotenv | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_mistralai.chat_models import ChatMistralAI | |
| from langchain_mistralai.embeddings import MistralAIEmbeddings | |
| from langchain.schema.output_parser import StrOutputParser | |
| from langchain_community.document_loaders import PyPDFLoader | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain.schema.runnable import RunnablePassthrough | |
| from langchain.prompts import PromptTemplate | |
| from langchain_community.vectorstores.utils import filter_complex_metadata | |
| #add new import | |
| from langchain_community.document_loaders.csv_loader import CSVLoader | |
| from prompt_template import base_template | |
| # load .env in local dev | |
| load_dotenv() | |
| env_api_key = os.environ.get("MISTRAL_API_KEY") | |
| llm_model = "open-mixtral-8x7b" | |
| class Rag: | |
| document_vector_store = None | |
| retriever = None | |
| chain = None | |
| def __init__(self, vectore_store=None): | |
| self.model = ChatMistralAI(model=llm_model) | |
| self.embedding = MistralAIEmbeddings(model="mistral-embed", mistral_api_key=env_api_key) | |
| self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100, length_function=len) | |
| self.prompt = PromptTemplate.from_template(base_template) | |
| self.vector_store = vectore_store | |
| def setModel(self, model): | |
| self.model = model | |
| def ingestToDb(self, file_path: str, filename: str): | |
| docs = PyPDFLoader(file_path=file_path).load() | |
| # Extract all text from the document | |
| text = "" | |
| for page in docs: | |
| text += page.page_content | |
| # Split the text into chunks | |
| chunks = self.text_splitter.split_text(text) | |
| return self.vector_store.addDoc(filename=filename, text_chunks=chunks, embedding=self.embedding) | |
| def getDbFiles(self): | |
| return self.vector_store.getDocs() | |
| 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) | |
| document_vector_store = FAISS.from_documents(chunks, self.embedding) | |
| self.retriever = document_vector_store.as_retriever( | |
| search_type="similarity_score_threshold", | |
| search_kwargs={ | |
| "k": 3, | |
| "score_threshold": 0.5, | |
| }, | |
| ) | |
| self.chain = self.prompt | self.model | StrOutputParser() | |
| def ask(self, query: str, messages: list): | |
| if not self.chain: | |
| return "Ajouter un document PDF d'abord." | |
| print("messages ", messages) | |
| # Retrieve the context document | |
| documentContext = self.retriever.invoke(query) | |
| # Retrieve the VectoreStore | |
| contextCommon = None | |
| return self.chain.invoke({ | |
| "query": query, | |
| "documentContext": documentContext, | |
| "commonContext": contextCommon, | |
| "messages": messages | |
| }) | |
| def clear(self): | |
| self.document_vector_store = None | |
| self.vector_store = None | |
| self.retriever = None | |
| self.chain = None |