fix(persona): correct typing, formatting, and improve docstrings

This commit is contained in:
Mikolaj Mikuliszyn 2025-04-27 18:00:14 +01:00
parent d6a032d8e7
commit 77196f40ee

View File

@ -2,12 +2,37 @@ import re
from typing import Any, Callable, Optional from typing import Any, Callable, Optional
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from langchain_core.memory import BaseMemory from langchain_core.memory import BaseMemory
class EnrichedMessage(BaseModel): class EnrichedMessage(BaseModel):
""" """A message enriched with persona traits and metadata.
A message enriched with persona traits and metadata.
This class extends a basic message with additional information about the
sender's personality traits and any relevant metadata. The traits are
automatically detected from the message content.
Examples:
Basic usage:
>>> message = EnrichedMessage(
... id="msg-001",
... content="I'm really excited about this project!",
... traits=["enthusiastic"], # Automatically detected
... metadata={"traits_count": {"enthusiastic": 1}}
... )
>>> message.traits
['enthusiastic']
With multiple traits:
>>> message = EnrichedMessage(
... id="msg-002",
... content="I'm not sure about this... maybe we should reconsider?",
... traits=["hesitant", "cautious"], # Multiple detected traits
... metadata={"traits_count": {"hesitant": 1, "cautious": 1}}
... )
>>> message.traits
['hesitant', 'cautious']
""" """
id: str id: str
@ -17,13 +42,42 @@ class EnrichedMessage(BaseModel):
class PersonaMemory(BaseMemory): class PersonaMemory(BaseMemory):
""" """Memory module that dynamically tracks emotional and behavioral traits
Memory module that dynamically tracks emotional and behavioral traits
exhibited by an agent over the course of a conversation. exhibited by an agent over the course of a conversation.
Traits are automatically detected from output messages and Traits are automatically detected from output messages and accumulated across
accumulated across interactions. Recent messages and detected traits interactions. Recent messages and detected traits can be retrieved to enrich
can be retrieved to enrich future prompts or modify agent behavior. future prompts or modify agent behavior.
The memory maintains a sliding window of the k most recent messages
(default k=10) and tracks all unique traits detected across these messages.
Examples:
Basic usage:
>>> memory = PersonaMemory()
>>> memory.save_context(
... {"input": "How are you?"},
... {"output": "I'm doing great! This is fantastic!"}
... )
>>> memory.load_memory_variables({})
{'persona': {'traits': ['enthusiastic']}}
Including recent messages:
>>> memory = PersonaMemory()
>>> memory.save_context(
... {"input": "What do you think?"},
... {"output": "I'm not sure... perhaps we should reconsider?"}
... )
>>> memory.load_memory_variables({}, include_messages=True)
{'persona': {
'traits': ['hesitant', 'cautious'],
'recent_messages': [{
'id': '1',
'content': "I'm not sure... perhaps we should reconsider?",
'traits': ['hesitant', 'cautious'],
'metadata': {'traits_count': {'hesitant': 1, 'cautious': 1}}
}]
}}
""" """
memory_key: str = "persona" memory_key: str = "persona"