Rename agent.py to knowledge_base.py

agent.py

@@ -1,26 +0,0 @@
-from meal_loader import documents
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain_community.llms import HuggingFaceHub
-from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
-
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-db = FAISS.from_documents(documents, embeddings)
-retriever = db.as_retriever(search_kwargs={"k": 3})
-llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.1", model_kwargs={"temperature": 0.3, "max_new_tokens": 500})
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-qa_chain = ConversationalRetrievalChain.from_llm(llm=llm, retriever=retriever, memory=memory)
-
-def generate_response(message, history, preferences):
-    prompt = f"""
-You are a meal plan assistant. The user has the following preferences:
-- Diet: {', '.join(preferences['diet'])}
-- Goal: {preferences['goal']}
-- Duration: {preferences['weeks']} week(s)
-
-User query: {message}
-"""
-    result = qa_chain({"question": prompt})
-    return result["answer"]

knowledge_base.py

@@ -0,0 +1,36 @@
+# knowledge_base.py
+import os
+import fitz  # PyMuPDF
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.docstore.document import Document
+
+CHROMA_DIR = "chroma"
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+
+
+def load_and_chunk_pdfs(folder_path):
+    documents = []
+    for filename in os.listdir(folder_path):
+        if filename.endswith(".pdf"):
+            path = os.path.join(folder_path, filename)
+            doc = fitz.open(path)
+            text = "\n".join(page.get_text() for page in doc)
+            documents.append(Document(page_content=text, metadata={"source": filename}))
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_documents(documents)
+    return chunks
+
+
+def create_vectorstore(chunks):
+    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
+    db = Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_DIR)
+    db.persist()
+    return db
+
+
+def load_vectorstore():
+    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
+    return Chroma(persist_directory=CHROMA_DIR, embedding_function=embeddings)