| from datetime import datetime, timezone |
| from typing import Optional, Dict, Any, List |
| from enum import Enum |
| from pydantic import BaseModel, Field, GetJsonSchemaHandler |
| from pydantic.json_schema import JsonSchemaValue |
| from pydantic_core import core_schema |
| from bson import ObjectId |
|
|
|
|
|
|
| class PaperSource(str, Enum): |
| UPLOAD = "upload" |
| ARXIV = "arxiv" |
| SEMANTIC_SCHOLAR = "semantic_scholar" |
| OPENALEX = "openalex" |
| OTHER = "other" |
|
|
|
|
| class PaperStatus(str, Enum): |
| UPLOADED = "uploaded" |
| PROCESSING = "processing" |
| INDEXED = "indexed" |
| FAILED = "failed" |
|
|
|
|
| class PyObjectId(ObjectId): |
| @classmethod |
| def __get_pydantic_core_schema__( |
| cls, source_type: Any, handler: Any |
| ) -> core_schema.CoreSchema: |
| def validate_from_str(input_value: str) -> ObjectId: |
| if isinstance(input_value, ObjectId): |
| return input_value |
| if isinstance(input_value, str): |
| if ObjectId.is_valid(input_value): |
| return ObjectId(input_value) |
| raise ValueError("Invalid ObjectId string") |
| raise ValueError("Invalid ObjectId type") |
|
|
| return core_schema.no_info_plain_validator_function(validate_from_str) |
|
|
| @classmethod |
| def __get_pydantic_json_schema__( |
| cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler |
| ) -> JsonSchemaValue: |
| return {"type": "string"} |
|
|
|
|
| class ResearchPaper(BaseModel): |
| """ |
| ResearchPaper model for MongoDB. |
| Represents metadata about an uploaded or fetched research paper. |
| """ |
| id: Optional[PyObjectId] = Field(default=None, alias="_id") |
| filename: str |
| stored_path: str |
| file_hash: str |
| source: PaperSource |
| status: PaperStatus |
| title: Optional[str] = None |
| authors: Optional[List[str]] = None |
| abstract: Optional[str] = None |
| publication_year: Optional[str] = None |
| doi: Optional[str] = None |
| arxiv_id: Optional[str] = None |
| semantic_scholar_id: Optional[str] = None |
| subject_areas: Optional[List[str]] = None |
| uploaded_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) |
| processed_at: Optional[datetime] = None |
| chunk_count: int = 0 |
| error_message: Optional[str] = None |
| metadata: Optional[Dict[str, Any]] = None |
| session_id: Optional[str] = None |
|
|
| class Config: |
| populate_by_name = True |
| arbitrary_types_allowed = True |
| json_encoders = {ObjectId: str} |
| json_schema_extra = { |
| "example": { |
| "filename": "attention_is_all_you_need.pdf", |
| "source": "upload", |
| "title": "Attention Is All You Need", |
| "authors": ["Vaswani et al."], |
| "publication_year": "2017" |
| } |
| } |
|
|
| def to_mongo(self) -> dict: |
| data = self.dict(by_alias=True, exclude={"id"}) |
| if self.id: |
| data["_id"] = self.id |
| return data |
|
|
| @classmethod |
| def from_mongo(cls, data: dict) -> "ResearchPaper": |
| if "_id" in data: |
| data = {**data, "_id": str(data["_id"])} |
| return cls(**data) |
