Modified CustomChain, code-linting

agent.py

@@ -28,30 +28,6 @@ llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", GEMINI_API_KEY=GEMINI_API
 
 
 print("Models initialized successfully.")
-def load_documents(directory):
-    loader = DirectoryLoader(
-        directory,
-        glob="**/*.txt",
-        loader_cls=TextLoader
-    )
-    documents = loader.load()
-
-    docx_loader = DirectoryLoader(
-        directory,
-        glob="**/*.docx",
-        loader_cls=UnstructuredWordDocumentLoader,
-        loader_kwargs={"mode": "elements"}
-    )
-    documents.extend(docx_loader.load())
-    print(f"Loaded {len(documents)} documents.")
-    pdf_loader = DirectoryLoader(
-        directory,
-        glob="**/*.pdf",
-        loader_cls=UnstructuredPDFLoader
-    )
-    documents.extend(pdf_loader.load())
-    print(f"Loaded {len(documents)} documents.")
-    return documents
 
 import os
 from dotenv import load_dotenv
@@ -200,7 +176,7 @@ def create_health_agent(vector_store):
 
         @property
         def input_keys(self):
-            return ['query', 'previous_conversation']
+            return ['query', 'previous_conversation', 'user_data']
 
         @property
         def output_keys(self):
@@ -209,6 +185,7 @@ def create_health_agent(vector_store):
         def _call(self, inputs):
             query = inputs['query']
             previous_conversation = inputs.get('previous_conversation', '')
+            user_data = inputs.get('user_data', '')
             
             # Retrieve relevant documents
             docs = retriever.get_relevant_documents(query)
@@ -218,7 +195,8 @@ def create_health_agent(vector_store):
             llm_inputs = {
                 'context': context,
                 'question': query,
-                'previous_conversation': previous_conversation
+                'previous_conversation': previous_conversation,
+                'user_data': user_data
             }
 
             # Generate response
@@ -229,7 +207,7 @@ def create_health_agent(vector_store):
     llm_chain = LLMChain(llm=llm, prompt=PROMPT)
 
     # Create and return the custom chain
-    return CustomRetrievalQA(retriever=retriever, llm_chain=llm_chain)
+    return CustomRetrievalQA(retriever=retriever, llm_chain=llm_chain, user_data=None)
 
 
 def agent_with_db():