Spaces:
Runtime error
Runtime error
File size: 3,533 Bytes
4962437 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import time
from typing import Any, List
from swarms.models.prompts.agent_prompt_generator import get_prompt
class TokenUtils:
@staticmethod
def count_tokens(text: str) -> int:
return len(text.split())
class PromptConstructor:
def __init__(self, ai_name: str, ai_role: str, tools):
self.ai_name = ai_name
self.ai_role = ai_role
self.tools = tools
def construct_full_prompt(self, goals: List[str]) -> str:
prompt_start = (
"""Your decisions must always be made independently
without seeking user assistance.\n
Play to your strengths as an LLM and pursue simple
strategies with no legal complications.\n
If you have completed all your tasks, make sure to
use the "finish" command."""
)
# Construct full prompt
full_prompt = (
f"You are {self.ai_name}, {self.ai_role}\n{prompt_start}\n\nGOALS:\n\n"
)
for i, goal in enumerate(goals):
full_prompt += f"{i+1}. {goal}\n"
full_prompt += f"\n\n{get_prompt(self.tools)}"
return full_prompt
class Message:
content: str
def count_tokens(self) -> int:
return TokenUtils.count_tokens(self.content)
def format_content(self) -> str:
return self.content
class SystemMessage(Message):
pass
class HumanMessage(Message):
pass
class MessageFormatter:
send_token_limit: int = 4196
def format_messages(self, **kwargs: Any) -> List[Message]:
prompt_constructor = PromptConstructor(ai_name=kwargs["ai_name"],
ai_role=kwargs["ai_role"],
tools=kwargs["tools"])
base_prompt = SystemMessage(content=prompt_constructor.construct_full_prompt(kwargs["goals"]))
time_prompt = SystemMessage(
content=f"The current time and date is {time.strftime('%c')}"
)
used_tokens = base_prompt.count_tokens() + time_prompt.count_tokens()
memory: VectorStoreRetriever = kwargs["memory"]
previous_messages = kwargs["messages"]
relevant_docs = memory.get_relevant_documents(str(previous_messages[-10:]))
relevant_memory = [d.page_content for d in relevant_docs]
relevant_memory_tokens = sum(
[TokenUtils.count_tokens(doc) for doc in relevant_memory]
)
while used_tokens + relevant_memory_tokens > 2500:
relevant_memory = relevant_memory[:-1]
relevant_memory_tokens = sum(
[TokenUtils.count_tokens(doc) for doc in relevant_memory]
)
content_format = (
f"This reminds you of these events "
f"from your past:\n{relevant_memory}\n\n"
)
memory_message = SystemMessage(content=content_format)
used_tokens += memory_message.count_tokens()
historical_messages: List[Message] = []
for message in previous_messages[-10:][::-1]:
message_tokens = message.count_tokens()
if used_tokens + message_tokens > self.send_token_limit - 1000:
break
historical_messages = [message] + historical_messages
used_tokens += message_tokens
input_message = HumanMessage(content=kwargs["user_input"])
messages: List[Message] = [base_prompt, time_prompt, memory_message]
messages += historical_messages
messages.append(input_message)
return messages
|