GPTRouter 0.3.7__tar.gz → 0.3.9__tar.gz

{GPTRouter-0.3.7 → GPTRouter-0.3.9}/GPTRouter.egg-info/PKG-INFO

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

{GPTRouter-0.3.7 → GPTRouter-0.3.9}/PKG-INFO

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

{GPTRouter-0.3.7 → GPTRouter-0.3.9}/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,
@@ -391,9 +419,18 @@ class GPTRouterClient:
         )
         generated_images = api_response.get("response", [])
         if isinstance(generated_images, dict):
-            generated_images = generated_images.get("artifacts", [])
+            generated_images = self._extract_images_from_dict(generated_images)
 
         return [
             parse_obj_as(ImageGenerationResponse, generated_img)
             for generated_img in generated_images
         ]
+
+    def _extract_images_from_dict(self, image_dict: dict) -> List[dict]:
+        if artifacts := image_dict.get("artifacts"):
+            return artifacts
+        
+        if image_urls := image_dict.get("imageUrls"):
+            return [{"url": url} for url in image_urls]
+
+        return []

{GPTRouter-0.3.7 → GPTRouter-0.3.9}/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"
@@ -15,6 +17,10 @@ class ProvidersEnum(Enum):
     STABLE_DIFFUSION = "stable-diffusion"
     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"
@@ -43,6 +49,11 @@ class ModelsEnum(Enum):
     STABLE_DIFFUSION_XL = "stable-diffusion-xl-1024-v1-0"
     STABLE_DIFFUSION_V6 = "stable-diffusion-v1-6"
 
+    REPLICATE_FLUX = "black-forest-labs/flux-dev"
+
+    OPENAI_EMBEDDINGS_ADA_002 = "text-embedding-ada-002"
+    AZURE_OPENAI_EMBEDDINGS_ADA_002 = "text-embedding-ada-002"
+
 class StreamingEventType(Enum):
     UPDATE = "update"
     META = "meta"
@@ -150,6 +161,9 @@ class GPTRouterEngines(Enum):
         ProvidersEnum.STABLE_DIFFUSION.value, "stable-diffusion-xl-1024-v1-0"
     )
     DALLE_3 = ModelDetails(ProvidersEnum.DALL_E.value, "dall-e-3")
+    REPLICATE_FLUX = ModelDetails(
+        ProvidersEnum.REPLICATE.value, "black-forest-labs/flux-dev"
+    )
 
     LLAMA_3_70B_NITRO = ModelDetails(
         ProvidersEnum.OPENROUTER, "meta-llama/llama-3-70b-instruct:nitro"

{GPTRouter-0.3.7 → GPTRouter-0.3.9}/gpt_router/models.py

@@ -109,3 +109,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.7 → GPTRouter-0.3.9}/setup.py

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