pycityagent 2.0.0a51__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a53__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/cityagent/blocks/cognition_block.py

@@ -1,4 +1,4 @@
-from pycityagent.llm.llm import LLM
+from pycityagent.llm import LLM
 from pycityagent.workflow.block import Block
 from pycityagent.memory import Memory
 from pycityagent.environment.simulator import Simulator
@@ -23,47 +23,60 @@ def extract_json(output_str):
         return None
 
 class CognitionBlock(Block):
+    configurable_fields = ["top_k"]
+    default_values = {"top_k": 20}
+    fields_description = {
+        "top_k": "Number of most relevant memories to return, defaults to 20"
+    }
+
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
-        super().__init__("CognitionBlock", llm, memory, simulator)
-        self.last_trigger_time = None
-        self.time_diff = 24*60*60 # 24小时
-        self.trigger_time = 0
-        self.token_consumption = 0
+        super().__init__("CognitionBlock", llm=llm, memory=memory, simulator=simulator)
+        self.top_k = 20
         
     async def set_status(self, status):
         for key in status:
-            await self.memory.update(key, status[key])
+            await self.memory.status.update(key, status[key])
         return
     
-    async def check_trigger(self):
-        now_time = await self.simulator.get_time()
-        if self.last_trigger_time is None or now_time - self.last_trigger_time > self.time_diff:
-            self.last_trigger_time = now_time
-            return True
-        return False
-    
-    async def attitude_update(self, topic):
-        description_prompt = """
-        You are a {gender}, aged {age}, belonging to the {race} race and identifying as {religion}. 
-        Your marital status is {marital_status}, and you currently reside in a {residence} area. 
-        Your occupation is {occupation}, and your education level is {education}. 
-        You are {personality}, with a consumption level of {consumption} and a family consumption level of {family_consumption}. 
-        Your income is {income}, and you are skilled in {skill}.
-        My current emotion intensities are (0 meaning not at all, 10 meaning very much):
-        sadness: {sadness}, joy: {joy}, fear: {fear}, disgust: {disgust}, anger: {anger}, surprise: {surprise}.
-        You have the following thoughts: {thought}.
-        In the following 21 words, I have chosen {emotion_types} to represent your current status:
-        Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
-        """
-        incident_list = await self.memory.get("incident")
-        if incident_list:
-            incident_prompt = "Today, these incidents happened:"
-            for incident in incident_list:
-                incident_prompt += incident
-        else:
-            incident_prompt = "No incidents happened today."
-        attitude = await self.memory.get("attitude")
-        if topic in attitude:
+    async def attitude_update(self):
+        """Cognition - attitude update workflow"""
+        attitude = await self.memory.status.get("attitude")
+        prompt_data = {
+            "gender": await self.memory.status.get("gender"),
+            "age": await self.memory.status.get("age"),
+            "race": await self.memory.status.get("race"),
+            "religion": await self.memory.status.get("religion"),
+            "marital_status": await self.memory.status.get("marital_status"),
+            "residence": await self.memory.status.get("residence"),
+            "occupation": await self.memory.status.get("occupation"),
+            "education": await self.memory.status.get("education"),
+            "personality": await self.memory.status.get("personality"),
+            "consumption": await self.memory.status.get("consumption"),
+            "family_consumption": await self.memory.status.get("family_consumption"),
+            "income": await self.memory.status.get("income"),
+            "skill": await self.memory.status.get("skill"),
+            "thought": await self.memory.status.get("thought"),
+            "emotion_types": await self.memory.status.get("emotion_types")
+        }
+        for topic in attitude:
+            description_prompt = """
+            You are a {gender}, aged {age}, belonging to the {race} race and identifying as {religion}. 
+            Your marital status is {marital_status}, and you currently reside in a {residence} area. 
+            Your occupation is {occupation}, and your education level is {education}. 
+            You are {personality}, with a consumption level of {consumption} and a family consumption level of {family_consumption}. 
+            Your income is {income}, and you are skilled in {skill}.
+            My current emotion intensities are (0 meaning not at all, 10 meaning very much):
+            sadness: {sadness}, joy: {joy}, fear: {fear}, disgust: {disgust}, anger: {anger}, surprise: {surprise}.
+            You have the following thoughts: {thought}.
+            In the following 21 words, I have chosen {emotion_types} to represent your current status:
+            Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
+            """
+            incident_str = await self.memory.stream.search_today(top_k=self.top_k)
+            if incident_str:
+                incident_prompt = "Today, these incidents happened:"
+                incident_prompt += incident_str
+            else:
+                incident_prompt = "No incidents happened today."
             previous_attitude = str(attitude[topic])  # Convert to string
             problem_prompt = (
                 f"You need to decide your attitude towards topic: {topic}, "
@@ -71,70 +84,39 @@ class CognitionBlock(Block):
                 "(0 meaning oppose, 10 meaning support). "
                 "Please return a new attitude rating (0-10, smaller meaning oppose, larger meaning support) in JSON format, and explain, e.g. {{\"attitude\": 5}}"
             )
-        else:
-            problem_prompt = (
-                f"You need to decide your attitude towards topic: {topic}, "
-                "which you have not rated your attitude towards this topic yet. "
-                "(0 meaning oppose, 10 meaning support). "
-                "Please return a new attitude rating (0-10, smaller meaning oppose, larger meaning support) in JSON format, and explain, e.g. {{\"attitude\": 5}}"
-            )
-        question_prompt = description_prompt + incident_prompt + problem_prompt
-        question_prompt = FormatPrompt(question_prompt)
-        emotion = await self.memory.get("emotion")
-        sadness = emotion["sadness"]
-        joy = emotion["joy"]
-        fear = emotion["fear"]
-        disgust = emotion["disgust"]
-        anger = emotion["anger"]
-        surprise = emotion["surprise"]
-        
-        question_prompt.format(
-            gender=await self.memory.get("gender"),
-            age=await self.memory.get("age"),
-            race=await self.memory.get("race"),
-            religion=await self.memory.get("religion"),
-            marital_status=await self.memory.get("marital_status"),
-            residence=await self.memory.get("residence"),
-            occupation=await self.memory.get("occupation"),
-            education=await self.memory.get("education"),
-            personality=await self.memory.get("personality"),
-            consumption=await self.memory.get("consumption"),
-            family_consumption=await self.memory.get("family_consumption"),
-            income=await self.memory.get("income"),
-            skill=await self.memory.get("skill"),
-            sadness=sadness,
-            joy=joy,
-            fear=fear,
-            disgust=disgust,
-            anger=anger,
-            surprise=surprise,
-            thought=await self.memory.get("thought"),
-            emotion_types=await self.memory.get("emotion_types")
-        )
-        evaluation = True
-        for retry in range(10):
-            try:
-                response = await self.llm.atext_request(question_prompt.to_dialog(), timeout=300)
-                response = json.loads(extract_json(response))
-                evaluation = False
-                break
-            except:
-                pass
-        if evaluation:
-            raise f"Request for attitude:{topic} update failed"
-        logger.info(f"""Cognition updated attitude:{topic}:
-            attitude: {response['attitude']}""")
-        attitude[topic] = response["attitude"]
-        await self.memory.update("attitude", attitude)
-        return
+            question_prompt = description_prompt + incident_prompt + problem_prompt
+            question_prompt = FormatPrompt(question_prompt)
+            emotion = await self.memory.status.get("emotion")
+            sadness = emotion["sadness"]
+            joy = emotion["joy"]
+            fear = emotion["fear"]
+            disgust = emotion["disgust"]
+            anger = emotion["anger"]
+            surprise = emotion["surprise"]
+            prompt_data["sadness"] = sadness
+            prompt_data["joy"] = joy
+            prompt_data["fear"] = fear
+            prompt_data["disgust"] = disgust
+            prompt_data["anger"] = anger
+            prompt_data["surprise"] = surprise
+            
+            question_prompt.format(**prompt_data)
+            evaluation = True
+            for retry in range(10):
+                try:
+                    response = await self.llm.atext_request(question_prompt.to_dialog(), timeout=300)
+                    response = json.loads(extract_json(response))
+                    evaluation = False
+                    break
+                except:
+                    pass
+            if evaluation:
+                raise f"Request for attitude:{topic} update failed"
+            attitude[topic] = response["attitude"]
+        await self.memory.status.update("attitude", attitude)
     
-    async def forward(self): #每日结算
-        whether_trigger = await self.check_trigger()
-        if not whether_trigger:
-            return
-        self.trigger_time += 1
-        consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
-
+    async def thought_update(self):
+        """Cognition - thought update workflow"""
         description_prompt = """
         You are a {gender}, aged {age}, belonging to the {race} race and identifying as {religion}. 
         Your marital status is {marital_status}, and you currently reside in a {residence} area. 
@@ -147,21 +129,20 @@ class CognitionBlock(Block):
         In the following 21 words, I have chosen {emotion_types} to represent your current status:
         Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
         """
-        incident_list = await self.memory.get("incident")
-        if incident_list:
-            incident_prompt = "Today, these incidents happened:"
-            for incident in incident_list:
-                incident_prompt += incident
+        incident_str = await self.memory.stream.search_today(top_k=20)
+        if incident_str:
+            incident_prompt = "Today, these incidents happened:\n" + incident_str
         else:
             incident_prompt = "No incidents happened today."
         question_prompt = """
-            Please reconsider your emotion intensities: 
-            sadness, joy, fear, disgust, anger, surprise (0 meaning not at all, 10 meaning very much).
-            Also summerize you current thoughts, and choose a word to describe your status: Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
-            Return in JSON format, e.g. {{"sadness": 5, "joy": 5, "fear": 5, "disgust": 5, "anger": 5, "surprise": 5, "thought": "Currently nothing good or bad is happening, I think ....", "word": "Relief"}}"""
+            Please review what happened today and share your thoughts and feelings about it.
+            Consider your current emotional state and experiences, then:
+            1. Summarize your thoughts and reflections on today's events
+            2. Choose one word that best describes your current emotional state from: Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
+            Return in JSON format, e.g. {{"thought": "Currently nothing good or bad is happening, I think ...."}}"""
         question_prompt = description_prompt + incident_prompt + question_prompt
         question_prompt = FormatPrompt(question_prompt)
-        emotion = await self.memory.get("emotion")
+        emotion = await self.memory.status.get("emotion")
         sadness = emotion["sadness"]
         joy = emotion["joy"]
         fear = emotion["fear"]
@@ -169,28 +150,28 @@ class CognitionBlock(Block):
         anger = emotion["anger"]
         surprise = emotion["surprise"]
         question_prompt.format(
-            gender=await self.memory.get("gender"),
-            age=await self.memory.get("age"),
-            race=await self.memory.get("race"),
-            religion=await self.memory.get("religion"),
-            marital_status=await self.memory.get("marital_status"),
-            residence=await self.memory.get("residence"),
-            occupation=await self.memory.get("occupation"),
-            education=await self.memory.get("education"),
-            personality=await self.memory.get("personality"),
-            consumption=await self.memory.get("consumption"),
-            family_consumption=await self.memory.get("family_consumption"),
-            income=await self.memory.get("income"),
-            skill=await self.memory.get("skill"),
+            gender=await self.memory.status.get("gender"),
+            age=await self.memory.status.get("age"),
+            race=await self.memory.status.get("race"),
+            religion=await self.memory.status.get("religion"),
+            marital_status=await self.memory.status.get("marital_status"),
+            residence=await self.memory.status.get("residence"),
+            occupation=await self.memory.status.get("occupation"),
+            education=await self.memory.status.get("education"),
+            personality=await self.memory.status.get("personality"),
+            consumption=await self.memory.status.get("consumption"),
+            family_consumption=await self.memory.status.get("family_consumption"),
+            income=await self.memory.status.get("income"),
+            skill=await self.memory.status.get("skill"),
             sadness=sadness,
             joy=joy,
             fear=fear,
             disgust=disgust,
             anger=anger,
             surprise=surprise,
-            emotion=await self.memory.get("emotion"),
-            thought=await self.memory.get("thought"),
-            emotion_types=await self.memory.get("emotion_types")
+            emotion=await self.memory.status.get("emotion"),
+            thought=await self.memory.status.get("thought"),
+            emotion_types=await self.memory.status.get("emotion_types")
         )
         
         evaluation = True
@@ -204,30 +185,32 @@ class CognitionBlock(Block):
                 pass
         if evaluation:
             raise Exception("Request for cognition update failed")
-        logger.info(f"""Cognition updated emotion intensities:
-            sadness: {response['sadness']},
-            joy: {response['joy']},
-            fear: {response['fear']},
-            disgust: {response['disgust']},
-            anger: {response['anger']},
-            surprise: {response['surprise']}"
-            thought: {response['thought']}
-            emotion_types: {response['word']}""")
         
-        await self.memory.update("incident", [])
-        await self.memory.update("emotion", {"sadness": int(response["sadness"]), "joy": int(response["joy"]), "fear": int(response["fear"]), "disgust": int(response["disgust"]), "anger": int(response["anger"]), "surprise": int(response["surprise"])})
-        await self.memory.update("thought", str(response["thought"]))
-        await self.memory.update("emotion_types", str(response["word"]))
-
-        consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
-        self.token_consumption += consumption_end - consumption_start
+        thought = str(response["thought"])
+        await self.memory.status.update("thought", thought)
+        await self.memory.stream.add_cognition(description=thought)
         return
     
-    async def emotion_update(self, incident): #每日结算
+    async def end_of_day(self):
+        """Cognition - end of day workflow"""
+        time = await self.simulator.get_simulator_second_from_start_of_day()
+        if time >= 86400 - 10 * 60:
+            return True
+        return False
+    
+    async def forward(self):
+        """Cognition workflow: Daily update"""
+        # cognition update: thought and attitude
+        if await self.end_of_day():
+            await self.thought_update()
+            await self.attitude_update()
+
+    async def emotion_update(self, incident):
+        """Cognition - emotion update workflow"""
         whether_trigger = await self.check_trigger()
         if not whether_trigger:
             return
-
+        print(f"Updating emotion for {incident}")
         description_prompt = """
         You are a {gender}, aged {age}, belonging to the {race} race and identifying as {religion}. 
         Your marital status is {marital_status}, and you currently reside in a {residence} area. 
@@ -241,14 +224,14 @@ class CognitionBlock(Block):
         Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
         """
 
-        incident_prompt = incident  #waiting for incident port
+        incident_prompt = f"{incident}"  #waiting for incident port
         question_prompt = """
             Please reconsider your emotion intensities: 
             sadness, joy, fear, disgust, anger, surprise (0 meaning not at all, 10 meaning very much).
-            Return in JSON format, e.g. {{"sadness": 5, "joy": 5, "fear": 5, "disgust": 5, "anger": 5, "surprise": 5}}"""
+            Return in JSON format, e.g. {{"sadness": 5, "joy": 5, "fear": 5, "disgust": 5, "anger": 5, "surprise": 5, "conclusion": "I feel ..."}}"""
         question_prompt = description_prompt + incident_prompt + question_prompt
         question_prompt = FormatPrompt(question_prompt)
-        emotion = await self.memory.get("emotion")
+        emotion = await self.memory.status.get("emotion")
         sadness = emotion["sadness"]
         joy = emotion["joy"]
         fear = emotion["fear"]
@@ -256,28 +239,28 @@ class CognitionBlock(Block):
         anger = emotion["anger"]
         surprise = emotion["surprise"]
         question_prompt.format(
-            gender=await self.memory.get("gender"),
-            age=await self.memory.get("age"),
-            race=await self.memory.get("race"),
-            religion=await self.memory.get("religion"),
-            marital_status=await self.memory.get("marital_status"),
-            residence=await self.memory.get("residence"),
-            occupation=await self.memory.get("occupation"),
-            education=await self.memory.get("education"),
-            personality=await self.memory.get("personality"),
-            consumption=await self.memory.get("consumption"),
-            family_consumption=await self.memory.get("family_consumption"),
-            income=await self.memory.get("income"),
-            skill=await self.memory.get("skill"),
+            gender=await self.memory.status.get("gender"),
+            age=await self.memory.status.get("age"),
+            race=await self.memory.status.get("race"),
+            religion=await self.memory.status.get("religion"),
+            marital_status=await self.memory.status.get("marital_status"),
+            residence=await self.memory.status.get("residence"),
+            occupation=await self.memory.status.get("occupation"),
+            education=await self.memory.status.get("education"),
+            personality=await self.memory.status.get("personality"),
+            consumption=await self.memory.status.get("consumption"),
+            family_consumption=await self.memory.status.get("family_consumption"),
+            income=await self.memory.status.get("income"),
+            skill=await self.memory.status.get("skill"),
             sadness=sadness,
             joy=joy,
             fear=fear,
             disgust=disgust,
             anger=anger,
             surprise=surprise,
-            emotion=await self.memory.get("emotion"),
-            thought=await self.memory.get("thought"),
-            emotion_types=await self.memory.get("emotion_types")
+            emotion=await self.memory.status.get("emotion"),
+            thought=await self.memory.status.get("thought"),
+            emotion_types=await self.memory.status.get("emotion_types")
         )
         
         evaluation = True
@@ -288,17 +271,11 @@ class CognitionBlock(Block):
                 evaluation = False
                 break
             except Exception as e:
-                print(e)
+                logger.warning(f"Request for cognition update failed: {e}")
                 pass
         if evaluation:
             raise Exception("Request for cognition update failed")
-        logger.info(f"""Cognition updated emotion intensities:
-            sadness: {response['sadness']},
-            joy: {response['joy']},
-            fear: {response['fear']},
-            disgust: {response['disgust']},
-            anger: {response['anger']},
-            surprise: {response['surprise']}""")
 
-        await self.memory.update("emotion", {"sadness": int(response["sadness"]), "joy": int(response["joy"]), "fear": int(response["fear"]), "disgust": int(response["disgust"]), "anger": int(response["anger"]), "surprise": int(response["surprise"])})
-        return
+        await self.memory.status.update("emotion", {"sadness": int(response["sadness"]), "joy": int(response["joy"]), "fear": int(response["fear"]), "disgust": int(response["disgust"]), "anger": int(response["anger"]), "surprise": int(response["surprise"])})
+        await self.memory.status.update("emotion_types", str(response["word"]))
+        return response["conclusion"]