From d05c81ab9fa7b74d1ced1be4905fa742f526def0 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 11:10:58 -0700
Subject: [PATCH 1/8] Refactor file handling and type annotations in API and
 context modules

- Updated type hints in `files.py` for better clarity and consistency.
- Enhanced error handling in `put_file` and `list_objects` functions.
- Improved metadata handling in `ObjectMetadata` and `ListObjectsResponse` classes.
- Cleaned up optional parameters and added type ignores where necessary for compatibility.
---
 .../extralit_server/api/handlers/v1/files.py  |  7 +++-
 .../extralit_server/api/schemas/v1/files.py   | 14 +++----
 .../src/extralit_server/contexts/files.py     | 41 +++++++++++--------
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/extralit-server/src/extralit_server/api/handlers/v1/files.py b/extralit-server/src/extralit_server/api/handlers/v1/files.py
index 7fc6e0d94..4706c15bc 100644
--- a/extralit-server/src/extralit_server/api/handlers/v1/files.py
+++ b/extralit-server/src/extralit_server/api/handlers/v1/files.py
@@ -72,7 +72,12 @@ async def put_file(
 
     try:
         response = files.put_object(
-            client, bucket, object, data=file.file, size=file.size, content_type=file.content_type
+            client,
+            bucket,
+            object,
+            data=file.file,
+            size=file.size,  # type: ignore
+            content_type=file.content_type,  # type: ignore
         )
         return response
     except S3Error as se:
diff --git a/extralit-server/src/extralit_server/api/schemas/v1/files.py b/extralit-server/src/extralit_server/api/schemas/v1/files.py
index 1af728715..bd7ff4f6b 100644
--- a/extralit-server/src/extralit_server/api/schemas/v1/files.py
+++ b/extralit-server/src/extralit_server/api/schemas/v1/files.py
@@ -44,10 +44,10 @@ def parse_metadata(cls, v):
         return v
 
     @classmethod
-    def from_minio_object(cls, minio_object: Object):
+    def from_minio_object(cls, minio_object: Object) -> "ObjectMetadata":
         return cls(
             bucket_name=minio_object.bucket_name,
-            object_name=minio_object.object_name,
+            object_name=minio_object.object_name or "",
             last_modified=minio_object.last_modified,
             is_latest=None if minio_object.is_latest is None else minio_object.is_latest.lower() == "true",
             etag=minio_object.etag,
@@ -58,7 +58,7 @@ def from_minio_object(cls, minio_object: Object):
         )
 
     @classmethod
-    def from_minio_write_response(cls, write_result: ObjectWriteResult):
+    def from_minio_write_response(cls, write_result: ObjectWriteResult) -> "ObjectMetadata":
         return cls(
             bucket_name=write_result.bucket_name,
             object_name=write_result.object_name,
@@ -76,10 +76,10 @@ class ListObjectsResponse(BaseModel):
     objects: Iterable[ObjectMetadata] = Field(default_factory=list)
 
     def __len__(self) -> int:
-        return len(self.objects)
+        return len(self.objects)  # type: ignore
 
     def __getitem__(self, index) -> ObjectMetadata:
-        return self.objects[index]
+        return self.objects[index]  # type: ignore
 
     def __iter__(self):
         return iter(self.objects)
@@ -114,14 +114,14 @@ def assign_version_id(cls, objects: List[ObjectMetadata]) -> List[ObjectMetadata
 
 class FileObjectResponse(BaseModel):
     response: HTTPResponse
-    metadata: Optional[ObjectMetadata]
+    metadata: ObjectMetadata
     versions: Optional[ListObjectsResponse]
 
     class Config:
         arbitrary_types_allowed = True
 
     @property
-    def version_tag(self) -> str:
+    def version_tag(self) -> Optional[str]:
         if not self.metadata or not self.versions:
             return ""
         else:
diff --git a/extralit-server/src/extralit_server/contexts/files.py b/extralit-server/src/extralit_server/contexts/files.py
index f16c6af33..463e03303 100644
--- a/extralit-server/src/extralit_server/contexts/files.py
+++ b/extralit-server/src/extralit_server/contexts/files.py
@@ -24,6 +24,7 @@
 from typing import Any, BinaryIO, Dict, List, Optional, Union
 from urllib.parse import urlparse
 from uuid import UUID
+from minio.datatypes import Object
 from urllib3 import HTTPResponse
 
 from fastapi import HTTPException
@@ -83,8 +84,8 @@ def put_object(
         data: Union[BinaryIO, bytes],
         length: Optional[int] = None,
         content_type: Optional[str] = None,
-        part_size: int = None,
-        metadata: Dict[str, Any] = None,
+        part_size: Optional[int] = None,
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> ObjectWriteResult:
         # Ensure bucket exists
         bucket_path = self._get_bucket_path(bucket_name)
@@ -126,7 +127,7 @@ def put_object(
             object_name=object_name,
             version_id=version_id,
             etag=content_hash,
-            http_headers={},
+            http_headers={},  # type: ignore
             last_modified=None,
             location=None,
         )
@@ -153,7 +154,7 @@ def get_object(self, bucket_name: str, object_name: str, version_id: Optional[st
         with open(meta_path, "r") as f:
             json.load(f)
 
-        return HTTPResponse(body=io.BytesIO(content), preload_content=False)
+        return HTTPResponse(body=io.BytesIO(content), preload_content=False)  # type: ignore
 
     def stat_object(self, bucket_name: str, object_name: str, version_id: Optional[str] = None) -> ObjectMetadata:
         if version_id:
@@ -177,15 +178,13 @@ def stat_object(self, bucket_name: str, object_name: str, version_id: Optional[s
 
         stats = path.stat()
 
-        last_modified = datetime.fromtimestamp(stats.st_mtime)
-
         return ObjectMetadata(
             bucket_name=bucket_name,
             object_name=object_name,
             version_id=version_id or metadata.get("version_id"),
             etag=metadata.get("etag"),
             size=stats.st_size,
-            last_modified=last_modified,
+            last_modified=datetime.fromtimestamp(stats.st_mtime),
             metadata=metadata,
             content_type=metadata.get("content_type", "application/octet-stream"),
         )
@@ -235,6 +234,7 @@ def list_objects(
         if start_after:
             files = [f for f in files if str(f.relative_to(bucket_path)) > start_after]
 
+        result = []
         for file_path in files:
             object_name = str(file_path.relative_to(bucket_path))
             stats = file_path.stat()
@@ -252,25 +252,27 @@ def list_objects(
                 object_name=object_name,
                 etag=metadata.get("etag"),
                 size=stats.st_size,
-                last_modified=stats.st_mtime,
+                last_modified=datetime.fromtimestamp(stats.st_mtime),
                 metadata=metadata,
                 content_type=metadata.get("content_type", "application/octet-stream"),
                 version_id=metadata.get("version_id") if include_version else None,
             )
 
-            yield obj
+            result.append(obj)
+
+        return result
 
 
 def get_minio_client() -> Optional[Union[Minio, LocalFileStorage]]:
     if None in [settings.s3_endpoint, settings.s3_access_key, settings.s3_secret_key]:
         # Use local file system storage if S3 settings are not provided
-        local_storage_path = os.path.join(settings.home_path, "storage")
+        local_storage_path = os.path.join(settings.home_path, "storage")  # type: ignore
         _LOGGER.info(f"Using local file storage at: {local_storage_path}")
         return LocalFileStorage(local_storage_path)
 
     try:
         parsed_url = urlparse(settings.s3_endpoint)
-        hostname = parsed_url.hostname
+        hostname: str = str(parsed_url.hostname)
         port = parsed_url.port
 
         if hostname is None:
@@ -318,10 +320,12 @@ def list_objects(
     recursive=True,
     start_after: Optional[str] = None,
 ) -> ListObjectsResponse:
-    objects = client.list_objects(
+    objects: List[ObjectMetadata | Object] = client.list_objects(  # type: ignore
         bucket, prefix=prefix, recursive=recursive, include_version=include_version, start_after=start_after
     )
-    objects = [ObjectMetadata.from_minio_object(obj) for obj in objects]
+    objects: List[ObjectMetadata] = [
+        obj if isinstance(obj, ObjectMetadata) else ObjectMetadata.from_minio_object(obj) for obj in objects
+    ]
     return ListObjectsResponse(objects=objects)
 
 
@@ -368,9 +372,9 @@ def put_object(
     bucket: str,
     object: str,
     data: Union[BinaryIO, bytes, str],
-    content_type: str = None,
-    size: int = None,
-    metadata: Dict[str, Any] = None,
+    size: int,
+    content_type: str = "application/octet-stream",
+    metadata: Optional[Dict[str, Any]] = None,
     part_size: int = 100 * 1024 * 1024,
 ) -> ObjectMetadata:
     if isinstance(data, bytes):
@@ -391,7 +395,7 @@ def put_object(
             content_type=content_type,
             length=size,
             part_size=part_size,
-            metadata=metadata,
+            metadata=metadata or {},
         )
 
         return ObjectMetadata.from_minio_write_response(response)
@@ -512,7 +516,8 @@ def delete_bucket(client: Union[Minio, LocalFileStorage], workspace_name: str):
             obj_list = list(objects)
             for obj in obj_list:
                 try:
-                    client.remove_object(workspace_name, obj.object_name, version_id=obj.version_id)
+                    if obj.object_name is not None:
+                        client.remove_object(workspace_name, obj.object_name, version_id=obj.version_id)
                 except S3Error as remove_err:
                     _LOGGER.warning(
                         f"Error removing object {obj.object_name} (version: {obj.version_id}) during bucket delete: {remove_err}"

From 8ca8056d11081af53f1f2d458b80c977ffe92ab6 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 11:11:35 -0700
Subject: [PATCH 2/8] Fix metadata handling in get_object function to ensure
 proper ObjectMetadata usage

---
 extralit-server/src/extralit_server/contexts/files.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/extralit-server/src/extralit_server/contexts/files.py b/extralit-server/src/extralit_server/contexts/files.py
index 463e03303..226b62061 100644
--- a/extralit-server/src/extralit_server/contexts/files.py
+++ b/extralit-server/src/extralit_server/contexts/files.py
@@ -357,7 +357,11 @@ def get_object(
         else:
             versions = None
 
-        return FileObjectResponse(response=obj, metadata=stat, versions=versions)
+        return FileObjectResponse(
+            response=obj,
+            metadata=stat if isinstance(stat, ObjectMetadata) else ObjectMetadata.from_minio_object(stat),
+            versions=versions,
+        )
 
     except S3Error as se:
         _LOGGER.error(f"Error getting object {object} from bucket {bucket}: {se}")

From 611eebf26b965df2c315c3dd7e62ce3032337da6 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 14:03:48 -0700
Subject: [PATCH 3/8] Refactor DocumentsAPI to support multi-criteria document
 retrieval

- Updated the `get` method to accept a dictionary of parameters for searching documents by workspace ID, ID, PMID, DOI, or reference.
- Enhanced error handling for cases with no documents found or multiple documents returned.
- Refactored related document model and repository methods to align with the new API structure.
- Updated frontend components to utilize the new document fetching logic, ensuring compatibility with the revised API.
---
 .../base/base-pdf-viewer/PDFViewer.vue        |   6 +-
 .../container/mode/useDocumentViewModel.ts    |  56 +++++----
 .../v1/domain/entities/document/Document.ts   |  15 +--
 .../usecases/get-document-by-id-use-case.ts   |  24 ++--
 .../repositories/DocumentRepository.ts        |  30 ++---
 .../api/handlers/v1/documents.py              |  84 +++++++-------
 .../unit/api/handlers/v1/test_documents.py    |  80 ++++++++++++-
 extralit/src/extralit/_api/_documents.py      |  51 +++-----
 extralit/src/extralit/_models/_documents.py   |   6 +-
 extralit/src/extralit/documents/_resource.py  | 109 +++++++-----------
 10 files changed, 254 insertions(+), 207 deletions(-)

diff --git a/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue b/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
index 1e2da70f2..84abea2f8 100644
--- a/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
+++ b/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
@@ -19,7 +19,7 @@
   </div>
 </template>
 
-<script>
+<script lang="ts">
 import { PDFView } from "@jonnytran/vue-pdf-viewer";
 
 export default {
@@ -44,8 +44,8 @@ export default {
   },
 
   watch: {
-    pageNumber(newPageNumber) {
-      this.currentPageNumber = newPageNumber;
+    pageNumber(newPageNumber: Number | String) {
+      this.currentPageNumber = newPageNumber ? Number(newPageNumber) : 1;
     },
   },
 
diff --git a/extralit-frontend/components/features/annotation/container/mode/useDocumentViewModel.ts b/extralit-frontend/components/features/annotation/container/mode/useDocumentViewModel.ts
index 0747fd31c..03d4ab8e7 100644
--- a/extralit-frontend/components/features/annotation/container/mode/useDocumentViewModel.ts
+++ b/extralit-frontend/components/features/annotation/container/mode/useDocumentViewModel.ts
@@ -6,10 +6,12 @@ import { Segment } from "@/v1/domain/entities/document/Document";
 import { useDataset } from "@/v1/infrastructure/storage/DatasetStorage";
 import { waitForAsyncValue } from "@/v1/infrastructure/services/useWait";
 import { useNotifications } from "~/v1/infrastructure/services/useNotifications";
+import { useWorkspaces } from "~/v1/infrastructure/storage/WorkspaceStorage";
 
 export const useDocumentViewModel = (props: { record: any }) => {
   const notification = useNotifications();
   const getDocument = useResolve(GetDocumentByIdUseCase);
+  const { state: workspaces } = useWorkspaces();
   const { state: dataset } = useDataset();
   const { state: document, set: setDocument, clear: clearDocument } = useDocument();
   const isLoading = ref(false);
@@ -19,28 +21,38 @@ export const useDocumentViewModel = (props: { record: any }) => {
   });
   const hasDocument = computed(() => {
     return (
-      props.record.metadata === null || props.record.metadata?.doc_id != null || props.record.metadata?.pmid != null
+      props.record.metadata === null || props.record.metadata?.reference != null || props.record.metadata?.pmid != null
     );
   });
 
-  const fetchDocumentByID = async (id: string) => {
+  const fetchDocument = async (metadata: any) => {
     try {
-      await getDocument.setDocumentByID(id);
-    } catch (e) {
-      notification.notify({
-        message: `Error fetching document with ID ${id}`,
-        type: "danger",
-      });
-      clearDocument();
-    }
-  };
+      await waitForAsyncValue(() => workspaces.selectedWorkspace?.id);
 
-  const fetchDocumentByPubmedID = async (pmid: string) => {
-    try {
-      await getDocument.setDocumentByPubmedID(pmid);
+      const params: {
+        workspace_id: string;
+        doc_id?: string;
+        pmid?: string;
+        doi?: string;
+        reference?: string;
+      } = { workspace_id: workspaces.selectedWorkspace!.id };
+
+      if (metadata?.reference) params.reference = metadata.reference;
+      if (metadata?.doc_id) params.doc_id = metadata.doc_id;
+      if (metadata?.pmid) params.pmid = metadata.pmid;
+      if (metadata?.doi) params.doi = metadata.doi;
+
+      // Ensure at least one identifier is provided
+      if (Object.keys(params).length === 0) {
+        throw new Error("No valid document identifier found in metadata");
+      }
+
+      await getDocument.setDocument(params);
     } catch (e) {
+      const identifier = metadata?.pmid || metadata?.doi || metadata?.doc_id || metadata?.reference || "unknown";
+      console.error(`Error fetching document with identifier "${identifier}":`, e);
       notification.notify({
-        message: `Error fetching document with pmid "${pmid}"`,
+        message: `Error fetching document with identifier "${identifier}"`,
         type: "danger",
       });
       clearDocument();
@@ -48,13 +60,13 @@ export const useDocumentViewModel = (props: { record: any }) => {
   };
 
   const updateDocument = async (metadata: any) => {
-    if (metadata?.pmid != null) {
-      if (document.pmid !== metadata.pmid) {
-        fetchDocumentByPubmedID(metadata.pmid);
-      }
+    if (metadata?.pmid != null && document.pmid !== metadata.pmid) {
+      fetchDocument(metadata);
+    } else if (metadata?.doi != null && document.doi !== metadata.doi) {
+      fetchDocument(metadata);
     } else if (metadata?.doc_id != null && document.id !== metadata.doc_id) {
-      fetchDocumentByID(metadata.doc_id);
-    } else if (!metadata?.pmid && !metadata?.doc_id && hasDocumentLoaded.value) {
+      fetchDocument(metadata);
+    } else if (!metadata?.pmid && !metadata?.doi && !metadata?.doc_id && hasDocumentLoaded.value) {
       clearDocument();
     }
 
@@ -107,8 +119,6 @@ export const useDocumentViewModel = (props: { record: any }) => {
 
   return {
     document,
-    fetchDocumentByID,
-    fetchDocumentByPubmedID,
     fetchDocumentSegments,
     focusDocumentPageNumber,
     clearDocument,
diff --git a/extralit-frontend/v1/domain/entities/document/Document.ts b/extralit-frontend/v1/domain/entities/document/Document.ts
index d70e80555..333198c3d 100644
--- a/extralit-frontend/v1/domain/entities/document/Document.ts
+++ b/extralit-frontend/v1/domain/entities/document/Document.ts
@@ -6,19 +6,19 @@ export class Segment {
 		public readonly header: string | null,
 		public readonly page_number: number | null,
 		public readonly type: string | null,
-	) {}
+	) { }
 
 	public static getDescription(segment: Segment): string {
-			let segmentDescription = Object.entries(segment)
-					.filter(([key, value]) => !['header', 'doc_id'].includes(key))
-					.map(([key, value]) => `${key}: ${value}`)
-					.join('\n');
-			return segmentDescription;
+		let segmentDescription = Object.entries(segment)
+			.filter(([key, value]) => !['header', 'doc_id'].includes(key))
+			.map(([key, value]) => `${key}: ${value}`)
+			.join('\n');
+		return segmentDescription;
 	}
 }
 
 export interface Segments {
-  items: Segment[];
+	items: Segment[];
 }
 
 export class Document {
@@ -27,6 +27,7 @@ export class Document {
 		public readonly url?: string,
 		public readonly file_name?: string,
 		public readonly pmid?: string,
+		public readonly doi?: string,
 		public readonly page_number?: number | string,
 		public reference?: string,
 		public segments?: Segment[],
diff --git a/extralit-frontend/v1/domain/usecases/get-document-by-id-use-case.ts b/extralit-frontend/v1/domain/usecases/get-document-by-id-use-case.ts
index 584b4af1c..cb69dbde1 100644
--- a/extralit-frontend/v1/domain/usecases/get-document-by-id-use-case.ts
+++ b/extralit-frontend/v1/domain/usecases/get-document-by-id-use-case.ts
@@ -8,16 +8,26 @@ export class GetDocumentByIdUseCase {
     private readonly documentStorage: IDocumentStorage
   ) {}
 
-  async setDocumentByID(id: string) {
-    const document = await this.documentRepository.getDocumentById(id);
+  async setDocument(params: {
+    workspace_id: string;
+    doc_id?: string;
+    pmid?: string;
+    doi?: string;
+    reference?: string;
+  }) {
+    const documents = await this.documentRepository.getDocuments(params);
 
-    this.documentStorage.set(document);
-  }
+    if (documents.length === 0) {
+      throw new Error("No documents found with the provided criteria");
+    }
 
-  async setDocumentByPubmedID(id: string) {
-    const document = await this.documentRepository.getDocumentByPubmedID(id);
+    // For now, we'll use the first document found
+    // In the future, you might want to add logic to handle multiple documents
+    if (documents.length > 1) {
+      console.warn(`Multiple documents found (${documents.length}). Using the first one.`);
+    }
 
-    this.documentStorage.set(document);
+    this.documentStorage.set(documents[0]);
   }
 
   async setSegments(workspace: string, reference: string): Promise<Segment[]> {
diff --git a/extralit-frontend/v1/infrastructure/repositories/DocumentRepository.ts b/extralit-frontend/v1/infrastructure/repositories/DocumentRepository.ts
index 7e75f2707..3b2e50911 100644
--- a/extralit-frontend/v1/infrastructure/repositories/DocumentRepository.ts
+++ b/extralit-frontend/v1/infrastructure/repositories/DocumentRepository.ts
@@ -3,25 +3,29 @@ import { Document, Segment, Segments } from "@/v1/domain/entities/document/Docum
 
 const DOCUMENT_API_ERRORS = {
   ERROR_FETCHING_DOCUMENT: "ERROR_FETCHING_DOCUMENT",
+  ERROR_FETCHING_SEGMENTS: "ERROR_FETCHING_SEGMENTS",
 };
 
 export class DocumentRepository {
   constructor(private readonly axios: NuxtAxiosInstance) {}
 
-  async getDocumentByPubmedID(pmid: string): Promise<Document> {
+  async getDocuments(params: {
+    workspace_id: string;
+    doc_id?: string;
+    pmid?: string;
+    doi?: string;
+    reference?: string;
+  }): Promise<Document[]> {
     try {
-      const { data } = await this.axios.get<Document>(`/v1/documents/by-pmid/${pmid}`);
-      return data;
-    } catch (error) {
-      throw {
-        response: DOCUMENT_API_ERRORS.ERROR_FETCHING_DOCUMENT,
-      };
-    }
-  }
+      const queryParams = Object.fromEntries(Object.entries(params).filter(([_, value]) => value !== undefined));
 
-  async getDocumentById(id: string): Promise<Document> {
-    try {
-      const { data } = await this.axios.get<Document>(`/v1/documents/by-id/${id}`);
+      if (Object.keys(queryParams).length === 0) {
+        throw new Error("At least one identifier parameter must be provided");
+      }
+
+      const { data } = await this.axios.get<Document[]>("/v1/documents", {
+        params: queryParams,
+      });
       return data;
     } catch (error) {
       throw {
@@ -39,7 +43,7 @@ export class DocumentRepository {
       return data.items;
     } catch (error) {
       throw {
-        response: DOCUMENT_API_ERRORS.ERROR_FETCHING_DOCUMENT,
+        response: DOCUMENT_API_ERRORS.ERROR_FETCHING_SEGMENTS,
       };
     }
   }
diff --git a/extralit-server/src/extralit_server/api/handlers/v1/documents.py b/extralit-server/src/extralit_server/api/handlers/v1/documents.py
index 5dd5e825b..8fc77ca2b 100644
--- a/extralit-server/src/extralit_server/api/handlers/v1/documents.py
+++ b/extralit-server/src/extralit_server/api/handlers/v1/documents.py
@@ -15,9 +15,9 @@
 import json
 import logging
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
-from fastapi import APIRouter, Body, Depends, File, Form, HTTPException, UploadFile, Path, status, Security
+from fastapi import APIRouter, Body, Depends, File, Form, HTTPException, UploadFile, Path, status, Security, Query
 from minio import Minio
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy import select
@@ -71,7 +71,7 @@ async def add_document(
         file_url = files.put_document_file(
             client=client,
             workspace_name=workspace.name,
-            document_id=document_create.id,
+            document_id=document_create.id,  # type: ignore[arg-type]
             file_data=file_data_bytes,
             filename=file_data.filename or "",
             metadata=document_create.dict(include={"file_name": True, "pmid": True, "doi": True}),
@@ -106,58 +106,57 @@ async def add_document(
     return document.id
 
 
-@router.get("/documents/by-pmid/{pmid}", response_model=DocumentListItem)
-async def get_document_by_pmid(
-    *, db: AsyncSession = Depends(get_async_db), pmid: str, current_user: User = Security(auth.get_current_user)
-):
-    if pmid is None or not isinstance(pmid, str) or not pmid.isnumeric():
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail=f"Document with pmid `{pmid}` not found",
-        )
-
-    query = await db.execute(select(Document).where(Document.pmid == pmid))
-    await authorize(current_user, DocumentPolicy.get())
-
-    documents = query.fetchone()
-    if documents is None or len(documents) == 0:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail=f"Document with pmid `{pmid}` not found",
-        )
-
-    document: Document = documents[0]
-    return DocumentListItem.model_validate(document)
-
-
-@router.get("/documents/by-id/{id}", response_model=DocumentListItem)
-async def get_document_by_id(
+@router.get(
+    "/documents", response_model=List[DocumentListItem], description="Get documents by ID, PMID, DOI, or reference."
+)
+async def get_document(
     *,
-    id: UUID = Path(..., title="The UUID of the document to get"),
+    workspace_id: UUID = Query(..., description="Workspace ID"),
+    id: Optional[UUID] = Query(None, description="Document ID"),
+    reference: Optional[str] = Query(None, description="Document reference"),
+    pmid: Optional[str] = Query(None, description="PubMed ID"),
+    doi: Optional[str] = Query(None, description="DOI"),
     db: AsyncSession = Depends(get_async_db),
     current_user: User = Security(auth.get_current_user),
-):
-    if id is None:
+) -> List[DocumentListItem]:
+    await authorize(current_user, DocumentPolicy.get())
+
+    if not any([id, pmid, doi, reference]):
         raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND,
-            detail=f"Document with id `{id}` not found",
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="At least one of id, pmid, doi, or reference must be provided",
         )
 
-    query = await db.execute(select(Document).where(Document.id == id))
-    await authorize(current_user, DocumentPolicy.get())
+    documents = await imports.find_existing_documents(
+        db=db,
+        workspace_id=workspace_id,
+        document_id=id,
+        pmid=pmid,
+        doi=doi,
+        reference=reference,
+    )
+
+    if not documents:
+        # If we still haven't found anything, raise 404
+        search_criteria = []
+        if id:
+            search_criteria.append(f"id={id}")
+        if pmid:
+            search_criteria.append(f"pmid={pmid}")
+        if doi:
+            search_criteria.append(f"doi={doi}")
+        if reference:
+            search_criteria.append(f"reference={reference}")
 
-    documents = query.fetchone()
-    if documents is None or len(documents) == 0:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
-            detail=f"Document with id `{id}` not found",
+            detail=f"No documents found with criteria: {', '.join(search_criteria)} in workspace {workspace_id}",
         )
 
-    document: Document = documents[0]
-    return DocumentListItem.model_validate(document)
+    return documents
 
 
-@router.patch("/documents/{id}", response_model=DocumentListItem)
+@router.patch("/documents/{id}", response_model=DocumentListItem, description="Update a document by ID.")
 async def update_document(
     *,
     id: UUID = Path(..., title="The UUID of the document to update"),
@@ -165,7 +164,6 @@ async def update_document(
     db: AsyncSession = Depends(get_async_db),
     current_user: User = Security(auth.get_current_user),
 ):
-    """Update a document by ID."""
     # First, get the document to ensure it exists and check permissions
     query = await db.execute(select(Document).where(Document.id == id))
     result = query.fetchone()
diff --git a/extralit-server/tests/unit/api/handlers/v1/test_documents.py b/extralit-server/tests/unit/api/handlers/v1/test_documents.py
index bc47b1767..d82dd77a2 100644
--- a/extralit-server/tests/unit/api/handlers/v1/test_documents.py
+++ b/extralit-server/tests/unit/api/handlers/v1/test_documents.py
@@ -129,22 +129,90 @@ async def test_get_document_by_pmid(async_client: "AsyncClient", db: "AsyncSessi
     workspace = await WorkspaceFactory.create()
     document = await DocumentFactory.create(pmid="123456", workspace=workspace, workspace_id=workspace.id)
 
-    response = await async_client.get(f"/api/v1/documents/by-pmid/{document.pmid}", headers=owner_auth_header)
+    response = await async_client.get(
+        "/api/v1/documents",
+        params={"pmid": document.pmid, "workspace_id": str(workspace.id)},
+        headers=owner_auth_header,
+    )
 
     assert response.status_code == 200
-    assert response.json()["pmid"] == document.pmid
+    response_data = response.json()
+    assert isinstance(response_data, list)
+    assert len(response_data) == 1
+    assert response_data[0]["pmid"] == document.pmid
 
 
 @pytest.mark.asyncio
 @pytest.mark.skip(reason="'coroutine' object has no attribute 'id'")
 async def test_get_document_by_id(async_client: AsyncClient, db: AsyncSession, owner_auth_header: dict):
-    document = await DocumentFactory.create()
+    workspace = await WorkspaceFactory.create()
+    document = await DocumentFactory.create(workspace=workspace, workspace_id=workspace.id)
+
+    response = await async_client.get(
+        "/api/v1/documents",
+        params={"id": str(document.id), "workspace_id": str(workspace.id)},
+        headers=owner_auth_header,
+    )
+
+    assert response.status_code == 200
+    response_data = response.json()
+    assert isinstance(response_data, list)
+    assert len(response_data) == 1
+    assert response_data[0]["id"] == str(document.id)
+
+
+@pytest.mark.asyncio
+async def test_get_document_by_doi(async_client: "AsyncClient", db: "AsyncSession", owner_auth_header: dict):
+    workspace = await WorkspaceFactory.create()
+    document = await DocumentFactory.create(doi="10.1234/test.doi", workspace=workspace, workspace_id=workspace.id)
+
+    response = await async_client.get(
+        "/api/v1/documents", params={"doi": document.doi, "workspace_id": str(workspace.id)}, headers=owner_auth_header
+    )
+
+    assert response.status_code == 200
+    response_data = response.json()
+    assert isinstance(response_data, list)
+    assert len(response_data) == 1
+    assert response_data[0]["doi"] == document.doi
 
-    response = await async_client.get(f"/api/v1/documents/by-id/{document.id}", headers=owner_auth_header)
+
+@pytest.mark.asyncio
+async def test_get_document_by_reference(async_client: "AsyncClient", db: "AsyncSession", owner_auth_header: dict):
+    workspace = await WorkspaceFactory.create()
+    document = await DocumentFactory.create(reference="test_ref_123", workspace=workspace, workspace_id=workspace.id)
+
+    response = await async_client.get(
+        "/api/v1/documents",
+        params={"reference": document.reference, "workspace_id": str(workspace.id)},
+        headers=owner_auth_header,
+    )
 
     assert response.status_code == 200
-    response_json = response.json()
-    assert response_json["id"] == str(document.id)
+    response_data = response.json()
+    assert isinstance(response_data, list)
+    assert len(response_data) == 1
+    assert response_data[0]["reference"] == document.reference
+
+
+@pytest.mark.asyncio
+async def test_get_document_workspace_id_only(async_client: "AsyncClient", owner_auth_header: dict):
+    """Test that requesting documents with only workspace_id returns a 400 error."""
+    response = await async_client.get(
+        "/api/v1/documents", params={"workspace_id": "123e4567-e89b-12d3-a456-426614174000"}, headers=owner_auth_header
+    )
+
+    assert response.status_code == 400
+    assert "At least one of id, pmid, doi, or reference must be provided" in response.json()["detail"]
+
+
+@pytest.mark.asyncio
+async def test_get_document_no_parameters(async_client: "AsyncClient", owner_auth_header: dict):
+    """Test that requesting documents without any parameters returns a 400 error."""
+    response = await async_client.get("/api/v1/documents", headers=owner_auth_header)
+
+    assert response.status_code == 400
+    assert "At least one of id, pmid, doi, or reference must be provided" in response.json()["detail"]
 
 
 @pytest.mark.skip(reason="Document delete API is failing with 500 error")
diff --git a/extralit/src/extralit/_api/_documents.py b/extralit/src/extralit/_api/_documents.py
index b4a140234..2d0da00fa 100644
--- a/extralit/src/extralit/_api/_documents.py
+++ b/extralit/src/extralit/_api/_documents.py
@@ -57,51 +57,35 @@ def create(self, model: "DocumentModel") -> "DocumentModel":
         return model
 
     @api_error_handler
-    def get(self, id: UUID) -> "DocumentModel":
-        """Get a document by ID.
+    def get(self, params: dict) -> "DocumentModel":
+        """Get a document using multiple search criteria.
 
         Args:
-            id: The document ID.
+            params: Dictionary containing any combination of:
+                - workspace_id: Workspace ID (required)
+                - id: Document ID
+                - pmid: PubMed ID
+                - doi: DOI
+                - reference: Document reference
 
         Returns:
             The document model.
         """
         from extralit._models._documents import DocumentModel
 
-        url = f"/api/v1/documents/by-id/{id}"
-        response = self.http_client.get(url=url)
+        url = "/api/v1/documents"
+        response = self.http_client.get(url=url, params=params)
         response.raise_for_status()
 
-        doc_data = response.json()
-        return DocumentModel(
-            id=doc_data.get("id"),
-            workspace_id=doc_data.get("workspace_id"),
-            file_name=doc_data.get("file_name"),
-            reference=doc_data.get("reference"),
-            url=doc_data.get("url"),
-            pmid=doc_data.get("pmid"),
-            doi=doc_data.get("doi"),
-            inserted_at=doc_data.get("inserted_at"),
-            updated_at=doc_data.get("updated_at"),
-        )
-
-    @api_error_handler
-    def get_by_pmid(self, pmid: str) -> "DocumentModel":
-        """Get a document by PMID.
+        doc_data_list = response.json()
 
-        Args:
-            pmid: The PubMed ID.
+        if not doc_data_list:
+            raise ValueError("No documents found with the provided criteria")
 
-        Returns:
-            The document model.
-        """
-        from extralit._models._documents import DocumentModel
-
-        url = f"/api/v1/documents/by-pmid/{pmid}"
-        response = self.http_client.get(url=url)
-        response.raise_for_status()
+        if len(doc_data_list) > 1:
+            print(f"Warning: Multiple documents found ({len(doc_data_list)}). Using the first one.")
 
-        doc_data = response.json()
+        doc_data = doc_data_list[0]
         return DocumentModel(
             id=doc_data.get("id"),
             workspace_id=doc_data.get("workspace_id"),
@@ -112,6 +96,7 @@ def get_by_pmid(self, pmid: str) -> "DocumentModel":
             doi=doc_data.get("doi"),
             inserted_at=doc_data.get("inserted_at"),
             updated_at=doc_data.get("updated_at"),
+            file_path=None,
         )
 
     @api_error_handler
@@ -142,6 +127,7 @@ def list(self, workspace_id: UUID) -> List["DocumentModel"]:
                 doi=doc_data.get("doi"),
                 inserted_at=doc_data.get("inserted_at"),
                 updated_at=doc_data.get("updated_at"),
+                file_path=None,
             )
             documents.append(doc)
 
@@ -199,4 +185,5 @@ def update(self, model: "DocumentModel") -> "DocumentModel":
             doi=doc_data.get("doi"),
             inserted_at=doc_data.get("inserted_at"),
             updated_at=doc_data.get("updated_at"),
+            file_path=None,
         )
diff --git a/extralit/src/extralit/_models/_documents.py b/extralit/src/extralit/_models/_documents.py
index bcdd4990a..d17cdb00c 100644
--- a/extralit/src/extralit/_models/_documents.py
+++ b/extralit/src/extralit/_models/_documents.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional
 from urllib.parse import unquote, urlparse
 from uuid import UUID
 
@@ -49,7 +49,7 @@ def from_file(
         file_path_or_url: str,
         *,
         reference: str,
-        workspace_id: Union[UUID, str],
+        workspace_id: UUID,
         id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> "DocumentModel":
@@ -60,7 +60,7 @@ def from_file(
 
         elif urlparse(file_path_or_url).scheme:
             url = file_path_or_url
-            file_path_or_url = None
+            file_path_or_url = None  # type: ignore[assignment]
             parsed_url = urlparse(url)
             path = parsed_url.path
             file_name = unquote(path).split("/")[-1]
diff --git a/extralit/src/extralit/documents/_resource.py b/extralit/src/extralit/documents/_resource.py
index d2b985299..4f8a4966c 100644
--- a/extralit/src/extralit/documents/_resource.py
+++ b/extralit/src/extralit/documents/_resource.py
@@ -51,10 +51,10 @@ class Document(Resource):
 
     def __init__(
         self,
-        workspace_id: Optional[UUID] = None,
+        workspace_id: UUID,
+        reference: str,
         file_name: Optional[str] = None,
         file_path: Optional[Union[str, Path]] = None,
-        reference: Optional[str] = None,
         url: Optional[str] = None,
         pmid: Optional[str] = None,
         doi: Optional[str] = None,
@@ -64,10 +64,10 @@ def __init__(
         """Initializes a Document object.
 
         Parameters:
-            workspace_id (UUID): The workspace ID to which this document belongs.
+            workspace_id (UUID): The workspace ID to which this document belongs (required).
+            reference (str): A reference identifier for the document (required).
             file_name (str): The name of the document file.
             file_path (Union[str, Path]): Local path to the document file.
-            reference (str): A reference identifier for the document.
             url (str): The URL of the document.
             pmid (str): The PubMed ID of the document.
             doi (str): The DOI of the document.
@@ -100,7 +100,7 @@ def from_file(
         file_path_or_url: Union[str, Path],
         *,
         reference: str,
-        workspace_id: Optional[UUID] = None,
+        workspace_id: UUID,
         pmid: Optional[str] = None,
         doi: Optional[str] = None,
         client: Optional["Extralit"] = None,
@@ -111,7 +111,6 @@ def from_file(
             file_path_or_url: Local file path or URL to the document.
             reference: A reference identifier for the document.
             workspace_id: The workspace ID to which this document belongs.
-            id: The document ID. If provided, the document will be created with this ID.
             pmid: The PubMed ID of the document.
             doi: The DOI of the document.
             client: The client used to interact with Extralit.
@@ -120,8 +119,14 @@ def from_file(
             Document: The created document object.
 
         Raises:
-            ValueError: If the file path does not exist or URL is invalid.
+            ValueError: If the file path does not exist or URL is invalid, or if required parameters are missing.
         """
+        if not workspace_id:
+            raise ValueError("workspace_id is required")
+
+        if not reference:
+            raise ValueError("reference is required")
+
         if isinstance(file_path_or_url, Path):
             file_path_or_url = str(file_path_or_url)
 
@@ -150,60 +155,6 @@ def from_file(
             client=client,
         )
 
-    @classmethod
-    def from_pmid(
-        cls,
-        pmid: str,
-        *,
-        reference: Optional[str] = None,
-        workspace_id: Optional[UUID] = None,
-        client: Optional["Extralit"] = None,
-    ) -> "Document":
-        """Create a Document from a PubMed ID.
-
-        Args:
-            pmid: The PubMed ID.
-            reference: A reference identifier for the document.
-            workspace_id: The workspace ID to which this document belongs.
-            client: The client used to interact with Extralit.
-
-        Returns:
-            Document: The created document object.
-        """
-        return cls(
-            workspace_id=workspace_id,
-            pmid=pmid,
-            reference=reference or f"PMID:{pmid}",
-            client=client,
-        )
-
-    @classmethod
-    def from_doi(
-        cls,
-        doi: str,
-        *,
-        reference: Optional[str] = None,
-        workspace_id: Optional[UUID] = None,
-        client: Optional["Extralit"] = None,
-    ) -> "Document":
-        """Create a Document from a DOI.
-
-        Args:
-            doi: The DOI.
-            reference: A reference identifier for the document.
-            workspace_id: The workspace ID to which this document belongs.
-            client: The client used to interact with Extralit.
-
-        Returns:
-            Document: The created document object.
-        """
-        return cls(
-            workspace_id=workspace_id,
-            doi=doi,
-            reference=reference or f"DOI:{doi}",
-            client=client,
-        )
-
     @classmethod
     def from_model(cls, model: DocumentModel, client: "Extralit") -> "Document":
         """Create a Document from a DocumentModel.
@@ -232,32 +183,50 @@ def from_model(cls, model: DocumentModel, client: "Extralit") -> "Document":
     @classmethod
     def get(
         cls,
+        workspace_id: UUID,
+        reference: Optional[str] = None,
         id: Optional[UUID] = None,
         pmid: Optional[str] = None,
+        doi: Optional[str] = None,
         client: Optional["Extralit"] = None,
     ) -> "Document":
-        """Get a document by ID or PMID.
+        """Get a document by ID, PMID, DOI, reference, or workspace_id.
 
         Args:
+            workspace_id: The workspace ID (required).
             id: The document ID.
             pmid: The PubMed ID.
+            doi: The DOI.
+            reference: The document reference.
             client: The client used to interact with Extralit.
 
         Returns:
             Document: The document object.
 
         Raises:
-            ValueError: If neither id nor pmid is provided.
+            ValueError: If workspace_id is not provided or if none of id, pmid, doi, or reference is provided.
         """
+        if not workspace_id:
+            raise ValueError("workspace_id is required")
+
         client = client or Extralit._get_default()
 
+        # Build parameters object
+        params = {"workspace_id": str(workspace_id)}
         if id:
-            model = client.api.documents.get(id)
-        elif pmid:
-            model = client.api.documents.get_by_pmid(pmid)
-        else:
-            raise ValueError("Either id or pmid must be provided")
-
+            params["id"] = str(id)
+        if pmid:
+            params["pmid"] = pmid
+        if doi:
+            params["doi"] = doi
+        if reference:
+            params["reference"] = reference
+
+        # Ensure at least one identifier is provided in addition to workspace_id
+        if len(params) <= 1:  # Only workspace_id is provided
+            raise ValueError("At least one of id, pmid, doi, or reference must be provided in addition to workspace_id")
+
+        model: DocumentModel = client.api.documents.get(params)
         return cls.from_model(model, client)
 
     ############################
@@ -340,4 +309,4 @@ def delete(self) -> None:
         self._log_message(f"Document deleted: {self}")
 
     def _with_client(self, client: "Extralit") -> "Self":
-        return Document.from_model(self._model, client)
+        return Document.from_model(self._model, client)  # type: ignore[return-value]

From 122b2046880f2dac7fa59110731f7accb1cb8deb Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 14:23:51 -0700
Subject: [PATCH 4/8] changelogs

---
 extralit-frontend/CHANGELOG.md | 4 ++++
 extralit-server/CHANGELOG.md   | 5 +++++
 2 files changed, 9 insertions(+)

diff --git a/extralit-frontend/CHANGELOG.md b/extralit-frontend/CHANGELOG.md
index 052c53a03..e8a4fa3ba 100644
--- a/extralit-frontend/CHANGELOG.md
+++ b/extralit-frontend/CHANGELOG.md
@@ -14,6 +14,10 @@ These are the section headers that we use:
 * "Security" in case of vulnerabilities.
 -->
 
+## [Extralit] [0.6.1](https://github.com/extralit/extralit/compare/v0.5.0...v0.6.1)
+### Changed
+- Refactored the frontend to use a single fetchDocument method that queries documents by any identifier and workspace, replacing the previous fetchDocumentByID and fetchDocumentByPubmedID methods. The view model and use case now expect and handle the new API response format
+
 ## [Extralit] [0.6.0](https://github.com/extralit/extralit/compare/v0.4.1...v0.6.0)
 
 ### Added
diff --git a/extralit-server/CHANGELOG.md b/extralit-server/CHANGELOG.md
index 5f3387a12..9a9afb27b 100644
--- a/extralit-server/CHANGELOG.md
+++ b/extralit-server/CHANGELOG.md
@@ -14,6 +14,11 @@ These are the section headers that we use:
 * "Security" in case of vulnerabilities.
 -->
 
+## [Extralit] [0.6.1](https://github.com/extralit/extralit/compare/v0.5.0...v0.6.1)
+
+### Changed
+- Replaced separate /documents/by-id/{id} and /documents/by-pmid/{pmid} endpoints with a single /documents endpoint that accepts workspace_id and one or more identifiers (id, pmid, doi, reference), returning a list of matching documents
+
 ## [Extralit] [0.6.0](https://github.com/extralit/extralit/compare/v0.5.0...v0.6.0)
 
 ### Added

From 1da050cd43b6007840c702556e6e1ceed2446ad3 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 14:24:35 -0700
Subject: [PATCH 5/8] changelogs

---
 extralit/CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/extralit/CHANGELOG.md b/extralit/CHANGELOG.md
index 44c824709..940c72271 100644
--- a/extralit/CHANGELOG.md
+++ b/extralit/CHANGELOG.md
@@ -14,6 +14,10 @@ These are the section headers that we use:
 * "Security" in case of vulnerabilities.
 -->
 
+## [Extralit] [0.6.1](https://github.com/extralit/extralit/compare/v0.5.0...v0.6.1)
+### Changed
+- Updated backend document fetching logic to use the new unified endpoint, including improved input validation and error messages
+
 ## [Extralit] [0.6.0](https://github.com/extralit/extralit/compare/v0.5.0...v0.6.0)
 
 ### Added

From 578a9e981fdfb5268180fb030d2b04aed69bd198 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 14:37:02 -0700
Subject: [PATCH 6/8] fix tests

---
 .../base/base-pdf-viewer/PDFViewer.vue        | 15 ------------
 .../unit/api/handlers/v1/test_documents.py    |  6 ++---
 .../unit/api/test_workspace_documents_api.py  | 24 ++++++++++++-------
 3 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue b/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
index 84abea2f8..bdf9463ca 100644
--- a/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
+++ b/extralit-frontend/components/base/base-pdf-viewer/PDFViewer.vue
@@ -56,21 +56,6 @@ export default {
     };
   },
 
-  mounted() {
-    window.addEventListener("hashchange", this.onHashChange);
-    this.onHashChange(); // Call on component mount to handle initial hash
-  },
-
-  methods: {
-    onHashChange() {
-      const hash = window.location.hash.substring(1); // Remove the '#' from the hash
-      const [key, value] = hash.split(".");
-      if (key === "page_number" && !isNaN(value)) {
-        this.currentPageNumber = Number(value);
-      }
-    },
-  },
-
   errorCaptured(err, component, info) {
     this.error = err;
     console.error(`Error caught from ${component}: ${err}`);
diff --git a/extralit-server/tests/unit/api/handlers/v1/test_documents.py b/extralit-server/tests/unit/api/handlers/v1/test_documents.py
index d82dd77a2..426344724 100644
--- a/extralit-server/tests/unit/api/handlers/v1/test_documents.py
+++ b/extralit-server/tests/unit/api/handlers/v1/test_documents.py
@@ -208,11 +208,11 @@ async def test_get_document_workspace_id_only(async_client: "AsyncClient", owner
 
 @pytest.mark.asyncio
 async def test_get_document_no_parameters(async_client: "AsyncClient", owner_auth_header: dict):
-    """Test that requesting documents without any parameters returns a 400 error."""
+    """Test that requesting documents without any parameters returns a 422 validation error."""
     response = await async_client.get("/api/v1/documents", headers=owner_auth_header)
 
-    assert response.status_code == 400
-    assert "At least one of id, pmid, doi, or reference must be provided" in response.json()["detail"]
+    assert response.status_code == 422
+    # FastAPI validation error for missing required parameter workspace_id
 
 
 @pytest.mark.skip(reason="Document delete API is failing with 500 error")
diff --git a/extralit/tests/unit/api/test_workspace_documents_api.py b/extralit/tests/unit/api/test_workspace_documents_api.py
index 9fc874620..c4a40369a 100644
--- a/extralit/tests/unit/api/test_workspace_documents_api.py
+++ b/extralit/tests/unit/api/test_workspace_documents_api.py
@@ -157,19 +157,25 @@ def test_document_create(self, mock_client, sample_workspace_id):
     def test_document_get(self, mock_client, sample_document_id):
         """Test Document.get() class method."""
         client, documents_api = mock_client
+        sample_workspace_id = uuid4()
 
-        # Mock the API response
+        # Mock the API response - now returns a list
         retrieved_model = DocumentModel(
-            id=sample_document_id, workspace_id=uuid4(), reference="Retrieved2023", file_name="retrieved.pdf"
+            id=sample_document_id,
+            workspace_id=sample_workspace_id,
+            reference="Retrieved2023",
+            file_name="retrieved.pdf",
         )
-        documents_api.get.return_value = retrieved_model
+        documents_api.get_document.return_value = retrieved_model
 
-        # Call get
+        # Call get with required workspace_id
         with patch("extralit.documents._resource.Extralit._get_default", return_value=client):
-            doc = Document.get(id=sample_document_id)
+            doc = Document.get(workspace_id=sample_workspace_id, id=sample_document_id)
 
-        # Verify API was called
-        documents_api.get.assert_called_once_with(sample_document_id)
+        # Verify API was called with the new unified method
+        documents_api.get_document.assert_called_once_with(
+            {"workspace_id": str(sample_workspace_id), "id": str(sample_document_id)}
+        )
 
         # Verify result
         assert doc.id == sample_document_id
@@ -224,12 +230,12 @@ def test_document_factory_methods(self, mock_client, sample_workspace_id):
         client, documents_api = mock_client
 
         with patch("extralit.documents._resource.Extralit._get_default", return_value=client):
-            # Test from_pmid
+            # Test from_pmid - now requires workspace_id as required parameter
             doc_pmid = Document.from_pmid(pmid="12345678", workspace_id=sample_workspace_id)
             assert doc_pmid.pmid == "12345678"
             assert doc_pmid.workspace_id == sample_workspace_id
 
-            # Test from_doi
+            # Test from_doi - now requires workspace_id as required parameter
             doc_doi = Document.from_doi(doi="10.1234/example", workspace_id=sample_workspace_id)
             assert doc_doi.doi == "10.1234/example"
             assert doc_doi.workspace_id == sample_workspace_id

From 2ad6ce515293687d4e4eb1da13c3ce33f8c482e5 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 15:00:21 -0700
Subject: [PATCH 7/8] fix tests

---
 .../tests/unit/api/test_workspace_documents_api.py   | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/extralit/tests/unit/api/test_workspace_documents_api.py b/extralit/tests/unit/api/test_workspace_documents_api.py
index c4a40369a..8ae736631 100644
--- a/extralit/tests/unit/api/test_workspace_documents_api.py
+++ b/extralit/tests/unit/api/test_workspace_documents_api.py
@@ -187,7 +187,7 @@ def test_document_update(self, mock_client, sample_document_id, sample_workspace
         client, documents_api = mock_client
 
         # Create a document with existing data
-        doc = Document(id=sample_document_id, workspace_id=sample_workspace_id, reference="Original2023", client=client)
+        doc = Document(workspace_id=sample_workspace_id, reference="Original2023", id=sample_document_id, client=client)
 
         # Update some fields
         doc.reference = "Updated2023"
@@ -217,7 +217,7 @@ def test_document_delete(self, mock_client, sample_document_id, sample_workspace
         client, documents_api = mock_client
 
         # Create a document
-        doc = Document(id=sample_document_id, workspace_id=sample_workspace_id, reference="ToDelete2023", client=client)
+        doc = Document(workspace_id=sample_workspace_id, reference="ToDelete2023", id=sample_document_id, client=client)
 
         # Call delete
         doc.delete()
@@ -230,13 +230,13 @@ def test_document_factory_methods(self, mock_client, sample_workspace_id):
         client, documents_api = mock_client
 
         with patch("extralit.documents._resource.Extralit._get_default", return_value=client):
-            # Test from_pmid - now requires workspace_id as required parameter
-            doc_pmid = Document.from_pmid(pmid="12345678", workspace_id=sample_workspace_id)
+            # Test the new unified get method with pmid
+            doc_pmid = Document.get(workspace_id=sample_workspace_id, pmid="12345678")
             assert doc_pmid.pmid == "12345678"
             assert doc_pmid.workspace_id == sample_workspace_id
 
-            # Test from_doi - now requires workspace_id as required parameter
-            doc_doi = Document.from_doi(doi="10.1234/example", workspace_id=sample_workspace_id)
+            # Test the new unified get method with doi
+            doc_doi = Document.get(workspace_id=sample_workspace_id, doi="10.1234/example")
             assert doc_doi.doi == "10.1234/example"
             assert doc_doi.workspace_id == sample_workspace_id
 

From 7913809e41ffda119bfb0c23288ca38028d1de48 Mon Sep 17 00:00:00 2001
From: JonnyTran <nhat.c.tran@gmail.com>
Date: Tue, 12 Aug 2025 15:06:43 -0700
Subject: [PATCH 8/8] fix tests

---
 .../unit/api/test_workspace_documents_api.py  | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/extralit/tests/unit/api/test_workspace_documents_api.py b/extralit/tests/unit/api/test_workspace_documents_api.py
index 8ae736631..14daf5a85 100644
--- a/extralit/tests/unit/api/test_workspace_documents_api.py
+++ b/extralit/tests/unit/api/test_workspace_documents_api.py
@@ -166,14 +166,14 @@ def test_document_get(self, mock_client, sample_document_id):
             reference="Retrieved2023",
             file_name="retrieved.pdf",
         )
-        documents_api.get_document.return_value = retrieved_model
+        documents_api.get.return_value = retrieved_model
 
         # Call get with required workspace_id
         with patch("extralit.documents._resource.Extralit._get_default", return_value=client):
             doc = Document.get(workspace_id=sample_workspace_id, id=sample_document_id)
 
         # Verify API was called with the new unified method
-        documents_api.get_document.assert_called_once_with(
+        documents_api.get.assert_called_once_with(
             {"workspace_id": str(sample_workspace_id), "id": str(sample_document_id)}
         )
 
@@ -229,13 +229,29 @@ def test_document_factory_methods(self, mock_client, sample_workspace_id):
         """Test Document factory methods."""
         client, documents_api = mock_client
 
+        # Mock the API responses for the get method
+        pmid_model = DocumentModel(
+            id=uuid4(),
+            workspace_id=sample_workspace_id,
+            reference="PMID:12345678",
+            pmid="12345678",
+        )
+        doi_model = DocumentModel(
+            id=uuid4(),
+            workspace_id=sample_workspace_id,
+            reference="DOI:10.1234/example",
+            doi="10.1234/example",
+        )
+
         with patch("extralit.documents._resource.Extralit._get_default", return_value=client):
             # Test the new unified get method with pmid
+            documents_api.get.return_value = pmid_model
             doc_pmid = Document.get(workspace_id=sample_workspace_id, pmid="12345678")
             assert doc_pmid.pmid == "12345678"
             assert doc_pmid.workspace_id == sample_workspace_id
 
             # Test the new unified get method with doi
+            documents_api.get.return_value = doi_model
             doc_doi = Document.get(workspace_id=sample_workspace_id, doi="10.1234/example")
             assert doc_doi.doi == "10.1234/example"
             assert doc_doi.workspace_id == sample_workspace_id