From 1f8fc6d2a7ad4d0801263f8394f87f76dc5fa1b3 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 2 Jan 2024 12:23:49 +0530
Subject: [PATCH] (feat) litellm add types for completion, embedding request

---
 litellm/types/completion.py | 36 ++++++++++++++++++++++++++++++++++++
 litellm/types/embedding.py  | 23 +++++++++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 litellm/types/completion.py
 create mode 100644 litellm/types/embedding.py

diff --git a/litellm/types/completion.py b/litellm/types/completion.py
new file mode 100644
index 000000000..5eac90575
--- /dev/null
+++ b/litellm/types/completion.py
@@ -0,0 +1,36 @@
+from typing import List, Optional, Union
+
+from pydantic import BaseModel, validator
+
+
+class CompletionRequest(BaseModel):
+    model: str
+    messages: List[str] = []
+    timeout: Optional[Union[float, int]] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    n: Optional[int] = None
+    stream: Optional[bool] = None
+    stop: Optional[dict] = None
+    max_tokens: Optional[float] = None
+    presence_penalty: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    logit_bias: Optional[dict] = None
+    user: Optional[str] = None
+    response_format: Optional[dict] = None
+    seed: Optional[int] = None
+    tools: Optional[List[str]] = None
+    tool_choice: Optional[str] = None
+    logprobs: Optional[bool] = None
+    top_logprobs: Optional[int] = None
+    deployment_id: Optional[str] = None
+    functions: Optional[List[str]] = None
+    function_call: Optional[str] = None
+    base_url: Optional[str] = None
+    api_version: Optional[str] = None
+    api_key: Optional[str] = None
+    model_list: Optional[List[str]] = None
+
+    class Config:
+        # allow kwargs
+        extra = "allow"
diff --git a/litellm/types/embedding.py b/litellm/types/embedding.py
new file mode 100644
index 000000000..9db0ef290
--- /dev/null
+++ b/litellm/types/embedding.py
@@ -0,0 +1,23 @@
+from typing import List, Optional, Union
+
+from pydantic import BaseModel, validator
+
+
+class EmbeddingRequest(BaseModel):
+    model: str
+    input: List[str] = []
+    timeout: int = 600
+    api_base: Optional[str] = None
+    api_version: Optional[str] = None
+    api_key: Optional[str] = None
+    api_type: Optional[str] = None
+    caching: bool = False
+    user: Optional[str] = None
+    custom_llm_provider: Optional[Union[str, dict]] = None
+    litellm_call_id: Optional[str] = None
+    litellm_logging_obj: Optional[dict] = None
+    logger_fn: Optional[str] = None
+
+    class Config:
+        # allow kwargs
+        extra = "allow"