Use StateKeys constants

Author: derailed-dash

src/llms_gen_agent/constants.py

@@ -0,0 +1,10 @@
+class StateKeys:
+    ALL_SUMMARIES = "all_summaries"
+    BATCHES = "batches"
+    BATCH_SUMMARIES = "batch_summaries"
+    CURRENT_BATCH = "current_batch"
+    DOC_SUMMARIES = "doc_summaries"
+    FILES = "files"
+    FILES_CONTENT = "files_content"
+    LOOP_ITERATION = "loop_iteration"
+    PROJECT_SUMMARY_RAW = "project_summary_raw"

src/llms_gen_agent/sub_agents/doc_summariser/tools.py

@@ -12,6 +12,7 @@
 from google.adk.tools import ToolContext
 
 from llms_gen_agent.config import logger
+from llms_gen_agent.constants import StateKeys
 
 
 def read_files(tool_context: ToolContext) -> dict:
@@ -20,7 +21,6 @@ def read_files(tool_context: ToolContext) -> dict:
     This tool retrieves a list of file paths from the `current_batch` key in the
     `tool_context.state`. It then iterates through this list, reads the
     content of each file, and stores it in a dictionary under the
-
     `files_content` key in the `tool_context.state`. The file path serves as
     the key for its content.
 
@@ -34,29 +34,30 @@ def read_files(tool_context: ToolContext) -> dict:
 
     # The files to read are either in 'current_batch' (for batched processing)
     # or in 'files' (for direct processing or initial setup).
-    file_paths = tool_context.state.get("current_batch", tool_context.state.get("files", []))
+    file_paths = tool_context.state.get(StateKeys.CURRENT_BATCH, 
+                                        tool_context.state.get(StateKeys.FILES, []))
     logger.debug(f"Got {len(file_paths)} files")
 
     # Initialise our session state key    
-    tool_context.state["files_content"] = {}
+    tool_context.state[StateKeys.FILES_CONTENT] = {}
 
     response = {"status": "success"}
     for file_path in file_paths:
-        if file_path not in tool_context.state["files_content"]:
+        if file_path not in tool_context.state[StateKeys.FILES_CONTENT]:
             try:
                 logger.debug(f"Reading file: {file_path}")
                 with open(file_path) as f:
                     content = f.read()
                     logger.debug(f"Read content: {content[:80]}...")
-                    tool_context.state["files_content"][file_path] = content
+                    tool_context.state[StateKeys.FILES_CONTENT][file_path] = content
             except (FileNotFoundError, PermissionError, UnicodeDecodeError) as e:
                 logger.warning("Could not read file %s: %s", file_path, e)
                 # Store an error message so the summarizer knows it failed
-                tool_context.state["files_content"][file_path] = f"Error: Could not read file. Reason: {e}"
+                tool_context.state[StateKeys.FILES_CONTENT][file_path] = f"Error: Could not read file. Reason: {e}"
                 response = {"status": "warnings"}
             except Exception as e:
                 logger.error("An unexpected error occurred while reading %s: %s", file_path, e)
-                tool_context.state["files_content"][file_path] = f"Error: An unexpected error occurred. Reason: {e}"
+                tool_context.state[StateKeys.FILES_CONTENT][file_path] = f"Error: An unexpected error occurred. Reason: {e}"
                 response = {"status": "warnings"}
 
     return response
@@ -68,38 +69,38 @@ def create_file_batches(tool_context: ToolContext, batch_size: int) -> list[list
     divides them into smaller batches, and stores these batches back into the
     session state for iterative processing by the LoopAgent.
     """
-    file_paths = tool_context.state.get("files", [])
+    file_paths = tool_context.state.get(StateKeys.FILES, [])
     logger.debug(f"create_file_batches: Received {len(file_paths)} files from session state.")
     logger.debug(f"Creating batches for {len(file_paths)} files with batch size {batch_size}")
     if not file_paths:
         logger.debug("No files to batch.")
-        tool_context.state["batches"] = [] # Ensure batches is set even if empty
+        tool_context.state[StateKeys.BATCHES] = [] # Ensure batches is set even if empty
         return []
     num_batches = math.ceil(len(file_paths) / batch_size)
     batches = [file_paths[i * batch_size:(i + 1) * batch_size] for i in range(num_batches)]
     logger.debug(f"Created {len(batches)} batches.")
-    tool_context.state["batches"] = batches # Store batches in session state
+    tool_context.state[StateKeys.BATCHES] = batches # Store batches in session state
     return batches
 
 
 def process_batch_selection(tool_context: ToolContext) -> dict:
     """Manages the batch selection for the loop, increments iteration counter, and logs batch info."""
     logger.debug("Executing process_batch_selection")
 
-    batches = tool_context.state.get("batches", [])
-    loop_iteration = tool_context.state.get("loop_iteration", 0)
+    batches = tool_context.state.get(StateKeys.BATCHES, [])
+    loop_iteration = tool_context.state.get(StateKeys.LOOP_ITERATION, 0)
 
     if not batches:
         logger.debug("No more batches to process. Exiting loop.")
         tool_context.actions.escalate = True # Signal to exit the loop
         return {"status": "no_more_batches"}
 
     current_batch = batches.pop(0) # Get the next batch
-    tool_context.state["batches"] = batches # Update batches in state
-    tool_context.state["current_batch"] = current_batch # Set current batch
+    tool_context.state[StateKeys.BATCHES] = batches # Update batches in state
+    tool_context.state[StateKeys.CURRENT_BATCH] = current_batch # Set current batch
 
     loop_iteration += 1
-    tool_context.state["loop_iteration"] = loop_iteration
+    tool_context.state[StateKeys.LOOP_ITERATION] = loop_iteration
 
     logger.debug(f"Processing batch {loop_iteration}. Files in batch: {len(current_batch)}. Remaining batches: {len(batches)}")
 
@@ -114,16 +115,16 @@ def update_summaries(tool_context: ToolContext) -> dict:
     """
     logger.debug("Executing update_summaries")
 
-    batch_summaries_output = tool_context.state.get("batch_summaries", {})
-    batch_summaries = batch_summaries_output.get("batch_summaries", {}) # Get the actual dict from the output_key
+    batch_summaries_output = tool_context.state.get(StateKeys.BATCH_SUMMARIES, {})
+    batch_summaries = batch_summaries_output.get(StateKeys.BATCH_SUMMARIES, {}) # Get the actual dict from the output_key
 
-    if "all_summaries" not in tool_context.state:
-        tool_context.state["all_summaries"] = {}
+    if StateKeys.ALL_SUMMARIES not in tool_context.state:
+        tool_context.state[StateKeys.ALL_SUMMARIES] = {}
 
-    tool_context.state["all_summaries"].update(batch_summaries)
+    tool_context.state[StateKeys.ALL_SUMMARIES].update(batch_summaries)
 
     logger.debug(f"Merged {len(batch_summaries)} summaries from current batch. "
-                 f"Total summaries collected: {len(tool_context.state['all_summaries'])}")
+                 f"Total summaries collected: {len(tool_context.state[StateKeys.ALL_SUMMARIES])}")
 
     return {"status": "success"}
 
@@ -136,15 +137,16 @@ def finalize_summaries(tool_context: ToolContext) -> dict:
     in `tool_context.state["doc_summaries"]`.
     """
     logger.debug("Executing finalize_summaries")
-    all_summaries = tool_context.state.get("all_summaries", {})
-    project_summary_raw = tool_context.state.get("project_summary_raw", {}).get("project_summary", "No project summary found.")
+    all_summaries = tool_context.state.get(StateKeys.ALL_SUMMARIES, {})
+    project_summary_raw = tool_context.state.get(StateKeys.PROJECT_SUMMARY_RAW, 
+                                                 {}).get("project_summary", "No project summary found.")
 
     final_doc_summaries = {
         "summaries": {
             **all_summaries,
             "project": project_summary_raw
         }
     }
-    tool_context.state["doc_summaries"] = final_doc_summaries
+    tool_context.state[StateKeys.DOC_SUMMARIES] = final_doc_summaries
 
     return {"status": "success"}