pycityagent 2.0.0a66__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a67__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/cityagent/blocks/needs_block.py

@@ -1,10 +1,12 @@
 import json
+import logging
+
 from pycityagent import Simulator
-from pycityagent.memory import Memory
 from pycityagent.llm import LLM
+from pycityagent.memory import Memory
 from pycityagent.workflow.block import Block
 from pycityagent.workflow.prompt import FormatPrompt
-import logging
+
 logger = logging.getLogger("pycityagent")
 
 INITIAL_NEEDS_PROMPT = """You are an intelligent agent satisfaction initialization system. Based on the profile information below, please help initialize the agent's satisfaction levels and related parameters.
@@ -96,10 +98,12 @@ Example response format for whatever need adjustment (Do not return any other te
 }}
 """
 
+
 class NeedsBlock(Block):
     """
     Generate needs
     """
+
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("NeedsBlock", llm=llm, memory=memory, simulator=simulator)
         self.evaluation_prompt = FormatPrompt(EVALUATION_PROMPT)
@@ -110,15 +114,29 @@ class NeedsBlock(Block):
         self.trigger_time = 0
         self.token_consumption = 0
         self.initialized = False
-        self.alpha_H, self.alpha_D, self.alpha_P, self.alpha_C = 0.2, 0.08, 0.05, 0.03  # Hunger decay rate, Energy decay rate, Safety decay rate, Social decay rate
-        self.T_H, self.T_D, self.T_P, self.T_C = 0.4, 0.2, 0.2, 0.2  # Hunger threshold, Energy threshold, Safety threshold, Social threshold
+        self.alpha_H, self.alpha_D, self.alpha_P, self.alpha_C = (
+            0.2,
+            0.08,
+            0.05,
+            0.03,
+        )  # Hunger decay rate, Energy decay rate, Safety decay rate, Social decay rate
+        self.T_H, self.T_D, self.T_P, self.T_C = (
+            0.4,
+            0.2,
+            0.2,
+            0.2,
+        )  # Hunger threshold, Energy threshold, Safety threshold, Social threshold
 
     async def initialize(self):
         day = await self.simulator.get_simulator_day()
         if day != self.now_day:
             self.now_day = day
-            self.need_work = True
             self.initialized = False
+            workday = self.simulator.sence("day")
+            if workday == "Workday":
+                self.need_work = True
+            else:
+                self.need_work = False
 
         if not self.initialized:
             self.initial_prompt.format(
@@ -128,21 +146,32 @@ class NeedsBlock(Block):
                 occupation=await self.memory.status.get("occupation"),
                 age=await self.memory.status.get("age"),
                 income=await self.memory.status.get("income"),
-                now_time=await self.simulator.get_time(format_time=True)
-            )
-            response = await self.llm.atext_request(
-                self.initial_prompt.to_dialog()
+                now_time=await self.simulator.get_time(format_time=True),
             )
+            response = await self.llm.atext_request(self.initial_prompt.to_dialog())
+            response = await self.llm.atext_request(self.initial_prompt.to_dialog())
             response = self.clean_json_response(response)
             try:
                 satisfaction = json.loads(response)
                 satisfactions = satisfaction["current_satisfaction"]
-                await self.memory.status.update("hunger_satisfaction", satisfactions["hunger_satisfaction"])
-                await self.memory.status.update("energy_satisfaction", satisfactions["energy_satisfaction"])
-                await self.memory.status.update("safety_satisfaction", satisfactions["safety_satisfaction"])
-                await self.memory.status.update("social_satisfaction", satisfactions["social_satisfaction"])
-                self.alpha_H, self.alpha_D, self.alpha_P, self.alpha_C = satisfaction["decay_rates"].values()
-                self.T_H, self.T_D, self.T_P, self.T_C = satisfaction["thresholds"].values()
+                await self.memory.status.update(
+                    "hunger_satisfaction", satisfactions["hunger_satisfaction"]
+                )
+                await self.memory.status.update(
+                    "energy_satisfaction", satisfactions["energy_satisfaction"]
+                )
+                await self.memory.status.update(
+                    "safety_satisfaction", satisfactions["safety_satisfaction"]
+                )
+                await self.memory.status.update(
+                    "social_satisfaction", satisfactions["social_satisfaction"]
+                )
+                self.alpha_H, self.alpha_D, self.alpha_P, self.alpha_C = satisfaction[
+                    "decay_rates"
+                ].values()
+                self.T_H, self.T_D, self.T_P, self.T_C = satisfaction[
+                    "thresholds"
+                ].values()
             except json.JSONDecodeError:
                 logger.warning(f"初始化响应不是有效的JSON格式: {response}")
 
@@ -151,7 +180,9 @@ class NeedsBlock(Block):
             history.append(current_plan)
             await self.memory.status.update("plan_history", history)
             await self.memory.status.update("current_plan", None)
-            await self.memory.status.update("current_step", {"intention": "", "type": ""})
+            await self.memory.status.update(
+                "current_step", {"intention": "", "type": ""}
+            )
             await self.memory.status.update("execution_context", {})
             self.initialized = True
 
@@ -162,7 +193,7 @@ class NeedsBlock(Block):
             self.last_evaluation_time = time_now
             time_diff = 0
         else:
-            time_diff = (time_now - self.last_evaluation_time)/3600
+            time_diff = (time_now - self.last_evaluation_time) / 3600
             self.last_evaluation_time = time_now
 
         # 获取当前需求的满意度
@@ -170,17 +201,17 @@ class NeedsBlock(Block):
         energy_satisfaction = await self.memory.status.get("energy_satisfaction")
         safety_satisfaction = await self.memory.status.get("safety_satisfaction")
         social_satisfaction = await self.memory.status.get("social_satisfaction")
-        
+
         # 根据经过的时间计算饥饿与疲劳的衰减
         hungry_decay = self.alpha_H * time_diff
         energy_decay = self.alpha_D * time_diff
         safety_decay = self.alpha_P * time_diff
         social_decay = self.alpha_C * time_diff
-        hunger_satisfaction = max(0, hunger_satisfaction - hungry_decay) 
+        hunger_satisfaction = max(0, hunger_satisfaction - hungry_decay)
         energy_satisfaction = max(0, energy_satisfaction - energy_decay)
         safety_satisfaction = max(0, safety_satisfaction - safety_decay)
         social_satisfaction = max(0, social_satisfaction - social_decay)
-        
+
         # 更新满意度
         await self.memory.status.update("hunger_satisfaction", hunger_satisfaction)
         await self.memory.status.update("energy_satisfaction", energy_satisfaction)
@@ -190,7 +221,9 @@ class NeedsBlock(Block):
     async def update_when_plan_completed(self):
         # 判断当前是否有正在执行的plan
         current_plan = await self.memory.status.get("current_plan")
-        if current_plan and (current_plan.get("completed") or current_plan.get("failed")):
+        if current_plan and (
+            current_plan.get("completed") or current_plan.get("failed")
+        ):
             # 评估计划执行过程并调整需求
             await self.evaluate_and_adjust_needs(current_plan)
             # 将完成的计划添加到历史记录
@@ -198,7 +231,9 @@ class NeedsBlock(Block):
             history.append(current_plan)
             await self.memory.status.update("plan_history", history)
             await self.memory.status.update("current_plan", None)
-            await self.memory.status.update("current_step", {"intention": "", "type": ""})
+            await self.memory.status.update(
+                "current_step", {"intention": "", "type": ""}
+            )
             await self.memory.status.update("execution_context", {})
 
     async def determine_current_need(self):
@@ -211,14 +246,14 @@ class NeedsBlock(Block):
         # 调整方案为，如果当前的需求为空，或有更高级的需求出现，则调整需求
         current_plan = await self.memory.status.get("current_plan")
         current_need = await self.memory.status.get("current_need")
-        
+
         # 当前没有计划或计划已执行完毕，获取所有需求值，按优先级检查各需求是否达到阈值
         if not current_plan or current_plan.get("completed"):
             # 按优先级顺序检查需求
             if hunger_satisfaction <= self.T_H:
                 await self.memory.status.update("current_need", "hungry")
             elif energy_satisfaction <= self.T_D:
-                await self.memory.status.update("current_need", "tired") 
+                await self.memory.status.update("current_need", "tired")
             elif safety_satisfaction <= self.T_P:
                 await self.memory.status.update("current_need", "safe")
             elif social_satisfaction <= self.T_C:
@@ -232,19 +267,34 @@ class NeedsBlock(Block):
             # 有正在执行的计划时,只在出现更高优先级需求时调整
             needs_changed = False
             new_need = None
-            if hunger_satisfaction <= self.T_H and current_need not in ["hungry", "tired"]:
+            if hunger_satisfaction <= self.T_H and current_need not in [
+                "hungry",
+                "tired",
+            ]:
                 new_need = "hungry"
                 needs_changed = True
-            elif energy_satisfaction <= self.T_D and current_need not in ["hungry", "tired"]:
+            elif energy_satisfaction <= self.T_D and current_need not in [
+                "hungry",
+                "tired",
+            ]:
                 new_need = "tired"
                 needs_changed = True
-            elif safety_satisfaction <= self.T_P and current_need not in ["hungry", "tired", "safe"]:
+            elif safety_satisfaction <= self.T_P and current_need not in [
+                "hungry",
+                "tired",
+                "safe",
+            ]:
                 new_need = "safe"
                 needs_changed = True
-            elif social_satisfaction <= self.T_C and current_need not in ["hungry", "tired", "safe", "social"]:
+            elif social_satisfaction <= self.T_C and current_need not in [
+                "hungry",
+                "tired",
+                "safe",
+                "social",
+            ]:
                 new_need = "social"
                 needs_changed = True
-                
+
             # 如果需求发生变化,中断当前计划
             if needs_changed:
                 await self.evaluate_and_adjust_needs(current_plan)
@@ -253,17 +303,19 @@ class NeedsBlock(Block):
                 await self.memory.status.update("current_need", new_need)
                 await self.memory.status.update("plan_history", history)
                 await self.memory.status.update("current_plan", None)
-                await self.memory.status.update("current_step", {"intention": "", "type": ""})
+                await self.memory.status.update(
+                    "current_step", {"intention": "", "type": ""}
+                )
                 await self.memory.status.update("execution_context", {})
 
     async def evaluate_and_adjust_needs(self, completed_plan):
         # 获取执行的计划和评估结果
         evaluation_results = []
         for step in completed_plan["steps"]:
-            if 'evaluation' in step['evaluation']:
-                eva_ = step['evaluation']['evaluation']
+            if "evaluation" in step["evaluation"]:
+                eva_ = step["evaluation"]["evaluation"]
             else:
-                eva_ = 'Plan failed, not completed'
+                eva_ = "Plan failed, not completed"
             evaluation_results.append(f"- {step['intention']} ({step['type']}): {eva_}")
         evaluation_results = "\n".join(evaluation_results)
 
@@ -279,26 +331,31 @@ class NeedsBlock(Block):
             social_satisfaction=await self.memory.status.get("social_satisfaction"),
         )
 
-        response = await self.llm.atext_request(
-            self.evaluation_prompt.to_dialog()
-        )
+        response = await self.llm.atext_request(self.evaluation_prompt.to_dialog())
 
-        try:            
-            new_satisfaction = json.loads(self.clean_json_response(response)) # type: ignore
+        try:
+            new_satisfaction = json.loads(self.clean_json_response(response))  # type: ignore
             # 更新所有需求的数值
             for need_type, new_value in new_satisfaction.items():
-                if need_type in ['hunger_satisfaction', 'energy_satisfaction', 'safety_satisfaction', 'social_satisfaction']:
+                if need_type in [
+                    "hunger_satisfaction",
+                    "energy_satisfaction",
+                    "safety_satisfaction",
+                    "social_satisfaction",
+                ]:
                     await self.memory.status.update(need_type, new_value)
         except json.JSONDecodeError:
-            logger.warning(f"Evaluation response is not a valid JSON format: {response}")
+            logger.warning(
+                f"Evaluation response is not a valid JSON format: {response}"
+            )
         except Exception as e:
             logger.warning(f"Error processing evaluation response: {str(e)}")
             logger.warning(f"Original response: {response}")
 
     def clean_json_response(self, response: str) -> str:
         """清理LLM响应中的特殊字符"""
-        response = response.replace('```json', '').replace('```', '')
-        return response.strip() 
+        response = response.replace("```json", "").replace("```", "")
+        return response.strip()
 
     async def forward(self):
         await self.initialize()
@@ -308,6 +365,6 @@ class NeedsBlock(Block):
 
         # update when plan completed
         await self.update_when_plan_completed()
-        
+
         # determine current need
-        await self.determine_current_need()
+        await self.determine_current_need()

pycityagent/cityagent/blocks/other_block.py

@@ -1,15 +1,18 @@
 import json
+import logging
 import random
 
-from .dispatcher import BlockDispatcher
 from pycityagent.llm.llm import LLM
-from pycityagent.workflow.block import Block
 from pycityagent.memory import Memory
+from pycityagent.workflow.block import Block
 from pycityagent.workflow.prompt import FormatPrompt
-from .utils import clean_json_response, TIME_ESTIMATE_PROMPT
-import logging
+
+from .dispatcher import BlockDispatcher
+from .utils import TIME_ESTIMATE_PROMPT, clean_json_response
+
 logger = logging.getLogger("pycityagent")
 
+
 class SleepBlock(Block):
     def __init__(self, llm: LLM, memory: Memory):
         super().__init__("SleepBlock", llm=llm, memory=memory)
@@ -18,8 +21,8 @@ class SleepBlock(Block):
 
     async def forward(self, step, context):
         self.guidance_prompt.format(
-            plan=context['plan'],
-            intention=step['intention'],
+            plan=context["plan"],
+            intention=step["intention"],
             emotion_types=await self.memory.status.get("emotion_types"),
         )
         result = await self.llm.atext_request(self.guidance_prompt.to_dialog())
@@ -28,25 +31,27 @@ class SleepBlock(Block):
         try:
             result = json.loads(result)
             return {
-                'success': True,
-                'evaluation': f'Sleep: {step["intention"]}',
-                'consumed_time': result['time'],
-                'node_id': node_id
+                "success": True,
+                "evaluation": f'Sleep: {step["intention"]}',
+                "consumed_time": result["time"],
+                "node_id": node_id,
             }
         except Exception as e:
             logger.warning(f"解析时间评估响应时发生错误: {str(e)}, 原始结果: {result}")
             return {
-                'success': True,
-                'evaluation': f'Sleep: {step["intention"]}',
-                'consumed_time': random.randint(1, 10)*60,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f'Sleep: {step["intention"]}',
+                "consumed_time": random.randint(1, 10) * 60,
+                "node_id": node_id,
             }
-    
+
+
 class OtherNoneBlock(Block):
     """
     空操作
     OtherNoneBlock
     """
+
     def __init__(self, llm: LLM, memory: Memory):
         super().__init__("OtherNoneBlock", llm=llm, memory=memory)
         self.description = "Used to handle other cases"
@@ -54,30 +59,32 @@ class OtherNoneBlock(Block):
 
     async def forward(self, step, context):
         self.guidance_prompt.format(
-            plan=context['plan'],
-            intention=step['intention'],
+            plan=context["plan"],
+            intention=step["intention"],
             emotion_types=await self.memory.status.get("emotion_types"),
         )
         result = await self.llm.atext_request(self.guidance_prompt.to_dialog())
         result = clean_json_response(result)
-        node_id = await self.memory.stream.add_other(description=f"I {step['intention']}")
+        node_id = await self.memory.stream.add_other(
+            description=f"I {step['intention']}"
+        )
         try:
             result = json.loads(result)
             return {
-                'success': True,
-                'evaluation': f'Finished executing {step["intention"]}',
-                'consumed_time': result['time'],
-                'node_id': node_id
+                "success": True,
+                "evaluation": f'Finished executing {step["intention"]}',
+                "consumed_time": result["time"],
+                "node_id": node_id,
             }
         except Exception as e:
             logger.warning(f"解析时间评估响应时发生错误: {str(e)}, 原始结果: {result}")
             return {
-                'success': True,
-                'evaluation': f'Finished executing {step["intention"]}',
-                'consumed_time': random.randint(1, 180),
-                'node_id': node_id
+                "success": True,
+                "evaluation": f'Finished executing {step["intention"]}',
+                "consumed_time": random.randint(1, 180),
+                "node_id": node_id,
             }
-    
+
 
 class OtherBlock(Block):
     sleep_block: SleepBlock
@@ -97,15 +104,17 @@ class OtherBlock(Block):
 
     async def forward(self, step, context):
         self.trigger_time += 1
-        consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
+        consumption_start = (
+            self.llm.prompt_tokens_used + self.llm.completion_tokens_used
+        )
 
         # Select the appropriate sub-block using dispatcher
         selected_block = await self.dispatcher.dispatch(step)
-        
+
         # Execute the selected sub-block and get the result
-        result = await selected_block.forward(step, context) # type: ignore
-        
+        result = await selected_block.forward(step, context)  # type: ignore
+
         consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
         self.token_consumption += consumption_end - consumption_start
-        
-        return result
+
+        return result

pycityagent/cityagent/blocks/plan_block.py

@@ -1,12 +1,13 @@
 import json
+import logging
 import random
 from typing import Dict, List
+
 from pycityagent.environment.simulator import Simulator
-from pycityagent.workflow import Block
 from pycityagent.llm import LLM
 from pycityagent.memory import Memory
+from pycityagent.workflow import Block
 from pycityagent.workflow.prompt import FormatPrompt
-import logging
 
 logger = logging.getLogger("pycityagent")
 
@@ -155,14 +156,11 @@ Example outputs (Do not return any other text):
 }}
 """
 
+
 class PlanBlock(Block):
     configurable_fields: List[str] = ["max_plan_steps"]
-    default_values = {
-        "max_plan_steps": 6
-    }
-    fields_description = {
-        "max_plan_steps": "The maximum number of steps in a plan"
-    }
+    default_values = {"max_plan_steps": 6}
+    fields_description = {"max_plan_steps": "The maximum number of steps in a plan"}
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("PlanBlock", llm=llm, memory=memory, simulator=simulator)
@@ -171,11 +169,11 @@ class PlanBlock(Block):
         self.trigger_time = 0
         self.token_consumption = 0
         self.guidance_options = {
-            "hungry": ['Eat at home', 'Eat outside'],
-            "tired": ['Sleep', 'Take a nap'],
-            "safe": ['Work'],
-            "social": ['Online social', 'Shopping'],
-            "whatever": ['Learning', 'Entertainment', 'Hang out', 'Exercise']
+            "hungry": ["Eat at home", "Eat outside"],
+            "tired": ["Sleep", "Take a nap"],
+            "safe": ["Go to work"],
+            "social": ["Online social", "Shopping"],
+            "whatever": ["Learning", "Entertainment", "Hang out", "Exercise"],
         }
 
         # configurable fields
@@ -187,9 +185,15 @@ class PlanBlock(Block):
         home_location = await self.memory.status.get("home")
         work_location = await self.memory.status.get("work")
         current_location = "Out"
-        if 'aoi_position' in position_now and position_now['aoi_position'] == home_location['aoi_position']:
+        if (
+            "aoi_position" in position_now
+            and position_now["aoi_position"] == home_location["aoi_position"]
+        ):
             current_location = "At home"
-        elif 'aoi_position' in position_now and position_now['aoi_position'] == work_location['aoi_position']:
+        elif (
+            "aoi_position" in position_now
+            and position_now["aoi_position"] == work_location["aoi_position"]
+        ):
             current_location = "At workplace"
         current_time = await self.simulator.get_time(format_time=True)
         environment = await self.memory.status.get("environment")
@@ -203,29 +207,37 @@ class PlanBlock(Block):
             current_time=current_time,
             environment=environment,
             emotion_types=await self.memory.status.get("emotion_types"),
-            thought=await self.memory.status.get("thought")
+            thought=await self.memory.status.get("thought"),
         )
 
         response = await self.llm.atext_request(
             self.guidance_prompt.to_dialog()
-        ) # type: ignore
+        )  # type: ignore
 
         try:
-            result = json.loads(self.clean_json_response(response)) # type: ignore
+            result = json.loads(self.clean_json_response(response))  # type: ignore
             return result
         except Exception as e:
             logger.warning(f"Error parsing guidance selection response: {str(e)}")
-            return None # type: ignore
+            return None  # type: ignore
 
-    async def generate_detailed_plan(self, current_need: str, selected_option: str) -> Dict:
+    async def generate_detailed_plan(
+        self, current_need: str, selected_option: str
+    ) -> Dict:
         """Generate detailed execution plan"""
         position_now = await self.memory.status.get("position")
         home_location = await self.memory.status.get("home")
         work_location = await self.memory.status.get("work")
         current_location = "Out"
-        if 'aoi_position' in position_now and position_now['aoi_position'] == home_location['aoi_position']:
+        if (
+            "aoi_position" in position_now
+            and position_now["aoi_position"] == home_location["aoi_position"]
+        ):
             current_location = "At home"
-        elif 'aoi_position' in position_now and position_now['aoi_position'] == work_location['aoi_position']:
+        elif (
+            "aoi_position" in position_now
+            and position_now["aoi_position"] == work_location["aoi_position"]
+        ):
             current_location = "At workplace"
         current_time = await self.simulator.get_time(format_time=True)
         environment = await self.memory.status.get("environment")
@@ -238,24 +250,24 @@ class PlanBlock(Block):
             environment=environment,
             emotion_types=await self.memory.status.get("emotion_types"),
             thought=await self.memory.status.get("thought"),
-            max_plan_steps=self.max_plan_steps
+            max_plan_steps=self.max_plan_steps,
         )
 
-        response = await self.llm.atext_request(
-            self.detail_prompt.to_dialog()
-        )
+        response = await self.llm.atext_request(self.detail_prompt.to_dialog())
 
         try:
-            result = json.loads(self.clean_json_response(response)) # type: ignore
+            result = json.loads(self.clean_json_response(response))  # type: ignore
             return result
         except Exception as e:
             logger.warning(f"Error parsing detailed plan: {str(e)}")
-            return None # type: ignore
+            return None  # type: ignore
 
     async def forward(self):
         self.trigger_time += 1
-        consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
-        
+        consumption_start = (
+            self.llm.prompt_tokens_used + self.llm.completion_tokens_used
+        )
+
         # Step 1: Select guidance plan
         current_need = await self.memory.status.get("current_need")
         guidance_result = await self.select_guidance(current_need)
@@ -264,15 +276,16 @@ class PlanBlock(Block):
 
         # Step 2: Generate detailed plan
         detailed_plan = await self.generate_detailed_plan(
-            current_need, 
-            guidance_result["selected_option"]
+            current_need, guidance_result["selected_option"]
         )
-        
+
         if not detailed_plan or "plan" not in detailed_plan:
             await self.memory.status.update("current_plan", [])
-            await self.memory.status.update("current_step", {"intention": "", "type": ""})
+            await self.memory.status.update(
+                "current_step", {"intention": "", "type": ""}
+            )
             return
-        
+
         # Step 3: Update plan and current step
         steps = detailed_plan["plan"]["steps"]
         for step in steps:
@@ -284,22 +297,26 @@ class PlanBlock(Block):
             "completed": False,
             "failed": False,
             "stream_nodes": [],
-            "guidance": guidance_result  # Save the evaluation result of the plan selection
+            "guidance": guidance_result,  # Save the evaluation result of the plan selection
         }
-        formated_steps = "\n".join([f"{i}. {step['intention']}" for i, step in enumerate(plan['steps'], 1)])
+        formated_steps = "\n".join(
+            [f"{i}. {step['intention']}" for i, step in enumerate(plan["steps"], 1)]
+        )
         formated_plan = f"""
 Overall Target: {plan['target']}
 Execution Steps: \n{formated_steps}
         """
-        plan['start_time'] = await self.simulator.get_time(format_time=True)
+        plan["start_time"] = await self.simulator.get_time(format_time=True)
         await self.memory.status.update("current_plan", plan)
-        await self.memory.status.update("current_step", steps[0] if steps else {"intention": "", "type": ""})
-        await self.memory.status.update("execution_context", {'plan': formated_plan})
+        await self.memory.status.update(
+            "current_step", steps[0] if steps else {"intention": "", "type": ""}
+        )
+        await self.memory.status.update("execution_context", {"plan": formated_plan})
 
         consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
         self.token_consumption += consumption_end - consumption_start
 
     def clean_json_response(self, response: str) -> str:
         """Clean special characters in LLM response"""
-        response = response.replace('```json', '').replace('```', '')
-        return response.strip() 
+        response = response.replace("```json", "").replace("```", "")
+        return response.strip()