Merge pull request #89 from FYP-2024-IQMA/SCRUM-128-Integrate-Agent-t…

…o-LangChain

[SCRUM-150] Scrum 128 integrate agent to lang chain + Scrum 150 user story

backend/__tests__/chat_endpoint_test.py

@@ -17,6 +17,10 @@
 logger.warning("Logging is active.")
 client = TestClient(app)
 
+#######################################################
+# OPENAI API TESTS
+#######################################################
+
 def test_generate_text_no_history():
     # Test without history
     response = client.post("/generate", json={"role": "user", 
@@ -51,10 +55,40 @@ def test_generate_text_invalid_prompt():
     logger.info(f"Response: {response.json()}")
     assert response.status_code == 422
 
+#######################################################
+# LANGCHAIN TESTS
+#######################################################
+
 def test_langchain_endpoint():
+    # Test langchain endpoint
     response = client.post("/langchain", json={"role": "user", 
                                               "content": "Reply with '42'. Do not add any other text. Stop generating after you have replied with '42'."
                                               })
     logger.info(f"Response: {response.json()}")
     assert response.status_code == 200, "Langchain endpoint is faulty."
-    assert "42" in response.json()["content"], "ChatGPT is not generating the correct response."
+    assert "42" in response.json()["content"]["final_answer"]["output"], "Langchain is not generating the correct response."
+
+def test_langchain_endpoint_with_history():
+    # Test with history
+    response = client.post("/langchain", json={"role": "assistant", 
+                                              "content": "Hello, how can I help you today?", 
+                                              "history": [{
+                                                  "role": "user", 
+                                                  "content": "Reply with '42'. Do not add any other text. Stop generating after you have replied with '42'."
+                                                  }]
+                                            })
+    logger.info(f"Response: {response.json()}")
+    assert response.status_code == 200, "Langchain endpoint is faulty."
+    assert "42" in response.json()["content"]["final_answer"]["output"], "Langchain is not generating the correct response."
+
+def test_langchain_endpoint_invalid_role():
+    # Test with missing role
+    response = client.post("/langchain", json={"prompt": "Hello, how can I help you today?"})
+    logger.info(f"Response: {response.json()}")
+    assert response.status_code == 422
+
+def test_langchain_endpoint_invalid_prompt():
+    # Test with missing prompt
+    response = client.post("/langchain", json={"role": "user"})
+    logger.info(f"Response: {response.json()}")
+    assert response.status_code == 422

backend/src/chatbot/README.md

@@ -0,0 +1,15 @@
+## This folder houses all the code for running the chatbot.
+
+# Here are a list of what it contains:
+- app.py : main script for running the FastAPI server to host the endpoints for the mobile app
+- llamap_notebook.ipynb : demo notebook for working with LlamaIdex & ChromaDB
+- streamlit_test.py : Streamlit demo
+
+- setup files:
+    - chatgpt.py : script for interfacing OpenAI API
+    - chroma_setup.py : for loading the vector DB and initializing the agent
+    - Dockerfile : script for starting FastAPI server
+    - langchain_setup.py : script for setting up Agent RAG with branching
+    - ollama_setup.py : for setting up the llama 3.1 8B model
+    - langchain_setup.py : for loading the agent-integrated chain
+

backend/src/chatbot/app.py

@@ -11,8 +11,10 @@
 # dealing with relative / absolute imports
 if __package__ is None or __package__ == '' or __name__ == '__main__':
     from chatgpt import ChatGPT
+    from langchain_setup import full_chain, full_chain_w_history
 else:
     from src.chatbot.chatgpt import ChatGPT
+    from src.chatbot.langchain_setup import full_chain, full_chain_w_history
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
@@ -56,24 +58,41 @@ async def generate_text(prompt: Prompt):
 @app.post("/langchain")
 async def langchain_text(prompt: Prompt):
     """
-    Generate a response from the ChatGPT object based on the role and prompt.
+    Generate a response from Agent-integrated chain based on the role and prompt.
     """
     logger.info("Endpoint '/langchain' has been called with prompt: %s", prompt)
     try:
-        llm = ChatOpenAI(
-            model="gpt-4o-mini",
-            api_key=os.environ.get("OPENAI_API_KEY"),
-        )
+        # llm = ChatOpenAI(
+        #     model="gpt-4o-mini",
+        #     api_key=os.environ.get("OPENAI_API_KEY"),
+        # )
+        config = {
+                "configurable": {
+                    "session_id": "abc123",
+                }
+            }
         if prompt.history:
+            # format response for langchain
             langchain_format = convert_openai_messages(prompt.history)
-            response = llm.invoke(langchain_format)
+            logger.info("Converted langchain messages: %s", langchain_format)
+            input = {
+                "input": prompt.content, 
+                "history": langchain_format,
+            }
+            response = full_chain_w_history.invoke(input, config=config)
         else:
-            # langchain_format = convert_openai_messages([{"role": prompt.role, "content": prompt.content}])
-            response = llm.invoke([{"role": prompt.role, "content": prompt.content}])
+            langchain_format = convert_openai_messages([{"role": prompt.role, "content": prompt.content}])
+            logger.info("Converted langchain messages: %s", langchain_format)
+            input = {
+                "input": prompt.content,
+                "history": langchain_format,
+            }
+            logger.info("Prompt: %s", prompt)
+            response = full_chain_w_history.invoke({"input": input}, config=config)
 
         return {
             "role": "assistant",
-            "content": response.content,
+            "content": response,
             }
     except Exception as e:
         logger.error("Error in '/langchain' endpoint: %s", str(e))

backend/src/chatbot/chroma_setup.py

@@ -20,6 +20,8 @@
 load_dotenv()
 nest_asyncio.apply()
 
+persist_dir = "./chroma_langchain_db"
+
 # attempt to load documents from persistent storage
 try:
     print(f"Attempting to load documents from persistent storage")
@@ -30,7 +32,7 @@
     vectorstore = ChromaVectorStore.from_params(
         collection_name="iqma_collection",
         embedding_function=embeddings,
-        persist_dir="./src/chatbot/chroma_langchain_db"
+        persist_dir=persist_dir,
     )
     print(f"Completed setting up vector store")
 
@@ -46,6 +48,7 @@
 
 # load, process and store documents in persistent storage if not existent
 if not index_loaded:
+    print(f"Loading documents from docs folder")
     # set up parser
     parser = LlamaParse(
         result_type="text",
@@ -58,7 +61,7 @@
     documents = document_reader.load_data()
 
     # set up chromadb
-    db = chromadb.PersistentClient(path="./src/chatbot/chroma_langchain_db")
+    db = chromadb.PersistentClient(path=persist_dir)
     chroma_collection = db.get_or_create_collection("iqma_collection")
 
     # create vector store

backend/src/chatbot/langchain_setup.py

@@ -0,0 +1,183 @@
+# script for creating chains
+
+from dotenv import load_dotenv
+from functools import partial
+from langchain_community.chat_message_histories import ChatMessageHistory
+from langchain_core.chat_history import BaseChatMessageHistory
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain.schema.runnable import RunnableBranch, RunnablePassthrough
+from langchain.schema.runnable.passthrough import RunnableAssign
+from langchain_openai import ChatOpenAI
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableLambda
+import os
+
+
+# dealing with relative / absolute imports
+if __package__ is None or __package__ == '' or __name__ == '__main__':
+    from chroma_setup import agent_executor
+else:
+    from src.chatbot.chroma_setup import agent_executor
+
+# functions for seeing runnables
+def RPrint(preface="State: "):
+    def print_and_return(x, preface=""):
+        print(f"{preface}{x}")
+        return x
+    return RunnableLambda(partial(print_and_return, preface=preface))
+
+def PPrint(preface="State: "):
+    def print_and_return(x, preface=""):
+        print(f"{preface}{x}")
+        return x
+    return RunnableLambda(partial(print_and_return, preface=preface))
+
+# functions for including memory and branching in chain
+def route(info):
+    if "question" in info["choice"].lower():
+        return agent_executor
+    else:
+        return answer_chain
+
+def get_session_history(session_id: str) -> BaseChatMessageHistory:
+    if session_id not in store:
+        store[session_id] = ChatMessageHistory()
+    return store[session_id]
+
+# Create chatbot chain for replying user
+answer_prompt = """<system_prompt>
+YOU ARE AN AWARD-WINNING CAREER COACH AND LEADERSHIP TRAINER, RECOGNIZED BY THE GLOBAL LEADERSHIP ASSOCIATION (2023) AS THE "TOP LEADERSHIP MENTOR OF THE YEAR." YOUR PRIMARY GOAL IS TO HELP USERS DEVELOP AND REFINE LEADERSHIP SKILLS IN THE WORKPLACE, GUIDING THEM THROUGH TRAINING, PRACTICAL EXERCISES, AND REAL-WORLD IMPLEMENTATION OF LEADERSHIP PRINCIPLES.
+
+###INSTRUCTIONS###
+
+- **SUMMARIZE AND ORGANIZE:** AFTER EVERY USER RESPONSE, YOU MUST CONCISELY SUMMARIZE KEY POINTS AND ORGANIZE THEM INTO CLEAN, ACTIONABLE NOTES THAT CAN HELP USERS REVISE THEIR LEADERSHIP TRAINING MATERIAL EFFECTIVELY. THIS IS A TOP PRIORITY. IF THE USER DID NOT PROVIDE MUCH, BUILD OFF WHAT THEY MENTIONED AND SHARE LEADERSHIP SKILL DEVELOPMENT COACHING ADVICE BASED ON THEIR PREVIOUS INPUT.
+  
+- **MENTORING AND GUIDANCE:** PROVIDE MENTORING AND GUIDANCE TO USERS, HELPING THEM IDENTIFY AREAS FOR LEADERSHIP DEVELOPMENT, WHETHER IT BE IN DECISION-MAKING, COMMUNICATION, TEAM MANAGEMENT, OR CONFLICT RESOLUTION.
+
+- **REAL-LIFE APPLICATION:** WHEN USERS ASK QUESTIONS OR SHARE THEIR EXPERIENCES, ENGAGE IN A CONVERSATIONAL MANNER. GUIDE THEM TO APPLY LEADERSHIP PRINCIPLES TO REAL-WORLD SCENARIOS BASED ON THEIR RESPONSES, ENSURING THAT THEY CAN REMEMBER AND IMPLEMENT THE KNOWLEDGE EFFECTIVELY.
+
+- **CHAIN OF THOUGHTS:**
+  1. **LISTEN AND UNDERSTAND:** Carefully analyze the user’s input to understand their current challenge, question, or reflection.
+  2. **SUMMARIZE KEY LEARNINGS:** Summarize the user’s response by highlighting critical points, and organize these into concise notes.
+  3. **PROVIDE ACTIONABLE FEEDBACK:** Offer specific, actionable advice on how to address the challenge or question, connecting it to fundamental leadership principles.
+  4. **REAL-WORLD IMPLEMENTATION:** Encourage users to practice or apply their learning in their current work environment, providing examples and scenarios that align with their input.
+  5. **FOLLOW-UP:** Suggest reflective questions or follow-up tasks to reinforce continuous leadership development and improvement.
+
+###FEW-SHOT EXAMPLES###
+
+**User Input Example 1:**  
+*"I'm struggling with delegating tasks to my team because I worry they won't do it as well as I would."*
+
+**Chatbot Response:**  
+1. **Summarized Key Points:**  
+   - Struggling with delegation due to concerns about task quality.
+   - Prefers to do tasks themselves to maintain high standards.
+   
+2. **Actionable Feedback:**  
+   - Recognize that effective delegation is key to leadership. Focus on training your team members to meet the standards you expect, and set clear guidelines.  
+   - Start by delegating smaller tasks and provide constructive feedback—this will help you build trust in their abilities over time.
+   
+3. **Real-Life Application:**  
+   - Try delegating one small task this week. Provide your team member with clear instructions and expectations, then review the results together. Offer feedback to help them improve.
+
+4. **Follow-Up:**  
+   - How did the delegation process go? What challenges did you face, and what improvements can be made next time?
+
+User input:
+{input}
+"""
+
+answer_prompt_template = ChatPromptTemplate.from_template(answer_prompt)
+answer_llm = ChatOpenAI(
+            model="gpt-4o-mini",
+            api_key=os.environ.get("OPENAI_API_KEY"),
+            temperature=1,
+        )
+answer_chain = answer_prompt_template | answer_llm | StrOutputParser()
+
+# create chain for deciding which chain to use
+branch_prompt = """
+Determine if the user is asking a question or looking for a summary based on the input. Choose the most likely topic based on the user input. Only one word is needed for the choice, no explanation is needed.\n[Available Options: summary, question]
+
+Input:
+{input}
+"""
+branch_prompt_template = ChatPromptTemplate.from_template(branch_prompt)
+branch_llm = ChatOpenAI(
+    model="gpt-4o-mini",
+    api_key=os.environ.get("OPENAI_API_KEY"),
+    temperature=0.2,
+)
+branch_chain = branch_prompt_template | branch_llm | StrOutputParser()
+
+# create chain for consolidating chain outputs and replying user
+reply_prompt = """
+Your job is to consolidate the outputs from the previous chains and reply to the user. If the choice is "question", compile a response using the input and the final answer from the agent. If the choice is "summary", compile a response using the input and the final answer from the answer chain.
+
+History:
+{history}
+
+Input:
+{input}
+
+Choice:
+{choice}
+
+Final Answer:
+{final_answer}
+"""
+
+reply_prompt_template = ChatPromptTemplate.from_template(reply_prompt)
+reply_llm = ChatOpenAI(
+    model="gpt-4o-mini",
+    api_key=os.environ.get("OPENAI_API_KEY"),
+    temperature=1,
+)
+reply_chain = reply_prompt_template | reply_llm | StrOutputParser()
+
+# Chain logic
+# branch_chain decides which chain to use
+# if user ask question, use agent to get a final answer
+# if user input, use chat_chain
+
+full_chain = (
+    PPrint()
+    | RunnableAssign({'choice': branch_chain})
+    | PPrint()
+    | RunnableAssign({'final_answer': route})
+    | PPrint()
+    | RunnableAssign({'output': reply_chain})
+)
+
+# full_chain = (
+#     PPrint()
+#     | RunnableAssign({'choice': branch_chain})
+#     | PPrint()
+#     | RunnableBranch(
+#         (lambda x: "question" in x["choice"].lower(), RunnableAssign({'final_answer': answer_llm})),
+#         # (lambda x: "summary" in x["choice"].lower(), RunnableAssign({'final_answer': answer_chain})),
+#         RunnableAssign({'final_answer': answer_chain})
+#     )
+#     | PPrint()
+#     | RunnableAssign({'output': reply_chain})
+# )
+
+store = {}
+
+full_chain_w_history = RunnableWithMessageHistory(
+    full_chain,
+    get_session_history,
+    input_messages_key="input",
+    history_messages_key="history",
+)
+
+if __name__ == "__main__":
+    # Test the chain
+    input = "Why does a leader need to be adaptable?"
+    # input = "This is not a question. This is a statement. I want you to summarize what I just said. Do not reply with a question."
+    # output = full_chain.invoke({"input": input})
+    output = full_chain_w_history.invoke(
+        {"ability": "math", "input": input},
+        config={"configurable": {"session_id": "abc123"}},
+    )