Add RawDocumentRequest

Signed-off-by: Hemslo Wang <[email protected]>

Author: hemslo

app/dependencies/html_preprocessor.py

@@ -1,6 +1,7 @@
 from bs4 import BeautifulSoup
 
-from app.models.document_model import DocumentMetadataModel, DocumentModel
+from app.models.document_metadata_model import DocumentMetadataModel
+from app.models.document_model import DocumentModel
 from app.models.html_document_request import HTMLDocumentRequest
 
 

app/dependencies/repository.py

@@ -8,6 +8,7 @@
 from app.dependencies.redis import RedisDep
 from app.models.document_model import DocumentModel
 from app.models.html_document_request import HTMLDocumentRequest
+from app.models.raw_document_request import RawDocumentRequest
 
 
 class Repository:
@@ -19,7 +20,7 @@ def get_digest(self, source_id: str) -> str | None:
         digest = self.redis.client.get(f"{config.DIGEST_PREFIX}:{source_id}")
         return digest.decode() if digest else None
 
-    def save(self, doc: HTMLDocumentRequest) -> None:
+    def save(self, doc: HTMLDocumentRequest | RawDocumentRequest) -> None:
         docs = self.preprocess_doc(doc)
         existing_keys = self.redis.client.keys(
             f"{self.redis.key_prefix}:{doc.source_id}:*"
@@ -42,8 +43,24 @@ def reset(self) -> None:
             self.redis.client.delete(*digest_keys)
         self.redis._create_index_if_not_exist(config.EMBEDDING_DIM)
 
-    def preprocess_doc(self, doc: HTMLDocumentRequest) -> list[DocumentModel]:
-        return self.document_transformer.transform_documents([preprocess(doc)])
+    @staticmethod
+    def _preprocess(
+        doc: HTMLDocumentRequest | RawDocumentRequest,
+    ) -> DocumentModel:
+        match doc:
+            case HTMLDocumentRequest():
+                return preprocess(doc)
+            case RawDocumentRequest():
+                return DocumentModel(
+                    metadata=doc.metadata,
+                    page_content=doc.page_content,
+                )
+
+    def preprocess_doc(
+        self,
+        doc: HTMLDocumentRequest | RawDocumentRequest,
+    ) -> list[DocumentModel]:
+        return self.document_transformer.transform_documents([self._preprocess(doc)])
 
 
 RepositoryDep = Annotated[Repository, Depends(Repository)]

app/models/base_document.py

@@ -1,10 +1,9 @@
 import hashlib
 from functools import cached_property
 from typing import TypedDict
-from urllib.parse import urlparse, urlunparse
 
 from langchain_core.documents import Document
-from pydantic import Field, validator
+from pydantic import Field
 
 
 def sha256(content: str) -> str:
@@ -26,24 +25,5 @@ def source_id(self):
     def page_content_digest(self):
         return sha256(self.page_content)
 
-    @validator("metadata", pre=True)
-    def validate_metadata(cls, metadata):
-        metadata["source"] = cls._normalize_url(metadata["source"])
-        return metadata
-
-    @staticmethod
-    def _normalize_url(url) -> str:
-        parsed_url = urlparse(url)
-        return urlunparse(
-            (
-                parsed_url.scheme,
-                parsed_url.netloc,
-                parsed_url.path,
-                parsed_url.params,
-                parsed_url.query,
-                "",
-            )
-        )
-
     class Config:
         keep_untouched = (cached_property,)

app/models/document_metadata_model.py

@@ -0,0 +1,8 @@
+from typing import NotRequired, TypedDict
+
+
+class DocumentMetadataModel(TypedDict):
+    source: str
+    title: NotRequired[str]
+    description: NotRequired[str]
+    language: NotRequired[str]

app/models/document_model.py

@@ -1,15 +1,7 @@
-from typing import NotRequired, TypedDict
-
 from pydantic import Field
 
 from app.models.base_document import BaseDocument
-
-
-class DocumentMetadataModel(TypedDict):
-    source: str
-    title: NotRequired[str]
-    description: NotRequired[str]
-    language: NotRequired[str]
+from app.models.document_metadata_model import DocumentMetadataModel
 
 
 class DocumentModel(BaseDocument):

app/models/html_document_request.py

@@ -1,7 +1,29 @@
 from typing import Literal
+from urllib.parse import urlparse, urlunparse
+
+from pydantic.class_validators import validator
 
 from app.models.base_document import BaseDocument
 
 
 class HTMLDocumentRequest(BaseDocument):
     type: Literal["HTML"] = "HTML"
+
+    @validator("metadata", pre=True)
+    def validate_metadata(cls, metadata):
+        metadata["source"] = cls._normalize_url(metadata["source"])
+        return metadata
+
+    @staticmethod
+    def _normalize_url(url) -> str:
+        parsed_url = urlparse(url)
+        return urlunparse(
+            (
+                parsed_url.scheme,
+                parsed_url.netloc,
+                parsed_url.path,
+                parsed_url.params,
+                parsed_url.query,
+                "",
+            )
+        )

app/models/raw_document_request.py

@@ -0,0 +1,12 @@
+from typing import Literal
+
+from pydantic.fields import Field
+
+from app.models.base_document import BaseDocument
+from app.models.document_metadata_model import DocumentMetadataModel
+
+
+class RawDocumentRequest(BaseDocument):
+    type: Literal["RAW"] = "RAW"
+
+    metadata: DocumentMetadataModel = Field(default_factory=DocumentMetadataModel)

app/routers/ingest.py

@@ -2,13 +2,14 @@
 
 from ..dependencies.repository import RepositoryDep
 from ..models.html_document_request import HTMLDocumentRequest
+from ..models.raw_document_request import RawDocumentRequest
 
 router = APIRouter()
 
 
 @router.post("/ingest", description="Ingest a document", status_code=201)
 def ingest_doc(
-    doc: HTMLDocumentRequest,
+    doc: HTMLDocumentRequest | RawDocumentRequest,
     response: Response,
     repository: RepositoryDep,
 ) -> str: