prepend system memories for plugin

This commit is contained in:
Kurt Heiden
2024-10-28 21:45:25 -06:00
parent cb6d5540cb
commit fb2460f907
2 changed files with 15 additions and 4 deletions
+3
View File
@@ -46,6 +46,9 @@ class Conversation(SMBaseModel):
def __str__(self):
return f"<Conversation id={self.id!r}>"
def prepend_system_message(self, role: str, text: str, meta: Dict[str, Any] = {}):
self.messages = [Message(role=role, text=text, meta=meta)] + self.messages
def add_message(self, role: str, text: str, meta: Dict[str, Any] = {}):
"""Add a new message to the conversation."""
self.messages.append(Message(role=role, text=text, meta=meta))
+12 -4
View File
@@ -6,6 +6,7 @@ class TestOllama(unittest.TestCase):
def test_generate_text(self):
result = sm.generate_text(prompt="What is the meaning of life?", llm_provider="ollama", llm_model="llama3.2")
self.assertGreater(len(result), 0)
self.assertIsNotNone(result)
def test_create_conversation(self):
@@ -13,14 +14,15 @@ class TestOllama(unittest.TestCase):
conversation.add_message("user", "Remember the number 42.")
result = conversation.send()
self.assertIsNotNone(result)
self.assertGreaterEqual(len(result.text), 0)
self.assertIsInstance(result, sm.models.Message)
# @unittest.skip('Not working...')
def test_memory(self):
conversation = sm.create_conversation(llm_provider="ollama", llm_model="llama3.2")
class SimpleMemoryPlugin:
def __init__(self):
self.memories = [
"the earth has fictionally beeen destroyed.",
"the earth has fictionally been destroyed.",
"the moon is made of cheese.",
]
@@ -29,15 +31,21 @@ class TestOllama(unittest.TestCase):
def send_hook(self, conversation: sm.Conversation):
for m in self.yield_memories():
conversation.add_message(role="system", text=m)
conversation.prepend_system_message(role="system", text=m)
conversation.add_plugin(SimpleMemoryPlugin())
conversation = sm.create_conversation(llm_provider="ollama", llm_model="llama3.2")
conversation.add_message(
role="user",
text="Write a poem about the moon",
)
self.assertGreater(len(conversation.messages), 0)
conversation.add_plugin(SimpleMemoryPlugin())
result = conversation.send()
self.assertGreater(len(conversation.messages), 2)
self.assertIsNotNone(result)
self.assertIsNotNone(result.text)
self.assertGreater(len(result.text), 0)
self.assertIsInstance(result, sm.models.Message)
def test_structure_response(self):