Add on_planning hook.

This commit is contained in:
BillSchumacher
2023-04-15 23:01:01 -05:00
parent 68e26bf9d6
commit 09a5b3149d
2 changed files with 23 additions and 2 deletions

View File

@@ -65,6 +65,7 @@ class Agent:
# Send message to AI, get response
with Spinner("Thinking... "):
assistant_reply = chat_with_ai(
self,
self.prompt,
self.user_input,
self.full_message_history,

View File

@@ -51,7 +51,7 @@ def generate_context(prompt, relevant_memory, full_message_history, model):
# TODO: Change debug from hardcode to argument
def chat_with_ai(
prompt, user_input, full_message_history, permanent_memory, token_limit
agent, prompt, user_input, full_message_history, permanent_memory, token_limit
):
"""Interact with the OpenAI API, sending the prompt, user input, message history,
and permanent memory."""
@@ -109,7 +109,7 @@ def chat_with_ai(
current_tokens_used += token_counter.count_message_tokens(
[create_chat_message("user", user_input)], model
) # Account for user input (appended later)
while next_message_to_add_index >= 0:
# print (f"CURRENT TOKENS USED: {current_tokens_used}")
message_to_add = full_message_history[next_message_to_add_index]
@@ -135,6 +135,26 @@ def chat_with_ai(
# Append user input, the length of this is accounted for above
current_context.extend([create_chat_message("user", user_input)])
plugin_count = len(cfg.plugins)
for i, plugin in enumerate(cfg.plugins):
plugin_response = plugin.on_planning(
agent.prompt_generator, current_context
)
if not plugin_response or plugin_response == "":
continue
tokens_to_add = token_counter.count_message_tokens(
[plugin_response], model
)
if current_tokens_used + tokens_to_add > send_token_limit:
if cfg.debug_mode:
print("Plugin response too long, skipping:",
plugin_response)
print("Plugins remaining at stop:", plugin_count - i)
break
current_context.append(
create_chat_message("system", plugin_response)
)
# Calculate remaining tokens
tokens_remaining = token_limit - current_tokens_used
# assert tokens_remaining >= 0, "Tokens remaining is negative.