Fixes after review

deepsense-ai · Oct 15, 2024 · f1094b8 · f1094b8
1 parent fb38633
commit f1094b8
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/packages/ragbits-document-search/src/ragbits/document_search/_main.py b/packages/ragbits-document-search/src/ragbits/document_search/_main.py
@@ -6,7 +6,7 @@
 from ragbits.core.vector_store import VectorStore
 from ragbits.document_search.documents.document import Document, DocumentMeta
 from ragbits.document_search.documents.element import Element
-from ragbits.document_search.documents.sources import Source
+from ragbits.document_search.documents.sources import GCSSource, LocalFileSource, Source
 from ragbits.document_search.ingestion.document_processor import DocumentProcessorRouter
 from ragbits.document_search.ingestion.providers.base import BaseProvider
 from ragbits.document_search.retrieval.rephrasers.base import QueryRephraser
@@ -80,7 +80,9 @@ async def search(self, query: str, search_config: SearchConfig = SearchConfig())
         return self.reranker.rerank(elements)
 
     async def ingest_document(
-        self, document: Union[DocumentMeta, Document, Source], document_processor: Optional[BaseProvider] = None
+        self,
+        document: Union[DocumentMeta, Document, Union[LocalFileSource, GCSSource]],
+        document_processor: Optional[BaseProvider] = None,
     ) -> None:
         """
         Ingest a document.
@@ -92,8 +94,7 @@ async def ingest_document(
         """
 
         if isinstance(document, Source):
-            local_path = await document.fetch()
-            document_meta = DocumentMeta.from_local_path(local_path)
+            document_meta = await DocumentMeta.from_source(document)
         elif isinstance(document, DocumentMeta):
             document_meta = document
         else:

diff --git a/packages/ragbits-document-search/src/ragbits/document_search/documents/document.py b/packages/ragbits-document-search/src/ragbits/document_search/documents/document.py
@@ -97,6 +97,24 @@ def from_local_path(cls, local_path: Path) -> "DocumentMeta":
             source=LocalFileSource(path=local_path),
         )
 
+    @classmethod
+    async def from_source(cls, source: Union[LocalFileSource, GCSSource]) -> "DocumentMeta":
+        """
+        Create a document metadata from a source.
+
+        Args:
+            source: The source from which the document is fetched.
+
+        Returns:
+            The document metadata.
+        """
+        path = await source.fetch()
+
+        return cls(
+            document_type=DocumentType(path.suffix[1:]),
+            source=source,
+        )
+
 
 class Document(BaseModel):
     """