From ac2bfcc52e518cde4f73f5bb3b7eae0b00287ec5 Mon Sep 17 00:00:00 2001
From: abhijit <abhijitastlar@gmail.com>
Date: Thu, 29 May 2025 19:21:13 +0530
Subject: [PATCH] fix: update route models

---
 javelin_sdk/client.py |  3 +++
 javelin_sdk/models.py | 54 ++++++++++++++++++-------------------------
 2 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/javelin_sdk/client.py b/javelin_sdk/client.py
index f8843d8..c552245 100644
--- a/javelin_sdk/client.py
+++ b/javelin_sdk/client.py
@@ -582,6 +582,9 @@ def register_bedrock(
         self.bedrock_session = bedrock_session
         self.bedrock_runtime_client = bedrock_runtime_client
 
+        if not route_name:
+            route_name = "awsbedrock"
+
         # Store the default bedrock route
         if route_name is not None:
             self.use_default_bedrock_route = True
diff --git a/javelin_sdk/models.py b/javelin_sdk/models.py
index 211a3d0..046cfc1 100644
--- a/javelin_sdk/models.py
+++ b/javelin_sdk/models.py
@@ -114,44 +114,32 @@ class ContentFilter(BaseModel):
     )
 
 
-class RouteConfig(BaseModel):
-    rate_limit: Optional[int] = Field(
-        default=None, description="Rate limit for the route"
-    )
-    owner: Optional[str] = Field(default=None, description="Owner of the route")
-    organization: Optional[str] = Field(
-        default=None, description="Organization associated with the route"
-    )
-    archive: Optional[bool] = Field(
-        default=None, description="Whether archiving is enabled"
-    )
-    retries: Optional[int] = Field(
-        default=None, description="Number of retries for the route"
-    )
-    llm_cache: bool = Field(False, description="Whether LLM cache is enabled")
-    role_to_assume: Optional[str] = Field(
-        None, description="Role to assume for the route"
-    )
-    enable_telemetry: Optional[bool] = Field(
-        None, description="Whether telemetry is enabled"
-    )
+class ArchivePolicy(BaseModel):
+    enabled: Optional[bool] = Field(default=None, description="Whether archiving is enabled")
     retention: Optional[int] = Field(default=None, description="Data retention period")
+
+
+class Policy(BaseModel):
+    dlp: Optional[Dlp] = Field(default=None, description="DLP configuration")
+    archive: Optional[ArchivePolicy] = Field(default=None, description="Archive policy configuration")
+    enabled: Optional[bool] = Field(default=None, description="Whether the policy is enabled")
+    prompt_safety: Optional[PromptSafety] = Field(default=None, description="Prompt Safety Description")
+    content_filter: Optional[ContentFilter] = Field(default=None, description="Content Filter Description")
+    security_filters: Optional[SecurityFilters] = Field(default=None, description="Security Filters Description")
+
+
+class RouteConfig(BaseModel):
+    budget: Optional[Budget] = Field(default=None, description="Budget configuration")
+    policy: Optional[Policy] = Field(default=None, description="Policy configuration")
+    retries: Optional[int] = Field(default=None, description="Number of retries for the route")
+    rate_limit: Optional[int] = Field(default=None, description="Rate limit for the route")
+    unified_endpoint: Optional[bool] = Field(default=None, description="Whether unified endpoint is enabled")
     request_chain: Optional[Dict[str, Any]] = Field(
         None, description="Request chain configuration"
     )
     response_chain: Optional[Dict[str, Any]] = Field(
         None, description="Response chain configuration"
     )
-    dlp: Optional[Dlp] = Field(default=None, description="DLP configuration")
-    content_filter: Optional[ContentFilter] = Field(
-        default=None, description="Content Filter Description"
-    )
-    prompt_safety: Optional[PromptSafety] = Field(
-        default=None, description="Prompt Safety Description"
-    )
-    security_filters: Optional[SecurityFilters] = Field(
-        default=None, description="Security Filters Description"
-    )
 
 
 class Model(BaseModel):
@@ -350,6 +338,10 @@ class Template(BaseModel):
 class Templates(BaseModel):
     templates: List[Template] = Field(default=[], description="List of templates")
 
+class SecretType(str, Enum):
+    AWS = "aws"
+    KUBERNETES = "kubernetes"
+
 
 class Secret(BaseModel):
     api_key: str = Field(default=None, description="Key of the Secret")