From bc66b3fb8d1e870cb17c9217b5e6643151f23251 Mon Sep 17 00:00:00 2001 From: Aditi Viswanathan Date: Wed, 31 May 2023 17:32:19 -0700 Subject: [PATCH] make BaseEntityStore inherit from BaseModel (#5478) # Make BaseEntityStore inherit from BaseModel This enables initializing InMemoryEntityStore by optionally passing in a value for the store field. ## Who can review? It's a small change so I think any of the reviewers can review, but tagging @dev2049 who seems most relevant since the change relates to Memory. --- langchain/memory/entity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langchain/memory/entity.py b/langchain/memory/entity.py index a7808e15..27f903b1 100644 --- a/langchain/memory/entity.py +++ b/langchain/memory/entity.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from itertools import islice from typing import Any, Dict, Iterable, List, Optional -from pydantic import Field +from pydantic import BaseModel, Field from langchain.base_language import BaseLanguageModel from langchain.chains.llm import LLMChain @@ -19,7 +19,7 @@ from langchain.schema import BaseMessage, get_buffer_string logger = logging.getLogger(__name__) -class BaseEntityStore(ABC): +class BaseEntityStore(BaseModel, ABC): @abstractmethod def get(self, key: str, default: Optional[str] = None) -> Optional[str]: """Get entity value from store."""