Merge branch 'main' into release/v3.3.1

awslabs · Dec 2, 2024 · d1a3e8d · d1a3e8d
2 parents 3f969c7 + 7613ced
commit d1a3e8d
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 16 deletions.
diff --git a/lib/rag/state_machine/ingest-pipeline.ts b/lib/rag/state_machine/ingest-pipeline.ts
@@ -37,6 +37,7 @@ import { Rule, Schedule, EventPattern, RuleTargetInput, EventField } from 'aws-c
 import { SfnStateMachine } from 'aws-cdk-lib/aws-events-targets';
 import { RagRepositoryType } from '../../schema';
 import * as kms from 'aws-cdk-lib/aws-kms';
+import * as cdk from 'aws-cdk-lib';
 
 type PipelineConfig = {
     chunkOverlap: number;
@@ -111,8 +112,8 @@ export class IngestPipelineStateMachine extends Construct {
             effect: Effect.ALLOW,
             actions: ['s3:GetObject', 's3:ListBucket'],
             resources: [
-                `arn:aws:s3:::${pipelineConfig.s3Bucket}`,
-                `arn:aws:s3:::${pipelineConfig.s3Bucket}/*`
+                `arn:${cdk.Aws.PARTITION}:s3:::${pipelineConfig.s3Bucket}`,
+                `arn:${cdk.Aws.PARTITION}:s3:::${pipelineConfig.s3Bucket}/*`
             ]
         });
 
@@ -176,10 +177,10 @@ export class IngestPipelineStateMachine extends Construct {
                     effect: Effect.ALLOW,
                     actions: ['ssm:GetParameter'],
                     resources: [
-                        `arn:aws:ssm:${process.env.CDK_DEFAULT_REGION}:${process.env.CDK_DEFAULT_ACCOUNT}:parameter${config.deploymentPrefix}/LisaServeRagPGVectorConnectionInfo`,
-                        `arn:aws:ssm:${process.env.CDK_DEFAULT_REGION}:${process.env.CDK_DEFAULT_ACCOUNT}:parameter${config.deploymentPrefix}/lisaServeRagRepositoryEndpoint`,
-                        `arn:aws:ssm:${process.env.CDK_DEFAULT_REGION}:${process.env.CDK_DEFAULT_ACCOUNT}:parameter${config.deploymentPrefix}/lisaServeRestApiUri`,
-                        `arn:aws:ssm:${process.env.CDK_DEFAULT_REGION}:${process.env.CDK_DEFAULT_ACCOUNT}:parameter${config.deploymentPrefix}/managementKeySecretName`
+                        `arn:${cdk.Aws.PARTITION}:ssm:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:parameter${config.deploymentPrefix}/LisaServeRagPGVectorConnectionInfo`,
+                        `arn:${cdk.Aws.PARTITION}:ssm:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:parameter${config.deploymentPrefix}/lisaServeRagRepositoryEndpoint`,
+                        `arn:${cdk.Aws.PARTITION}:ssm:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:parameter${config.deploymentPrefix}/lisaServeRestApiUri`,
+                        `arn:${cdk.Aws.PARTITION}:ssm:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:parameter${config.deploymentPrefix}/managementKeySecretName`
                     ]
                 }),
                 new PolicyStatement({

diff --git a/lib/schema.ts b/lib/schema.ts
@@ -537,6 +537,10 @@ const LiteLLMConfig = z.object({
         'Key string must be defined for model management operations, and it must start with "sk-".' +
       'This can be any string, and a random UUID is recommended. Example: sk-f132c7cc-059c-481b-b5ca-a42e191672aa',
     ),
+    general_settings: z.any().optional(),
+    litellm_settings: z.any().optional(),
+    router_settings: z.any().optional(),
+    environment_variables: z.any().optional()
 })
     .describe('Core LiteLLM configuration - see https://litellm.vercel.app/docs/proxy/configs#all-settings for more details about each field.');
 

diff --git a/lib/serve/rest-api/src/utils/generate_litellm_config.py b/lib/serve/rest-api/src/utils/generate_litellm_config.py
@@ -50,10 +50,14 @@ def generate_config(filepath: str) -> None:
     config_models = []  # ensure config_models is a list and not None
     config_models.extend(litellm_model_params)
     config_contents["model_list"] = config_models
-    config_contents["litellm_settings"] = {
-        "drop_params": True,  # drop unrecognized param instead of failing the request on it
-        "request_timeout": 600,
-    }
+    if "litellm_settings" not in config_contents:
+        config_contents["litellm_settings"] = {}
+    config_contents["litellm_settings"].update(
+        {
+            "drop_params": True,  # drop unrecognized param instead of failing the request on it
+            "request_timeout": 600,
+        }
+    )
 
     # Get database connection info
     db_param_response = ssm_client.get_parameter(Name=os.environ["LITELLM_DB_INFO_PS_NAME"])
@@ -65,13 +69,14 @@ def generate_config(filepath: str) -> None:
         f"/{db_params['dbName']}"
     )
 
-    config_contents.update(
+    if "general_settings" not in config_contents:
+        config_contents["general_settings"] = {}
+
+    config_contents["general_settings"].update(
         {
-            "general_settings": {
-                "store_model_in_db": True,
-                "database_url": connection_str,
-                "master_key": config_contents["db_key"],
-            }
+            "store_model_in_db": True,
+            "database_url": connection_str,
+            "master_key": config_contents["db_key"],
         }
     )