From b9515c3b9641ecd4df294bac04e2863bb02e8b0b Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 28 May 2024 12:53:46 -0700
Subject: [PATCH] feat - add batches types

---
 litellm/types/llms/openai.py | 52 +++++++++++++++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
index 1c60ad6dbb..60f7a90417 100644
--- a/litellm/types/llms/openai.py
+++ b/litellm/types/llms/openai.py
@@ -6,7 +6,7 @@ from typing import (
     Literal,
     Iterable,
 )
-from typing_extensions import override, Required
+from typing_extensions import override, Required, Dict
 from pydantic import BaseModel
 
 from openai.types.beta.threads.message_content import MessageContent
@@ -146,3 +146,53 @@ class Thread(BaseModel):
 
     object: Literal["thread"]
     """The object type, which is always `thread`."""
+
+
+# OpenAI Batches Types
+class CreateBatchRequest(BaseModel):
+    """
+    CreateBatchRequest
+    """
+
+    completion_window: Literal["24h"]
+    endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"]
+    input_file_id: str
+    metadata: Optional[Dict[str, str]] = None
+    extra_headers: Optional[Dict[str, str]] = None
+    extra_body: Optional[Dict[str, str]] = None
+    timeout: Optional[float] = None
+
+
+class RetrieveBatchRequest(BaseModel):
+    """
+    RetrieveBatchRequest
+    """
+
+    batch_id: str
+    extra_headers: Optional[Dict[str, str]] = None
+    extra_body: Optional[Dict[str, str]] = None
+    timeout: Optional[float] = None
+
+
+class CancelBatchRequest(BaseModel):
+    """
+    CancelBatchRequest
+    """
+
+    batch_id: str
+    extra_headers: Optional[Dict[str, str]] = None
+    extra_body: Optional[Dict[str, str]] = None
+    timeout: Optional[float] = None
+
+
+class ListBatchRequest(BaseModel):
+    """
+    ListBatchRequest - List your organization's batches
+    Calls https://api.openai.com/v1/batches
+    """
+
+    after: Optional[str] = None
+    limit: Optional[int] = 20
+    extra_headers: Optional[Dict[str, str]] = None
+    extra_body: Optional[Dict[str, str]] = None
+    timeout: Optional[float] = None