refactor: update document store import to use BaseDocumentStore in DocumentContextExtractor (#19466)

Author: Kaushal-26

llama-index-core/llama_index/core/extractors/document_context.py

@@ -30,7 +30,7 @@
     DocumentBlock,
 )
 from llama_index.core.schema import BaseNode, Node, TextNode
-from llama_index.core.storage.docstore.simple_docstore import DocumentStore
+from llama_index.core.storage.docstore.types import BaseDocumentStore
 
 
 def is_text_node(node: BaseNode) -> TypeGuard[Union[Node, TextNode]]:
@@ -64,7 +64,7 @@ class DocumentContextExtractor(BaseExtractor):
 
     Attributes:
         llm (LLM): Language model instance for generating context
-        docstore (DocumentStore): Storage for parent documents
+        docstore (BaseDocumentStore): Storage for parent documents
         key (str): Metadata key for storing extracted context
         prompt (str): Prompt template for context generation
         doc_ids (Set[str]): Set of processed document IDs
@@ -87,7 +87,7 @@ class DocumentContextExtractor(BaseExtractor):
 
     # Pydantic fields
     llm: LLM
-    docstore: DocumentStore
+    docstore: BaseDocumentStore
     key: str
     prompt: str
     doc_ids: Set[str]
@@ -103,7 +103,7 @@ class DocumentContextExtractor(BaseExtractor):
 
     def __init__(
         self,
-        docstore: DocumentStore,
+        docstore: BaseDocumentStore,
         llm: Optional[LLM] = None,
         max_context_length: int = 1000,
         key: str = DEFAULT_KEY,