Feature/rag_agent

agents/basic_agent.py

@@ -8,7 +8,7 @@
 from models.conversation import Conversation
 from models.generation_parameters import GenerationParameters
 from tenacity import wait_exponential, retry_if_exception_type, wait_random_exponential
-
+from agents.utils.prompt_utils import get_default_system_prompt
 from namespace import FactorioNamespace
 
 GENERAL_INSTRUCTIONS = \
@@ -175,8 +175,9 @@ def find_idle_furnaces(entities):
 
 
 class BasicAgent(AgentABC):
-   def __init__(self, model, system_prompt, task, *args, **kwargs):
-        instructions = GENERAL_INSTRUCTIONS+system_prompt+FINAL_INSTRUCTION
+   def __init__(self, model, system_prompt_parts, task, *args, **kwargs):
+
+        instructions = GENERAL_INSTRUCTIONS+get_default_system_prompt(system_prompt_parts)+FINAL_INSTRUCTION
         self.task = task
         instructions += f"\n\n### Goal\n{task.goal_description}\n\n"
         super().__init__( model, instructions, *args, **kwargs)

agents/query_agent.py

@@ -0,0 +1,236 @@
+import tenacity
+
+from agents import Response, CompletionResult, Policy
+from agents.agent_abc import AgentABC
+from agents.utils.formatters.recursive_report_formatter import RecursiveReportFormatter
+from agents.utils.llm_factory import LLMFactory
+from agents.utils.parse_response import parse_response
+from models.conversation import Conversation
+from models.generation_parameters import GenerationParameters
+from tenacity import wait_exponential, retry_if_exception_type, wait_random_exponential
+from agents.utils.prompt_utils import get_rag_system_prompt
+
+from namespace import FactorioNamespace
+
+GENERAL_INSTRUCTIONS = \
+"""
+# Factorio LLM Agent Instructions
+
+## Overview
+You are an AI agent designed to play Factorio, specializing in:
+- Long-horizon planning
+- Spatial reasoning 
+- Systematic automation
+
+## Environment Structure
+- Operates like an interactive Python shell
+- Agent messages = Python programs to execute
+- User responses = STDOUT/STDERR from REPL
+- Interacts through 27 core API methods (to be specified)
+
+## Response Format
+
+### 1. PLANNING Stage
+Think through each step extensively in natural language
+There are 2 types of policies you should write, 1) <information_lookup> (using prints and the query_information tool) or 2) <action>
+Your planning stage should address the following:
+1. Error Analysis
+   - Was there an error in the previous execution?
+   - If yes, what was the problem?
+2. Next Step classification
+   - What is the most useful next step, is it an <information_lookup> or an <action>? 
+   - What will this step achieve?
+3. Action Planning
+   if <information_lookup>:
+   - what information needs to be printed
+   - what environment information needs to be printed
+   If <action>:
+   - What specific actions are needed?
+   - What resources are required?
+
+Remember: If you don't have some information regarding howto use the API, use the <information_lookup> policy to print the information you need.
+DO NOT TRY TO USE THE API WITHOUT KNOWING HOW IT WORKS.
+
+### 2. POLICY Stage
+Write Python code to execute the <action> or <information_lookup> plan:
+```python
+# Code must be enclosed in Python tags
+your_code_here
+```
+MAXIMUM 20 LINES OF CODE PER POLICY
+
+## Best Practices
+
+### Querying Information
+- Make sure to query information using the query_information tool
+- Use the tool like you would research different areas on the wiki
+- The pages have info regarding how to use the api, how to carry out different actions and general factorio knowledge
+- Always wait for the information to be printed before executing steps that rely on this information. Do not execute steps that rely on information that has not been printed yet.
+
+
+### Modularity
+- Create small, modular policies
+- Each policy should have a single clear purpose
+- Keep policies easy to debug and modify
+- Avoid breaking existing automated structures
+- Encapsulate working logic into functions if needed
+
+### Debugging & Verification
+- Use print statements to monitor important state
+- Implement assert statements for self-verification
+- Use specific, parameterized assertion messages
+- Example: `assert condition, f"Expected {expected}, got {actual}"`
+
+### State Management
+- Consider entities needed for each step
+- Track entities across different inventories
+- Monitor missing requirements
+- Preserve working automated structures
+
+### Error Handling
+- Fix errors as they occur
+- Don't repeat previous steps
+- Continue from last successful execution
+- Avoid unnecessary state changes
+- Analyze the root cause of entities that aren't working, and prioritize automated solutions (like transport belts) above manual triage
+- Do not just pick up all entities and start from scratch, try to fix the modular small problems
+
+### Code Structure
+- Write code as direct Python interpreter commands
+- Only encapsulate reusable utility code into functions 
+- Use appropriate spacing and formatting
+
+## Understanding Output
+
+### Error Messages
+```stderr
+Error: 1: ("Initial Inventory: {...}")
+10: ("Error occurred in following lines...")
+```
+- Numbers indicate line of execution
+- Previous lines executed successfully
+- Fix errors at indicated line
+
+### Status Updates
+```stdout
+23: ('Resource collection completed...')
+78: ('Entities on map: [...]')
+```
+- Shows execution progress
+- Provides entity status
+- Lists warnings and conditions
+
+### Entity Status Checking
+- Monitor entity `warnings` field
+- Check entity `status` field
+- Verify resource levels
+- Track production states
+
+## Game Progression
+- Think about long term objectives, and break them down into smaller, manageable steps.
+- Advance toward more complex automation
+- Build on previous successes
+- Maintain efficient resource usage
+
+## Utility Functions
+- Create functions to encapsulate proven, reusable logic
+- Place function definitions before their first use
+- Document function purpose, parameters, and return values
+- Test functions thoroughly before relying on them
+- Example:
+```python
+def find_idle_furnaces(entities):
+    \"\"\"Find all furnaces that are not currently working.
+
+    Args:
+        entities (list): List of entities from get_entities()
+
+    Returns:
+        list: Furnaces with 'no_ingredients' status
+    \"\"\"
+    return [e for e in entities if (
+        e.name == 'stone-furnace' and 
+        e.status == EntityStatus.NO_INGREDIENTS
+    )]
+```
+
+## Data Structures
+- Use Python's built-in data structures to organize entities
+- Sets for unique entity collections:
+```python
+working_furnaces = {e for e in get_entities() 
+                   if e.status == EntityStatus.WORKING}
+```
+- Dictionaries for entity mapping:
+```python
+furnace_by_position = {
+    (e.position.x, e.position.y): e 
+    for e in get_entities() 
+    if isinstance(e, Furnace)
+}
+```
+- Lists for ordered operations:
+```python
+sorted_furnaces = sorted(
+    get_entities(),
+    key=lambda e: (e.position.x, e.position.y)
+)
+```
+
+## Important Notes
+- Use transport belts to keep burners fed with coal
+- Always inspect game state before making changes
+- Consider long-term implications of actions
+- Maintain working systems, and clear entities that aren't working or don't have a clear purpose
+- Build incrementally and verify each step
+- DON'T REPEAT YOUR PREVIOUS STEPS - just continue from where you left off. Take into account what was the last action that was executed and continue from there. If there was a error previously, do not repeat your last lines - as this will alter the game state unnecessarily.
+- Do not encapsulate your code in a function _unless_ you are writing a utility for future use - just write it as if you were typing directly into the Python interpreter.
+- Your inventory has space for ~2000 items. If it fills up, insert the items into a chest.
+- Ensure that your factory is arranged in a grid, as this will make things easier.
+"""
+
+FINAL_INSTRUCTION = "\n\nALWAYS WRITE VALID PYTHON. YOUR WEIGHTS WILL BE ERASED IF YOU DON'T USE PYTHON.\nWHEN YOU SEE AN ERROR, DO NOT CLEAN EVERYTHING UP. FIX ERRORS MODULARLY, NOT BY STARTING FRESH EVERY TIME. CREATE SMALL MODULAR POLICIES TO FIX ERRORS AS THEY COME" # Annoying how effective this is
+
+
+class QueryAgent(AgentABC):
+   def __init__(self, model, system_prompt_parts, task, *args, **kwargs):
+        instructions = GENERAL_INSTRUCTIONS+get_rag_system_prompt(system_prompt_parts)+FINAL_INSTRUCTION
+        self.task = task
+        instructions += f"\n\n### Goal\n{task.goal_description}\n\n"
+        super().__init__( model, instructions, *args, **kwargs)
+        self.llm_factory = LLMFactory(model)
+        self.formatter = RecursiveReportFormatter(chunk_size=16,llm_call=self.llm_factory.acall,cache_dir='summary_cache')
+        self.generation_params = GenerationParameters(n=1, max_tokens=4096, model=model)
+
+   async def step(self, conversation: Conversation, response: Response, namespace: FactorioNamespace) -> Policy:
+
+       # We format the conversation every N steps to add a context summary to the system prompt
+       formatted_conversation = await self.formatter.format_conversation(conversation, namespace)
+       # We set the new conversation state for external use
+       self.set_conversation(formatted_conversation)
+
+       return await self._get_policy(formatted_conversation)
+
+   @tenacity.retry(
+       retry=retry_if_exception_type(Exception),
+       wait=wait_exponential(multiplier=1, min=4, max=10)
+   )
+   async def _get_policy(self, conversation: Conversation):
+       response = await self.llm_factory.acall(
+           messages=self.formatter.to_llm_messages(conversation),
+           n_samples=1,  # We only need one program per iteration
+           temperature=self.generation_params.temperature,
+           max_tokens=self.generation_params.max_tokens,
+           model=self.generation_params.model,
+       )
+
+       policy = parse_response(response)
+       if not policy:
+           raise Exception("Not a valid Python policy")
+
+       return policy
+
+   async def end(self, conversation: Conversation, completion: CompletionResult):
+       pass
+
+

agents/utils/prompt_utils.py

@@ -0,0 +1,23 @@
+import re
+from typing import Optional, Tuple
+import ast
+
+
+def get_default_system_prompt(prompt_object):
+        # Combine all parts into final prompt
+        return (
+            f"```types\n{prompt_object['type_defs']}\n{prompt_object['entity_defs']}\n```\n"
+            f"```methods\n{prompt_object['schema']}\n```"
+            f"{prompt_object['manual_defs']}\n{prompt_object['examples']}\n"
+        )
+
+def get_rag_system_prompt(prompt_object):
+        # Combine all parts into final prompt
+        return (
+            f"```types\n{prompt_object['type_defs']}\n{prompt_object['entity_defs']}\n```\n"
+            f"```methods\n{prompt_object['schema']}\n```"
+            f"{prompt_object['manual_defs']}\n"
+            f"Information gathering tools available to you\n"
+            f"```rag_methods\n{prompt_object['rag_schema']}\n```"
+            f"{prompt_object['rag_manual_defs']}\n"
+        )

env/src/entities.py

@@ -520,7 +520,7 @@ class ElectricMiningDrill(MiningDrill, Electric):
     pass
 
 class BurnerInserter(Inserter, BurnerType):
-    """An inserter powered by burnable fuel."""
+    """An inserter powered by burnable fuel, does not require electricity, runs on burner fuels."""
     _height: float = 1
     _width: float = 1
     pass

env/src/instance.py

@@ -30,7 +30,8 @@
 from rcon.factorio_rcon import RCONClient
 from models.game_state import GameState
 from utils.controller_loader.system_prompt_generator import SystemPromptGenerator
-
+# import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
 CHUNK_SIZE = 32
 MAX_SAMPLES = 5000
 
@@ -190,7 +191,7 @@ def get_elapsed_ticks(self):
         if not response: return 0
         return int(response)
 
-    def get_system_prompt(self) -> str:
+    def get_system_prompt(self) -> dict:
         """
         Get the system prompt for the Factorio environment.
         This includes all the available actions, objects, and entities that the agent can interact with.
@@ -315,10 +316,11 @@ def wrapper(*args, **kwargs):
 
     def eval_with_error(self, expr, timeout=60):
         """ Evaluate an expression with a timeout, and return the result without error handling"""
-        # with ThreadPoolExecutor(max_workers=1) as executor:
-        #     future = executor.submit(self._eval_with_timeout, expr)
-        #     score, goal, result = future.result(timeout)
-        #     return score, goal, result
+        #with ThreadPoolExecutor(max_workers=1) as executor:
+        #    future = executor.submit(self.namespace.eval_with_timeout, expr)
+        #    score, goal, result = future.result(timeout)
+        #    return score, goal, result
+
         def handler(signum, frame):
             raise TimeoutError()