Update custom_llm.py

custom_llm.py

@@ -270,15 +270,29 @@ def custom_unique_df_chain(llm, df):
 
 def custom_combined_chain(llm, df_chain, memory_chain):
 
-    prompt = PromptTemplate.from_template("""<s><INST> Given the following question, classify it as either being more relevant with a dataframe object of ticket submissions' history or several documents of user guide and general knowledge:
+    # prompt = PromptTemplate.from_template("""<s><INST> Given the following question, classify it as either being more relevant with a dataframe object of ticket submissions' history or several documents of user guide and general knowledge:
 
+    # <question>
+    # {question}
+    # </question>
+
+    # Respond with ONLY one word either "ticket" or "knowledge"
+    
+    # </s></INST>""")
+
+    prompt = PromptTemplate.from_template("""<s><INST> You have access to the following data sources:
+    1. Dataframe : use this data source to retrieve anything about ticket submission history
+    2. Documents : use this data source to retrieve anything related to user guide and work instruction
+    
     <question>
     {question}
     </question>
-
-    Respond with ONLY one word either "ticket" or "knowledge"
     
-    </s></INST>""")
+    
+    Respond with ONLY one word either "dataframe" or "documents"
+    
+    </s></INST>
+    """)
     
     # def route(info):
     #     if 'ticket' in info['topic']:
@@ -289,7 +303,7 @@ def custom_combined_chain(llm, df_chain, memory_chain):
     
     # full_chain = RunnablePassthrough.assign(topic= (prompt | llm)) | RunnableLambda(route)
     # combined_chain = prompt | llm
-    return RunnablePassthrough.assign(topic=prompt | llm) | RunnableBranch( (lambda x: "ticket" in x['topic'].lower(), df_chain), memory_chain )
+    return RunnablePassthrough.assign(topic=prompt | llm) | RunnableBranch( (lambda x: "dataframe" in x['topic'].lower(), df_chain), memory_chain )
 
 
 class CustomLLM(LLM):