diff --git a/lost_in_the_middle.py b/lost_in_the_middle.py index 779cb24..077a29f 100644 --- a/lost_in_the_middle.py +++ b/lost_in_the_middle.py @@ -1,31 +1,51 @@ -from cat.mad_hatter.decorators import tool, hook, plugin -from pydantic import BaseModel -from datetime import datetime, date - -class MySettings(BaseModel): - required_int: int - optional_int: int = 69 - required_str: str - optional_str: str = "meow" - required_date: date - optional_date: date = 1679616000 - -@plugin -def settings_schema(): - return MySettings.schema() - -@tool -def get_the_day(tool_input, cat): - """Get the day of the week. Input is always None.""" - - dt = datetime.now() - - return dt.strftime('%A') - -@hook -def before_cat_sends_message(message, cat): - - prompt = f'Rephrase the following sentence in a grumpy way: {message["content"]}' - message["content"] = cat.llm(prompt) - - return message +from cat.mad_hatter.decorators import hook + +def litm(documents): + """ + Function based on Haystack's LITM ranker: + https://github.com/deepset-ai/haystack/blob/main/haystack/nodes/ranker/lost_in_the_middle.py + + Lost In The Middle is based on the paper https://arxiv.org/abs/2307.03172 + Check it for mor details. + + + Parameters + ---------- + documents: List of documents (the declarative working memories) + + Returns + ---------- + litm_docs: The same list but reordered + """ + if len(documents) == 1: + return documents + + document_index = list(range(len(documents))) + lost_in_the_middle_indices = [0] + + for doc_idx in document_index[1:]: + insertion_index = len(lost_in_the_middle_indices) // 2 + len(lost_in_the_middle_indices) % 2 + lost_in_the_middle_indices.insert(insertion_index, doc_idx) + litm_docs = [documents[idx] for idx in lost_in_the_middle_indices] + return litm_docs + + +@hook(priority=1) +def after_cat_recalls_memories(cat) -> None: + """Hook after semantic search in memories. + + The hook is executed just after the Cat searches for the meaningful context in memories + and stores it in the *Working Memory*. + + Parameters + ---------- + cat : CheshireCat + Cheshire Cat instance. + + """ + if cat.working_memory['declarative_memories']: + litm_docs = litm(cat.working_memory['declarative_memories']) + cat.working_memory['declarative_memories'] = litm_docs + else: + print("#HicSuntGattones") + pass # do nothing diff --git a/plugin.json b/plugin.json index 8fd7571..cd72c31 100644 --- a/plugin.json +++ b/plugin.json @@ -2,9 +2,8 @@ "name": "Lost In The Middle", "version": "0.0.1", "description": "Description of lost_in_the_middle.", - "author_name": "Me", - "author_url": "https://mywebsite.me", - "plugin_url": "https://github.com/my_name/lost_in_the_middle", - "tags": "cat, template, example", - "thumb": "https://raw.githubusercontent.com/my_repo_path/lost_in_the_middle.png" + "author_name": "nickprock", + "plugin_url": "https://github.com/nickprock/lost_in_the_middle", + "tags": "vector-db, embeddings, lost_in_the_middle", + "thumb": "https://raw.githubusercontent.com/nickprock/lost-in-the-middle/main/stregatto_space.png" } diff --git a/stregatto_space.png b/stregatto_space.png new file mode 100644 index 0000000..7eede8a Binary files /dev/null and b/stregatto_space.png differ