Feature/rest api sdk (#220)

Add RAG SDK functions

Author: bedanley

Makefile

@@ -154,6 +154,7 @@ createPythonEnvironment:
 installPythonRequirements:
 	pip3 install pip --upgrade
 	pip3 install -r requirements-dev.txt
+	pip3 install -e lisa-sdk
 
 
 ## Set up TypeScript interpreter environment

lambda/authorizer/lambda_functions.py

@@ -54,18 +54,22 @@ def lambda_handler(event: Dict[str, Any], context) -> Dict[str, Any]:  # type: i
     jwt_groups_property = os.environ.get("JWT_GROUPS_PROP", "")
 
     deny_policy = generate_policy(effect="Deny", resource=event["methodArn"])
-
+    groups: str
     if id_token in get_management_tokens():
-        allow_policy = generate_policy(effect="Allow", resource=event["methodArn"], username="lisa-management-token")
+        username = "lisa-management-token"
+        # Add management token to Admin groups
+        groups = json.dumps([admin_group])
+        allow_policy = generate_policy(effect="Allow", resource=event["methodArn"], username=username)
+        allow_policy["context"] = {"username": username, "groups": groups}
         logger.debug(f"Generated policy: {allow_policy}")
         return allow_policy
 
     if jwt_data := id_token_is_valid(id_token=id_token, client_id=client_id, authority=authority):
         is_admin_user = is_admin(jwt_data, admin_group, jwt_groups_property)
-        groups = get_property_path(jwt_data, jwt_groups_property)
+        groups = json.dumps(get_property_path(jwt_data, jwt_groups_property) or [])
         username = find_jwt_username(jwt_data)
         allow_policy = generate_policy(effect="Allow", resource=event["methodArn"], username=username)
-        allow_policy["context"] = {"username": username, "groups": json.dumps(groups or [])}
+        allow_policy["context"] = {"username": username, "groups": groups}
 
         if requested_resource.startswith("/models") and not is_admin_user:
             # non-admin users can still list models
@@ -185,10 +189,13 @@ def get_management_tokens() -> list[str]:
         secret_tokens.append(
             secrets_manager.get_secret_value(SecretId=secret_id, VersionStage="AWSCURRENT")["SecretString"]
         )
-        secret_tokens.append(
-            secrets_manager.get_secret_value(SecretId=secret_id, VersionStage="AWSPREVIOUS")["SecretString"]
-        )
+        try:
+            secret_tokens.append(
+                secrets_manager.get_secret_value(SecretId=secret_id, VersionStage="AWSPREVIOUS")["SecretString"]
+            )
+        except Exception:
+            logger.info("No previous management token version found")
     except ClientError as e:
-        logger.warn(f"Unable to fetch {secret_id}. {e.response['Error']['Code']}: {e.response['Error']['Message']}")
+        logger.warning(f"Unable to fetch {secret_id}. {e.response['Error']['Code']}: {e.response['Error']['Message']}")
 
     return secret_tokens

lambda/repository/lambda_functions.py

@@ -24,7 +24,7 @@
 from lisapy.langchain import LisaOpenAIEmbeddings
 from models.domain_objects import IngestionType, RagDocument
 from repository.rag_document_repo import RagDocumentRepository
-from utilities.common_functions import api_wrapper, get_cert_path, get_id_token, get_username, retry_config
+from utilities.common_functions import api_wrapper, get_cert_path, get_groups, get_id_token, get_username, retry_config
 from utilities.exceptions import HTTPException
 from utilities.file_processing import process_record
 from utilities.validation import validate_model_name, ValidationError
@@ -186,22 +186,14 @@ def _get_embeddings_pipeline(model_name: str) -> Any:
 
 @api_wrapper
 def list_all(event: dict, context: dict) -> List[Dict[str, Any]]:
-    """Return info on all available repositories.
-
-    Currently, there is no support for dynamic repositories so only a single OpenSearch repository
-    is returned.
-    """
-
-    user_groups = json.loads(event["requestContext"]["authorizer"]["groups"]) or []
+    """Return info on all available repositories."""
+    user_groups = get_groups(event)
     registered_repositories = get_registered_repositories()
-
-    return list(
-        filter(lambda repository: user_has_group(user_groups, repository["allowedGroups"]), registered_repositories)
-    )
+    return [repo for repo in registered_repositories if user_has_group(user_groups, repo["allowedGroups"])]
 
 
 def user_has_group(user_groups: List[str], allowed_groups: List[str]) -> bool:
-    """Returns if user groups has at least one intersections with allowed groups.
+    """Returns if user groups has at least one intersection with allowed groups.
 
     If allowed groups is empty this will return True.
     """
@@ -290,8 +282,8 @@ def delete_document(event: dict, context: dict) -> Dict[str, Any]:
     path_params = event.get("pathParameters", {})
     repository_id = path_params.get("repositoryId")
 
-    query_string_params = event["queryStringParameters"]
-    collection_id = query_string_params["collectionId"]
+    query_string_params = event.get("queryStringParameters", {})
+    collection_id = query_string_params.get("collectionId")
     document_id = query_string_params.get("documentId")
     document_name = query_string_params.get("documentName")
 

lambda/utilities/common_functions.py

@@ -360,7 +360,7 @@ def get_session_id(event: dict) -> str:
 
 def get_groups(event: Any) -> List[str]:
     """Get user groups from event."""
-    groups: List[str] = event.get("requestContext", {}).get("authorizer", {}).get("groups", [])
+    groups: List[str] = json.loads(event.get("requestContext", {}).get("authorizer", {}).get("groups", "[]"))
     return groups
 
 

lib/core/api_deployment.ts

@@ -14,11 +14,12 @@
   limitations under the License.
 */
 
-import { CfnOutput, Stack, StackProps, Aws } from 'aws-cdk-lib';
+import { Aws, CfnOutput, Stack, StackProps } from 'aws-cdk-lib';
 import { Deployment, RestApi } from 'aws-cdk-lib/aws-apigateway';
 import { Construct } from 'constructs';
 
 import { BaseProps } from '../schema';
+import { StringParameter } from 'aws-cdk-lib/aws-ssm';
 
 type LisaApiDeploymentStackProps = {
     restApiId: string;
@@ -43,8 +44,14 @@ export class LisaApiDeploymentStack extends Stack {
         // https://github.com/aws/aws-cdk/issues/25582
         (deployment as any).resource.stageName = config.deploymentStage;
 
+        const api_url = `https://${restApiId}.execute-api.${this.region}.${Aws.URL_SUFFIX}/${config.deploymentStage}`;
+        new StringParameter(this, 'LisaApiDeploymentStringParameter', {
+            parameterName: `${config.deploymentPrefix}/${config.deploymentName}/${config.appName}/LisaApiUrl`,
+            stringValue: api_url,
+            description: 'API Gateway URL for LISA',
+        });
         new CfnOutput(this, 'ApiUrl', {
-            value: `https://${restApiId}.execute-api.${this.region}.${Aws.URL_SUFFIX}/${config.deploymentStage}`,
+            value: api_url,
             description: 'API Gateway URL'
         });
     }

lib/core/layers/authorizer/requirements.txt

@@ -1,4 +1,4 @@
+# urllib3<2 // Provided by Lambda
+requests==2.32.3
 cryptography==43.0.1
 PyJWT==2.9.0
-requests==2.32.3
-urllib3<2

lib/core/layers/common/requirements.txt

@@ -1,3 +1,3 @@
-boto3>=1.34.131
-botocore>=1.34.131
-urllib3<2
+# boto3>=1.34.131 // Provided by Lambda
+# botocore>=1.34.131 // Provided by Lambda
+# urllib3<2 // Provided by Lambda

lib/core/layers/fastapi/requirements.txt

@@ -1,4 +1,4 @@
-boto3==1.34.131
+# boto3==1.34.131 // Provided by Lambda
 fastapi==0.111.0
 mangum==0.17.0
 pydantic==2.8.2

lib/models/model-api.ts

@@ -86,6 +86,7 @@ export class ModelsApi extends Construct {
             StringParameter.valueForStringParameter(this, `${config.deploymentPrefix}/layerVersion/fastapi`),
         );
 
+        const lambdaLayers = [commonLambdaLayer, fastapiLambdaLayer];
         const restApi = RestApi.fromRestApiAttributes(this, 'RestApi', {
             restApiId: restApiId,
             rootResourceId: rootResourceId,
@@ -140,7 +141,7 @@ export class ModelsApi extends Construct {
         const createModelStateMachine = new CreateModelStateMachine(this, 'CreateModelWorkflow', {
             config: config,
             modelTable: modelTable,
-            lambdaLayers: [commonLambdaLayer, fastapiLambdaLayer],
+            lambdaLayers: lambdaLayers,
             role: stateMachinesLambdaRole,
             vpc: vpc,
             securityGroups: securityGroups,
@@ -155,7 +156,7 @@ export class ModelsApi extends Construct {
         const deleteModelStateMachine = new DeleteModelStateMachine(this, 'DeleteModelWorkflow', {
             config: config,
             modelTable: modelTable,
-            lambdaLayers: [commonLambdaLayer, fastapiLambdaLayer],
+            lambdaLayers: lambdaLayers,
             role: stateMachinesLambdaRole,
             vpc: vpc,
             securityGroups: securityGroups,
@@ -167,7 +168,7 @@ export class ModelsApi extends Construct {
         const updateModelStateMachine = new UpdateModelStateMachine(this, 'UpdateModelWorkflow', {
             config: config,
             modelTable: modelTable,
-            lambdaLayers: [commonLambdaLayer, fastapiLambdaLayer],
+            lambdaLayers: lambdaLayers,
             role: stateMachinesLambdaRole,
             vpc: vpc,
             securityGroups: securityGroups,
@@ -193,7 +194,7 @@ export class ModelsApi extends Construct {
             restApi,
             authorizer,
             './lambda',
-            [commonLambdaLayer, fastapiLambdaLayer],
+            lambdaLayers,
             {
                 name: 'handler',
                 resource: 'models',
@@ -272,7 +273,7 @@ export class ModelsApi extends Construct {
                 restApi,
                 authorizer,
                 './lambda',
-                [commonLambdaLayer],
+                lambdaLayers,
                 f,
                 getDefaultRuntime(),
                 vpc,

lib/rag/layer/requirements.txt

@@ -1,5 +1,6 @@
-boto3>=1.34.131
-botocore>=1.34.131
+# boto3>=1.34.131 // Provided by Lambda
+# botocore>=1.34.131 // Provided by Lambda
+# urllib3<2 // Provided by Lambda
 langchain==0.3.9
 langchain-community==0.3.9
 langchain-openai==0.2.11
@@ -9,4 +10,3 @@ psycopg2-binary==2.9.9
 pypdf==4.3.1
 python-docx==1.1.0
 requests-aws4auth==1.2.3
-urllib3<2

lib/serve/rest-api/src/api/endpoints/v2/litellm_passthrough.py

@@ -39,6 +39,9 @@
     # List models
     "models",
     "v1/models",
+    # Model Info
+    "model/info",
+    "v1/model/info",
     # Text completions
     "chat/completions",
     "v1/chat/completions",

lisa-sdk/lisapy/__init__.py

@@ -12,5 +12,5 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-# flake8: noqa
-from .main import Lisa
+from .api import LisaApi  # noqa: F401
+from .main import LisaLlm  # noqa: F401

lisa-sdk/lisapy/api.py

@@ -0,0 +1,45 @@
+#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License").
+#   You may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from typing import Any, Dict, Optional, Union
+
+from pydantic import BaseModel, Field
+from requests import Session
+
+from .config import ConfigMixin
+from .doc import DocsMixin
+from .model import ModelMixin
+from .rag import RagMixin
+from .repository import RepositoryMixin
+from .session import SessionMixin
+
+
+class LisaApi(BaseModel, RepositoryMixin, ModelMixin, ConfigMixin, DocsMixin, RagMixin, SessionMixin):
+    url: str = Field(..., description="REST API url for LiteLLM")
+    headers: Optional[Dict[str, str]] = Field(None, description="Headers for request.")
+    cookies: Optional[Dict[str, str]] = Field(None, description="Cookies for request.")
+    verify: Optional[Union[str, bool]] = Field(None, description="Whether to verify SSL certificates.")
+    timeout: int = Field(10, description="Timeout in minutes request.")
+    _session: Session
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self._session = Session()
+        if self.headers:
+            self._session.headers = self.headers  # type: ignore
+        if self.verify is not None:
+            self._session.verify = self.verify
+        if self.cookies:
+            self._session.cookies = self.cookies  # type: ignore

lisa-sdk/lisapy/common.py

@@ -0,0 +1,26 @@
+#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License").
+#   You may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from typing import Dict, Optional, Union
+
+from requests import Session
+
+
+class BaseMixin:
+    url: str
+    headers: Optional[Dict[str, str]]
+    cookies: Optional[Dict[str, str]]
+    verify: Optional[Union[str, bool]]
+    timeout: int
+    _session: Session

lisa-sdk/lisapy/config.py

@@ -0,0 +1,30 @@
+#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License").
+#   You may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from typing import Dict, List
+
+from .common import BaseMixin
+from .errors import parse_error
+
+
+class ConfigMixin(BaseMixin):
+    """Mixin for config-related operations."""
+
+    def get_configs(self, config_scope: str = "global") -> List[Dict]:
+        response = self._session.get(f"{self.url}/configuration?configScope={config_scope}")
+        if response.status_code == 200:
+            json_configs: List[Dict] = response.json()
+            return json_configs
+        else:
+            raise parse_error(response.status_code, response)

lisa-sdk/lisapy/doc.py

@@ -0,0 +1,28 @@
+#   Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#   Licensed under the Apache License, Version 2.0 (the "License").
+#   You may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from .common import BaseMixin
+from .errors import parse_error
+
+
+class DocsMixin(BaseMixin):
+    """Mixin for doc-related operations."""
+
+    def list_docs(self) -> str:
+        response = self._session.get(f"{self.url}/docs")
+        if response.status_code == 200:
+            html: str = response.text
+            return html
+        else:
+            raise parse_error(response.status_code, response)