from simulator_env import SwarmAgent from huggingface_hub import hf_hub_download from llama_cpp import Llama import textwrap import re # Download only the behavior-tree model shard model_path = hf_hub_download( repo_id="Inventors-Hub/SwarmChat-models", repo_type="model", filename="Falcon3-10B-Instruct-BehaviorTree-3epochs.Q4_K_M.gguf", ) # llm = Llama(model_path=model_path, n_ctx=1024*4) llm = Llama( model_path=model_path, n_ctx=1024*4, # down from 4096 low_vram=True, # llama.cpp low-vram mode f16_kv=True, # half-precision kv cache use_mmap=True, # mmap file use_mlock=False, ) def call_behaviors() -> dict: behavior_dict = {} for name, attribute in SwarmAgent.__dict__.items(): if callable(attribute) and not name.startswith("_") \ and not name.startswith("update") and not name.startswith("obstacle"): doc = attribute.__doc__ if doc is not None: # Dedent, strip, and join into one line by replacing newlines and tabs cleaned_doc = " ".join(textwrap.dedent(doc).strip().split()) else: cleaned_doc = "" behavior_dict[name] = cleaned_doc return behavior_dict def extract_behavior_tree(response: str) -> str: """ Extracts an XML behavior tree from the given response text. Looks for a block of XML enclosed in tags. """ pattern = re.compile(r'()', re.DOTALL) match = pattern.search(response) if match: return match.group(1).strip() else: # If no valid XML block is found, return the original response. return response.strip() def save_behavior_tree(tree_xml: str, file_name: str = "tree.xml") -> None: """ Saves the behavior tree XML to a file. """ with open(file_name, "w", encoding="utf-8") as f: f.write(tree_xml) def construct_prompt(prompt: str, prompt_type: str="one") -> str: behaviors = call_behaviors() behaviors_text = "\n".join(f"{name}: {doc}" for name, doc in behaviors.items()) plan_prompt = f""" <>You are a helpful, respectful, and honest AI assistant. Your task is to generate well-structured XML code for behavior trees based on the provided instructions.<> INSTRUCTIONS: It is CRITICAL to use only the following behaviors structured as a dictionary: {behaviors_text} to construct behavior tree in XML format for the following command. Including any behavior that is not in the provided dictionary can result in damage to the agents and potentially humans, therefore you are not allowed to do so. AVOID AT ALL COSTS. USER COMMAND: generate behavior tree to "{prompt}". Take a step back and think deeply about the behavior you need for this command. Consider the XML structure and the behaviors you use. The output MUST follow this XML structure exactly, including: - A root element with element with an inner structure of Sequences, Fallback, Conditions, and Actions. - A section listing all node models. - No additional text or commentary outside the XML. Output only the XML behavior tree without extra text. OUTPUT: """ if prompt_type == "zero": return plan_prompt elif prompt_type == "one": path = "Prompt_One_shot.txt" with open(path, "r", encoding="utf-8") as file: file_content = file.read() return f"{file_content} {plan_prompt}" elif prompt_type == "two": path = "Prompt_Two_Shot.txt" with open(path, "r", encoding="utf-8") as file: file_content = file.read() return f"{file_content} {plan_prompt}" else: raise ValueError("Unknown prompt type provided.") def generate_behavior_tree(task_prompt: str) -> str: prompt = construct_prompt(task_prompt) print("\n\n",prompt,"\n\n") output = llm( prompt, temperature=0, max_tokens=1024, top_p=0.95, top_k=50, repeat_penalty=1.1 ) response = output.get("choices", [{}])[0].get("text", "").strip() tree_xml = extract_behavior_tree(response) save_behavior_tree(tree_xml) print("\n response: \n", response) return tree_xml # Example usage: if __name__ == "__main__": task = "Generate a behavior tree to just form a line." response = generate_behavior_tree(task) print("Generated behavior tree response:") print(response)