ConversationEntityMemory

Purpose

Tracks named entities (people, places, organizations, or "important concepts") mentioned across a conversation, accumulating context snippets per entity so that later turns mentioning the same entity can retrieve everything previously said about it. Suitable for conversations with many recurring named referents (CRM-style assistants, relationship-tracking bots, etc.).

Constructor

def __init__(
    self,
    llm=None,
    memory_key: str = "entity_info",
    input_key: str = "input",
    output_key: str = "output",
    lang: str = "en",
    entity_extraction_prompt: str = None,
)

Public methods

save_context(inputs, outputs) -> None

1. Combines input_text + " " + output_text.
2. Calls _extract_entities(combined_text).
3. For each extracted entity name not already in self.entity_store, creates an Entity(name=entity_name, entity_type="unknown") — entity_type is always the literal string "unknown"; nothing in this class ever classifies entity type (person/place/organization), despite the docstring/prompt implying such categories exist.
4. Calls entity.add_context(combined_text) on every extracted entity (new or existing) — the entire combined turn text, not just the sentence mentioning the entity, is stored as that entity's context snippet. If a turn mentions 3 entities, all 3 get the same full-turn text appended.
5. Appends {"input", "output", "entities"} to self.chat_memory.

_extract_entities(text: str) -> List[str] (internal)

LLM path (if self.llm is set): formats entity_extraction_prompt.format(text=text), calls self.llm.generate(prompt, max_tokens=100), then splits the response on commas, stripping whitespace and discarding any resulting fragment <= 2 characters long.

Fallback path (if self.llm is None): lazily imports stanza and constructs stanza.Pipeline(lang=self.lang, processors='tokenize,ner') on first use (this can be slow — stanza downloads/loads language models). Filters doc.ents to those that are either alphabetic or longer than 3 characters.

Both paths catch all exceptions, log an error, and return [] on failure — extraction never raises to the caller.

[!WARNING] stanza fallback requires a one-time model download stanza.Pipeline(...) will attempt to download language models on first construction unless you've already run stanza.download(lang) and the models are cached locally. In an offline or restricted-network environment, the fallback path's first call may hang or raise inside the try/except (caught and logged, silently returning []) — you won't get an explicit signal that entity extraction is non-functional beyond an empty entities list and a logged error.

load_memory_variables(inputs) -> Dict[str, Any]

1. Extracts entities from the current input text only (inputs.get(self.input_key, "") — the assistant output is not considered here, unlike save_context).
2. For each extracted entity found in self.entity_store, retrieves its last 3 contexts (get_recent_contexts(limit=3)) and formats:

f"معلومات سابقة عن {entity_name}:\n{info}"

[!DANGER] Hardcoded Arabic prefix regardless of lang This string — "معلومات سابقة عن" ("previous information about") — is a hardcoded literal, not conditioned on self.lang in any way. Even if you construct ConversationEntityMemory(lang="en") and use an English-only LLM for extraction, every retrieved entity context returned to your prompt template will be prefixed with this Arabic phrase. If your downstream LLM or prompt template assumes English-only output, you'll get an unexpected mixed-language string. See Known Issues.

3. Joins all entity blocks with "\n\n" and returns {self.memory_key: entity_info} (empty string if no entities matched).

clear() -> None

Clears both entity_store and chat_memory.

get_memory_stats() -> Dict[str, Any]

Returns type, entity_count, message_count, total_contexts (sum of len(entity.contexts) across all entities).

get_entity_info(entity_name: str) -> Optional[Dict[str, Any]]

Returns None if the entity isn't in entity_store; otherwise a dict with name, type, contexts_count, first_seen, last_seen, recent_contexts (via get_recent_contexts(), default limit=3). Lookup is by exact string match against however the entity name was extracted — no normalization (casing, whitespace) is applied, so "Sarah" and "sarah" would be tracked as two separate entities if extraction is inconsistent.

list_entities() -> List[str]

Returns list(self.entity_store.keys()).

Async methods

asave_context / aload_memory_variables — local asyncio.to_thread(...) wrappers.

Raises

_extract_entities never raises (catches internally). Public methods generally don't raise beyond standard dict/str errors from malformed inputs.

Internal workflow

Thread safety

Not thread-safe. No locking around entity_store/chat_memory mutation, and the lazily-constructed self._nlp (stanza pipeline) is not guarded against concurrent first-initialization from multiple threads.

Complexity

• save_context: O(extraction cost) + O(e) where e = number of entities in the turn.
• load_memory_variables: O(extraction cost) + O(e × c) where c = contexts per matched entity (capped at 3 by get_recent_contexts).
• Entity lookup (entity_store dict): O(1) average.

Examples

from zerda.memory import ConversationEntityMemory

class FakeLLM:
    def generate(self, prompt: str, max_tokens: int = 256) -> str:
        return "Sarah, Acme Corp"

# 1. Basic
mem = ConversationEntityMemory(llm=FakeLLM())
mem.save_context(
    {"input": "Sarah just joined Acme Corp as an engineer."},
    {"output": "Got it, noted!"},
)

# 2. Intermediate — retrieve stored context on a later, related turn
result = mem.load_memory_variables({"input": "What do we know about Sarah?"})
print(result["entity_info"])
# "معلومات سابقة عن Sarah:\nSarah just joined Acme Corp as an engineer. Got it, noted!"
# (note: entire turn text is stored per entity, and the Arabic prefix appears
#  regardless of the `lang` you configured)

# 3. Advanced — introspect a specific entity
info = mem.get_entity_info("Sarah")
print(info["contexts_count"], info["first_seen"])

# 4. Error handling — extraction failure never raises, just returns no entities
class BrokenLLM:
    def generate(self, prompt, max_tokens=256):
        raise RuntimeError("timeout")

mem2 = ConversationEntityMemory(llm=BrokenLLM())
mem2.save_context({"input": "Sarah again"}, {"output": "ok"})
print(mem2.list_entities())  # [] — the error was logged, not raised

# 5. Production — without an LLM, relies on stanza (must be pre-downloaded)
# mem3 = ConversationEntityMemory(lang="en")  # no llm -> stanza fallback
# mem3.save_context({"input": "Barack Obama visited Paris."}, {"output": "Interesting!"})

Best practices

• Always supply an llm in production unless you've verified stanza and its language models are available in your deployment environment.
• If your application is not Arabic-facing, be prepared to post-process or strip the hardcoded Arabic prefix from load_memory_variables's output, or don't rely on this class if a clean English-only string is required.
• Normalize entity names yourself (e.g. lowercase + strip) before comparing against list_entities() output, since this class does no normalization internally.

Common mistakes

• Assuming entity_type reflects an actual classification (person/place/org) — it is always "unknown".
• Assuming lang affects LLM-based extraction — it only affects the stanza fallback path.
• Assuming per-entity contexts are scoped to the sentence mentioning that entity — they're the full combined turn text.

Related APIs

• Entity — the record type stored per tracked entity.
• MemoryCompressor — for a different approach to keeping long-running context compact.