Source code for langchain.memory.entity
import logging
from abc import ABC, abstractmethod
from itertools import islice
from typing import Any, Dict, Iterable, List, Optional
from langchain.chains.llm import LLMChain
from langchain.memory.chat_memory import BaseChatMemory
from langchain.memory.prompt import (
    ENTITY_EXTRACTION_PROMPT,
    ENTITY_SUMMARIZATION_PROMPT,
)
from langchain.memory.utils import get_prompt_input_key
from langchain.pydantic_v1 import BaseModel, Field
from langchain.schema import BasePromptTemplate
from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.messages import BaseMessage, get_buffer_string
from langchain.utilities.redis import get_client
logger = logging.getLogger(__name__)
[docs]class BaseEntityStore(BaseModel, ABC):
    """Abstract base class for Entity store."""
[docs]    @abstractmethod
    def get(self, key: str, default: Optional[str] = None) -> Optional[str]:
        """Get entity value from store."""
        pass
[docs]    @abstractmethod
    def set(self, key: str, value: Optional[str]) -> None:
        """Set entity value in store."""
        pass
[docs]    @abstractmethod
    def delete(self, key: str) -> None:
        """Delete entity value from store."""
        pass
[docs]    @abstractmethod
    def exists(self, key: str) -> bool:
        """Check if entity exists in store."""
        pass
[docs]    @abstractmethod
    def clear(self) -> None:
        """Delete all entities from store."""
        pass
[docs]class InMemoryEntityStore(BaseEntityStore):
    """In-memory Entity store."""
    store: Dict[str, Optional[str]] = {}