GPTRouter 0.3.8__tar.gz → 0.4.0__tar.gz

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/GPTRouter.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: GPTRouter
-Version: 0.3.8
+Version: 0.4.0
 Summary: A Python package for working with GPTRouter APIs
 Author-email: sirjanpreet@writesonic.com
 Classifier: Development Status :: 3 - Alpha

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: GPTRouter
-Version: 0.3.8
+Version: 0.4.0
 Summary: A Python package for working with GPTRouter APIs
 Author-email: sirjanpreet@writesonic.com
 Classifier: Development Status :: 3 - Alpha

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/gpt_router/client.py

@@ -13,6 +13,9 @@ from gpt_router.models import (
     ModelGenerationRequest,
     GenerationResponse,
     ChunkedGenerationResponse,
+    ModelRouterEmbeddingsGenerationRequest,
+    ModelRouterGenerationResponse,
+    ModelRouterEmbeddingsGenerationResponse
 )
 from gpt_router.exceptions import (
     GPTRouterApiTimeoutError,
@@ -71,6 +74,31 @@ class GPTRouterClient:
             request_timeout if stream_read_timeout is None else stream_read_timeout
         )
 
+    async def agenerate_embeddings(
+        self,
+        *,
+        ordered_generation_requests: List[
+            ModelRouterEmbeddingsGenerationRequest
+        ],
+        model_router_metadata: Optional[GPTRouterMetadata] = None,
+    ) -> ModelRouterGenerationResponse:
+        api_payload = {
+            "data": [
+                {
+                    **request.dict(),
+                    "order": index + 1,
+                }
+                for index, request in enumerate(ordered_generation_requests)
+            ],
+        }
+        api_payload = self.add_metadata_info(api_payload, model_router_metadata)
+        response = await self._async_api_call(
+            path="/v1/generate/embeddings",
+            method="POST",
+            payload=api_payload,
+        )
+        return ModelRouterEmbeddingsGenerationResponse.parse_obj(response)
+
     def add_metadata_info(
         self,
         payload: dict,

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/gpt_router/enums.py

@@ -1,6 +1,8 @@
 from collections import namedtuple
 from enum import Enum
 
+
+
 class ProvidersEnum(Enum):
     OPENAI = "openai"
     CHAT_OPENAI = "chat_openai"
@@ -16,6 +18,9 @@ class ProvidersEnum(Enum):
     DALL_E = "dall-e"
     OPENROUTER = "openrouter"
     REPLICATE = "replicate"
+    OPENAI_EMBEDDINGS = "openai_embeddings"
+    AZURE_OPENAI_EMBEDDINGS = "azure_openai_embeddings"
+
 
 class ModelsEnum(Enum):
     OPENAI_GPT_35_TURBO_16K = "gpt-3.5-turbo-16k"
@@ -45,6 +50,12 @@ class ModelsEnum(Enum):
     STABLE_DIFFUSION_V6 = "stable-diffusion-v1-6"
 
     REPLICATE_FLUX = "black-forest-labs/flux-dev"
+    REPLICATE_FLUX_PRO = "black-forest-labs/flux-pro"
+    REPLICATE_FLUX_11_PRO = "black-forest-labs/flux-1.1-pro"
+    REPLICATE_FLUX_SCHNELL = "black-forest-labs/flux-schnell"
+
+    OPENAI_EMBEDDINGS_ADA_002 = "text-embedding-ada-002"
+    AZURE_OPENAI_EMBEDDINGS_ADA_002 = "text-embedding-ada-002"
 
 class StreamingEventType(Enum):
     UPDATE = "update"
@@ -156,6 +167,15 @@ class GPTRouterEngines(Enum):
     REPLICATE_FLUX = ModelDetails(
         ProvidersEnum.REPLICATE.value, "black-forest-labs/flux-dev"
     )
+    REPLICATE_FLUX_PRO = ModelDetails(
+        ProvidersEnum.REPLICATE.value, "black-forest-labs/flux-pro"
+    )
+    REPLICATE_FLUX_11_PRO = ModelDetails(
+        ProvidersEnum.REPLICATE.value, "black-forest-labs/flux-1.1-pro"
+    )
+    REPLICATE_FLUX_SCHNELL = ModelDetails(
+        ProvidersEnum.REPLICATE.value, "black-forest-labs/flux-schnell"
+    )
 
     LLAMA_3_70B_NITRO = ModelDetails(
         ProvidersEnum.OPENROUTER, "meta-llama/llama-3-70b-instruct:nitro"

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/gpt_router/models.py

@@ -89,6 +89,7 @@ class ImageGenerationRequest(BaseModel):
     num_images: int = Field(alias="numImages", default=1)
     width: Optional[int] = Field(None, alias="width")
     height: Optional[int] = Field(None, alias="height")
+    num_inference_steps: Optional[int] = Field(None, alias="numInferenceSteps")
 
     class Config:
         allow_population_by_field_name = True
@@ -109,3 +110,47 @@ class ImageGenerationResponse(BaseModel):
     url: str = None
     base64: str = None
     finish_reason: Optional[str] = Field(default="SUCCESS", alias='finishReason')
+
+
+class ModelRouterGenerationRequest(BaseModel):
+    model_name: str
+    model_provider: str
+    prompt_params: GenerationParams = Field(alias="promptParams")
+    should_validate_token_count: Optional[bool] = Field(
+        default=False, alias="shouldValidateTokenCount"
+    )
+
+    class Config:
+        allow_population_by_field_name = True
+
+    def dict(self, *args, **kwargs):
+        # Generate the standard dict first.
+        result = super().dict(*args, **kwargs, exclude_none=True, by_alias=True)
+
+        result["providerName"] = self.model_provider
+        result["modelName"] = self.model_name
+
+        return result
+
+class ModelRouterEmbeddingsGenerationParams(BaseModel):
+    input: List[str]
+
+
+class ModelRouterEmbeddingsGenerationRequest(ModelRouterGenerationRequest):
+    prompt_params: ModelRouterEmbeddingsGenerationParams = Field(alias="promptParams")
+
+class ModelRouterGenerationResponse(BaseModel):
+    id: str
+    choices: List[Choice]
+    model: str
+    provider_id: Optional[str] = Field(None, alias="providerId")
+    model_id: Optional[str] = Field(None, alias="modelId")
+    meta: Optional[Usage]
+
+
+
+class EmbeddingsChoice(BaseModel):
+    embedding: List[float]
+
+class ModelRouterEmbeddingsGenerationResponse(BaseModel):
+    choices: List[EmbeddingsChoice]

{GPTRouter-0.3.8 → GPTRouter-0.4.0}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name="GPTRouter",
-    version="0.3.8",
+    version="0.4.0",
     description="A Python package for working with GPTRouter APIs",
     long_description=long_description,
     long_description_content_type="text/markdown",