Merge branch 'main' of https://github.com/AIRI-Institute/AriGraph int…

…o main

Conflicts:
	graphs/parent_graph.py
	requirements.txt

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 AIRI - Artificial Intelligence Research Institute
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,54 @@
+# AriGraph: Learning Knowledge Graph World Models with Episodic Memory for LLM Agents
+
+AriGraph functions as the external memory architecture for large language models (LLMs), featuring a knowledge graph that is built from the ground up. This memory, configured as a semantic knowledge graph with added episodic vertices and edges, greatly improves the performance of Retrieval-Augmented Generation (RAG) in text-based games. Currently, AriGraph is a key component of the Ariadne agent, crafted to navigate the challenges of text-based games within the [TextWorld](https://github.com/microsoft/TextWorld) framework. This agent markedly surpasses all pre-existing baselines in these scenarios and showcases exceptional scalability in more expansive environments. For more detailed information about AriGraph and the Ariadne agent, please refer to the accompanying [paper](https://arxiv.org/abs/2407.04363). Experience the games used to evaluate our agent by following the provided [link](http://158.255.5.225/).
+
+![**Ariadne agent and its results**](img/Architecture.png?raw=True)
+
+## Performance
+We implemented five TextWorld environments for three distinct tasks: Treasure Hunt, Cleaning, and Cooking. The Treasure Hunt task requires navigating a maze and searching for treasure, while the Cleaning task involves tidying up a house by placing items in their designated spots. The Cooking task focuses on gathering ingredients and preparing a meal. Each LLM agent tested had the same decision-making module, differing only in memory implementation. We reported average human scores across all runs and for the top-3 performing runs. The table below presents the mean normalized game scores: 
+Type of memory | Treasure Hunt | Cleaning | Cooking | Treasure Hunt Hard | Cooking Hard
+-- | -- | -- | -- | -- | -- 
+AriGraph (ours) | 1.0 | 0.79 | 1.0 | 1.0 | 1.0
+Human Players Top-3 | 1.0 | 0.85 | 1.0 | - | -
+Human Players All | 0.96 | 0.59 | 0.32 | - | -
+Full History | 0.49 | 0.05 | 0.18 | - | -
+Summary | 0.33 | 0.39 | 0.52 | 0.17 | 0.21
+RAG | 0.33 | 0.35 | 0.36 | 0.17 | 0.17
+
+## Requirements
+Due to dependencies required by TextWorld, our code can only be executed on Linux systems, specifically after installing certain system libraries.
+On a Debian/Ubuntu-based system, these can be installed with
+
+    sudo apt update && sudo apt install build-essential libffi-dev python3-dev curl git
+
+And on macOS, with
+
+    brew install libffi curl git
+
+To complete requirements installation, you need Python 3.11+ and to run 
+
+    pip install -r requirements.txt
+
+## Repository structure
+- **agents** contains GPTagent.
+- **envs** contains TextWorld files for environment loading.
+- **graphs** contains TripletGraph in parent_graph.py and other graphs which inherit it.
+- **logs** contains logs of every reported run of our agent.
+- **prompts** contains prompts used in pipelines.
+- **src** and **utils** contains service classes and functions.
+
+The other Python files contain a pipeline for each agent as well as code for running the game in interactive mode in the console. Each pipeline includes a highlighted, changeable section with parameters for execution.
+
+## Citation
+If you find our work useful, please cite the [AriGraph](https://arxiv.org/abs/2407.04363):
+```
+@misc{anokhin2024arigraphlearningknowledgegraph,
+      title={AriGraph: Learning Knowledge Graph World Models with Episodic Memory for LLM Agents}, 
+      author={Petr Anokhin and Nikita Semenov and Artyom Sorokin and Dmitry Evseev and Mikhail Burtsev and Evgeny Burnaev},
+      year={2024},
+      eprint={2407.04363},
+      archivePrefix={arXiv},
+      primaryClass={cs.AI},
+      url={https://arxiv.org/abs/2407.04363}, 
+}
+```

graphs/parent_graph.py

@@ -15,6 +15,41 @@ def __init__(self, model, system_prompt, api_key):
             api_key=api_key,
         )
 
+    def generate(self, prompt, jsn = False, t = 0.7):
+        if jsn:   
+            chat_completion = self.client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": self.system_prompt,
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                model=self.model,
+                response_format={"type": "json_object"},
+                temperature=t
+            )
+        else:
+            chat_completion = self.client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": self.system_prompt,
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                model=self.model,
+                temperature=t
+            )
+        response = chat_completion.choices[0].message.content
+        prompt_tokens = chat_completion.usage.prompt_tokens
+        completion_tokens = chat_completion.usage.completion_tokens
     def generate(self, prompt, jsn = False, t = 0.7):
         if jsn:   
             chat_completion = self.client.chat.completions.create(
@@ -51,6 +86,7 @@ def generate(self, prompt, jsn = False, t = 0.7):
         prompt_tokens = chat_completion.usage.prompt_tokens
         completion_tokens = chat_completion.usage.completion_tokens
 
+        cost = completion_tokens * 3 / 100000 + prompt_tokens * 1 / 100000
         cost = completion_tokens * 3 / 100000 + prompt_tokens * 1 / 100000
         self.total_amount += cost
         return response, cost
@@ -205,4 +241,4 @@ def print_graph(self):
         print("Triplets in the graph:")
         for triplet_str in self.get_all_triplets():
             print(triplet_str)
-
+

pipeline_arigraph.py

@@ -178,6 +178,7 @@ def choose_action(observations, observation, subgraph, top_episodic, plan0, all_
 \n2. History of {n_prev} last observations and actions: {observations} 
 \n3. Your current observation: {observation}
 \n4. Information from the memory module that can be relevant to current situation:  {subgraph}
+\n5. Your {topk_episodic} most relevant episodic memories from the past for the current situation: {top_episodic}.
 \n6. Your current plan: {plan0}'''
 
     if if_explore:
@@ -206,6 +207,7 @@ def planning(observations, observation, plan0, subgraph, top_episodic, if_explor
 \n2. History of {n_prev} last observations and actions: {observations} 
 \n3. Your current observation: {observation}
 \n4. Information from the memory module that can be relevant to current situation: {subgraph}
+\n5. Your {topk_episodic} most relevant episodic memories from the past for the current situation: {top_episodic}.
 \n6. Your previous plan: {plan0}'''
 
     if if_explore:
@@ -227,4 +229,4 @@ def get_unexpl_exits(locations, graph):
 
 
 if __name__ == "__main__":
-    run()
+    run()