From bd838434a3a5b6972b52d2a8361a5e6ebed8ff57 Mon Sep 17 00:00:00 2001 From: Uday-sidagana Date: Fri, 12 Jun 2026 22:26:50 +0530 Subject: [PATCH] Add SDK support for enhanced task status --- sdks/python/README.md | 8 ++++ sdks/python/openrag_sdk/__init__.py | 2 + sdks/python/openrag_sdk/documents.py | 35 +++++++++++++++- sdks/python/openrag_sdk/models.py | 6 +++ sdks/typescript/README.md | 8 ++++ sdks/typescript/src/documents.ts | 49 +++++++++++++++++++---- sdks/typescript/src/index.ts | 2 + sdks/typescript/src/types.ts | 4 ++ sdks/typescript/tests/integration.test.ts | 15 +++++++ tests/integration/sdk/test_documents.py | 8 ++++ 10 files changed, 129 insertions(+), 8 deletions(-) diff --git a/sdks/python/README.md b/sdks/python/README.md index 664a81d6b..1afe0f9ed 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -161,6 +161,14 @@ final_status = await client.documents.wait_for_task(result.task_id) print(f"Status: {final_status.status}") print(f"Successful files: {final_status.successful_files}") +# Get structured failure metadata for a task +enhanced_status = await client.documents.get_task_status_enhanced(result.task_id) +print(enhanced_status.files) + +# List tasks with structured failure metadata +enhanced_tasks = await client.documents.list_tasks_enhanced() +print(len(enhanced_tasks.tasks)) + # Ingest from file object with open("./report.pdf", "rb") as f: result = await client.documents.ingest(file=f, filename="report.pdf") diff --git a/sdks/python/openrag_sdk/__init__.py b/sdks/python/openrag_sdk/__init__.py index 86bddffc2..e9bfdebc3 100644 --- a/sdks/python/openrag_sdk/__init__.py +++ b/sdks/python/openrag_sdk/__init__.py @@ -65,6 +65,7 @@ Source, SourcesEvent, StreamEvent, + TaskListResponse, UpdateKnowledgeFilterOptions, ) @@ -93,6 +94,7 @@ "SearchResult", "SearchFilters", "IngestResponse", + "TaskListResponse", "DeleteDocumentResponse", "Conversation", "ConversationDetail", diff --git a/sdks/python/openrag_sdk/documents.py b/sdks/python/openrag_sdk/documents.py index 414144c3d..42740b82b 100644 --- a/sdks/python/openrag_sdk/documents.py +++ b/sdks/python/openrag_sdk/documents.py @@ -5,7 +5,12 @@ from typing import TYPE_CHECKING, BinaryIO from .exceptions import NotFoundError -from .models import DeleteDocumentResponse, IngestResponse, IngestTaskStatus +from .models import ( + DeleteDocumentResponse, + IngestResponse, + IngestTaskStatus, + TaskListResponse, +) if TYPE_CHECKING: from .client import OpenRAGClient @@ -97,6 +102,34 @@ async def get_task_status(self, task_id: str) -> IngestTaskStatus: data = response.json() return IngestTaskStatus(**data) + async def get_task_status_enhanced(self, task_id: str) -> IngestTaskStatus: + """ + Get the status of an ingestion task with structured failure metadata. + + Args: + task_id: The task ID returned from ingest(). + + Returns: + IngestTaskStatus with current task status and classified failure details. + """ + response = await self._client._request( + "GET", + f"/api/v1/tasks/{task_id}/enhanced", + ) + data = response.json() + return IngestTaskStatus(**data) + + async def list_tasks_enhanced(self) -> TaskListResponse: + """ + List ingestion tasks with structured failure metadata. + + Returns: + TaskListResponse with enhanced task status entries. + """ + response = await self._client._request("GET", "/api/v1/tasks/enhanced") + data = response.json() + return TaskListResponse(**data) + async def wait_for_task( self, task_id: str, diff --git a/sdks/python/openrag_sdk/models.py b/sdks/python/openrag_sdk/models.py index b0eee0864..c486fcecf 100644 --- a/sdks/python/openrag_sdk/models.py +++ b/sdks/python/openrag_sdk/models.py @@ -89,6 +89,12 @@ class IngestTaskStatus(BaseModel): files: dict = {} # Detailed per-file status +class TaskListResponse(BaseModel): + """Response from listing ingestion tasks.""" + + tasks: list[IngestTaskStatus] + + class DeleteDocumentResponse(BaseModel): """Response from document deletion.""" diff --git a/sdks/typescript/README.md b/sdks/typescript/README.md index 9526ca923..16a6edc98 100644 --- a/sdks/typescript/README.md +++ b/sdks/typescript/README.md @@ -188,6 +188,14 @@ const finalStatus = await client.documents.waitForTask(result.task_id); console.log(`Status: ${finalStatus.status}`); console.log(`Successful files: ${finalStatus.successful_files}`); +// Get structured failure metadata for a task +const enhancedStatus = await client.documents.getTaskStatusEnhanced(result.task_id); +console.log(enhancedStatus.files); + +// List tasks with structured failure metadata +const enhancedTasks = await client.documents.listTasksEnhanced(); +console.log(enhancedTasks.tasks.length); + // Ingest from File object (browser) const file = new File([...], "report.pdf"); const result = await client.documents.ingest({ diff --git a/sdks/typescript/src/documents.ts b/sdks/typescript/src/documents.ts index 274fce9a2..978d14af2 100644 --- a/sdks/typescript/src/documents.ts +++ b/sdks/typescript/src/documents.ts @@ -9,6 +9,7 @@ import type { IngestResponse, IngestTaskStatus, NotFoundError, + TaskListResponse, } from "./types"; export interface IngestOptions { @@ -100,14 +101,48 @@ export class DocumentsClient { `/api/v1/tasks/${taskId}` ); const data = await response.json(); + return this.parseTaskStatus(data); + } + + /** + * Get the status of an ingestion task with structured failure metadata. + * + * @param taskId - The task ID returned from ingest(). + * @returns IngestTaskStatus with current task status and classified failure details. + */ + async getTaskStatusEnhanced(taskId: string): Promise { + const response = await this.client._request( + "GET", + `/api/v1/tasks/${taskId}/enhanced` + ); + const data = await response.json(); + return this.parseTaskStatus(data); + } + + /** + * List ingestion tasks with structured failure metadata. + * + * @returns TaskListResponse with enhanced task status entries. + */ + async listTasksEnhanced(): Promise { + const response = await this.client._request("GET", "/api/v1/tasks/enhanced"); + const data = await response.json(); return { - task_id: data.task_id, - status: data.status, - total_files: data.total_files ?? 0, - processed_files: data.processed_files ?? 0, - successful_files: data.successful_files ?? 0, - failed_files: data.failed_files ?? 0, - files: data.files ?? {}, + tasks: (data.tasks ?? []).map((task: Record) => + this.parseTaskStatus(task) + ), + }; + } + + private parseTaskStatus(data: Record): IngestTaskStatus { + return { + task_id: data["task_id"] as string, + status: data["status"] as string, + total_files: (data["total_files"] as number | undefined) ?? 0, + processed_files: (data["processed_files"] as number | undefined) ?? 0, + successful_files: (data["successful_files"] as number | undefined) ?? 0, + failed_files: (data["failed_files"] as number | undefined) ?? 0, + files: (data["files"] as Record | undefined) ?? {}, }; } diff --git a/sdks/typescript/src/index.ts b/sdks/typescript/src/index.ts index 3283f118c..cf516c037 100644 --- a/sdks/typescript/src/index.ts +++ b/sdks/typescript/src/index.ts @@ -64,6 +64,8 @@ export { SearchResult, // Document types IngestResponse, + IngestTaskStatus, + TaskListResponse, DeleteDocumentResponse, // Conversation types Conversation, diff --git a/sdks/typescript/src/types.ts b/sdks/typescript/src/types.ts index 2d2336ee8..183adc654 100644 --- a/sdks/typescript/src/types.ts +++ b/sdks/typescript/src/types.ts @@ -71,6 +71,10 @@ export interface IngestTaskStatus { files: Record; } +export interface TaskListResponse { + tasks: IngestTaskStatus[]; +} + export interface DeleteDocumentResponse { success: boolean; deleted_chunks: number; diff --git a/sdks/typescript/tests/integration.test.ts b/sdks/typescript/tests/integration.test.ts index 221151edc..1bf64c0a4 100644 --- a/sdks/typescript/tests/integration.test.ts +++ b/sdks/typescript/tests/integration.test.ts @@ -362,6 +362,21 @@ describe.skipIf(SKIP_TESTS)("OpenRAG TypeScript SDK Integration", () => { expect(finalStatus.status).toBeDefined(); }); + it("should get enhanced task status", async () => { + const result = await client.documents.ingest({ + filePath: testFilePath, + wait: false, + }); + const taskId = (result as any).task_id; + expect(taskId).toBeDefined(); + + const enhancedStatus = await client.documents.getTaskStatusEnhanced(taskId); + expect(enhancedStatus.status).toBeDefined(); + + const enhancedTasks = await client.documents.listTasksEnhanced(); + expect(Array.isArray(enhancedTasks.tasks)).toBe(true); + }); + it("should delete a document", async () => { // First ingest (wait for completion) const ingestResult = await client.documents.ingest({ filePath: testFilePath }); diff --git a/tests/integration/sdk/test_documents.py b/tests/integration/sdk/test_documents.py index f46ad25fd..5c787d6d3 100644 --- a/tests/integration/sdk/test_documents.py +++ b/tests/integration/sdk/test_documents.py @@ -116,6 +116,14 @@ async def test_task_status_polling(self, client, tmp_path): status = await client.documents.get_task_status(task_response.task_id) assert status.status is not None + enhanced_status = await client.documents.get_task_status_enhanced( + task_response.task_id + ) + assert enhanced_status.status is not None + + enhanced_tasks = await client.documents.list_tasks_enhanced() + assert isinstance(enhanced_tasks.tasks, list) + final = await client.documents.wait_for_task(task_response.task_id) assert final.status in ("completed", "failed")