Included the alternative retrieval process

my_agent/utils/business_interaction.py

@@ -5,12 +5,15 @@ from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from pydantic import BaseModel, ConfigDict, Field
 from typing import Optional, List
-from .models_loader import llm
-from .prompts import introduction_prompt , business_interaction_prompt
+from .models_loader import llm,ST
+from .prompts import introduction_prompt , business_interaction_prompt, business_retrieval_prompt
 from .tools import retrieve_tool
 from langgraph.prebuilt import create_react_agent
 from langmem.short_term import SummarizationNode
 from langchain_core.messages.utils import count_tokens_approximately
+from langchain_core.messages import RemoveMessage
+from .data_loader import load_influencer_data
+
 
 
 
@@ -24,51 +27,81 @@ business_state = State()
 
 class BusinessInteractionChatbot:
     def __init__(self):
+        self.messages = []
+
         self.react_agent=create_react_agent(
         model=llm.bind_tools([retrieve_tool]),
         tools=[retrieve_tool]
         )
-        self.summarization_model = llm.bind(max_tokens=400)
+        self.summarization_model = llm.bind(max_tokens=128)
 
         self.summarization_node = SummarizationNode(
         token_counter=count_tokens_approximately,
         model=self.summarization_model,
-        max_tokens=256,
+        max_tokens=300,
         max_tokens_before_summary=256,
         max_summary_tokens=128,
         )
-
         self.memory = MemorySaver()
         # self.llm = ChatGroq(model_name="Gemma2-9b-It")
         self.workflow = self._initialize_workflow()
         self.interact_agent = self.workflow.compile(checkpointer=self.memory)
-        self.messages = []
 
 
     def _initialize_workflow(self):
         workflow = StateGraph(MessagesState)
         workflow.add_node("chatbot", self._call_model)
         workflow.add_node("summarize",self.summarization_node)
+        workflow.add_node("remove_message",self.delete_messages)
+
         workflow.add_edge(START, "summarize")
         workflow.add_edge("summarize", "chatbot")
+        workflow.add_edge("chatbot","remove_message")
         workflow.add_edge("chatbot", END)
         return workflow
+    
+    def delete_messages(self,state):
+        print('Entered message deletion....')
+        if len(self.messages) > 5:
+            print('satisfied...')
+            self.messages = self.messages[2:]
+
+    def manual_retrieval(self):
+        embedded_query = ST.encode(str([msg['content'] for msg in self.messages if msg['role'] == 'user']))  # Embed each topic
+        data = load_influencer_data()
+        scores, retrieved_examples = data.get_nearest_examples("embeddings", embedded_query, k=1)
+
+        # Construct a list of dictionaries for this topic
+        result = [{user: story} for user, story in zip(retrieved_examples['username'], retrieved_examples['agentic_story'])]
+        return result
+
 
     def _call_model(self, state):
         print('Entered into callmodel')
         template = business_interaction_prompt
         messages = [SystemMessage(content=template)] + state["messages"]
-        tool_response = self.react_agent.invoke({'messages':messages})['messages'][-2]
-        response = self.react_agent.invoke({'messages':messages})['messages'][-1]
-        print('Tool response:',tool_response)
-        return {"messages": [response]}
+        # response = self.react_agent.invoke({'messages':messages})['messages'][-2]
+        response = self.react_agent.invoke({'messages':messages})['messages']
+        if response [-2].name == None:
+            print('Entered into manual retrieval')
+            retrievals = self.manual_retrieval()
+            template = business_retrieval_prompt(retrievals)
+            messages = [SystemMessage(content=template)] + state["messages"]
+            backup_response = self.react_agent.invoke({'messages':messages})['messages'][-1]
+            print('Backup response:',backup_response.content)
+            return {"messages": [backup_response.content]}
+
+        else:
+            return {"messages": [response[-1]]}
 
     def chat(self, user_input: str):
         print('Entered into chat')
+        
         self.messages.append({"role": "user", "content": user_input})
         config = {"configurable": {"thread_id": "2"}}
         response = self.interact_agent.invoke({"messages":self.messages}, config)['messages'][-1].content
         print('The response:',response)
         self.messages.append({"role": "assistant", "content": response})
+        print('The message_history:',self.messages)
         business_state.interactions.append({'user': user_input, 'agent_response': response})
         return response