Vectore store improvements

lambda/authorizer/lambda_functions.py

@@ -13,11 +13,12 @@
 #   limitations under the License.
 
 """Authorize for REST API."""
+import json
 import logging
 import os
 import ssl
 from functools import cache
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import boto3
 import create_env_variables  # noqa: F401
@@ -61,8 +62,9 @@ def lambda_handler(event: Dict[str, Any], context) -> Dict[str, Any]:  # type: i
 
     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)
         allow_policy = generate_policy(effect="Allow", resource=event["methodArn"], username=jwt_data["sub"])
-        allow_policy["context"] = {"username": jwt_data["sub"]}
+        allow_policy["context"] = {"username": jwt_data["sub"], "groups": json.dumps(groups or [])}
 
         if requested_resource.startswith("/models") and not is_admin_user:
             # non-admin users can still list models
@@ -142,14 +144,20 @@ def id_token_is_valid(*, id_token: str, client_id: str, authority: str) -> Dict[
 
 def is_admin(jwt_data: dict[str, Any], admin_group: str, jwt_groups_property: str) -> bool:
     """Check if the user is an admin."""
-    props = jwt_groups_property.split(".")
-    current_node = jwt_data
+    return admin_group in (get_property_path(jwt_data, jwt_groups_property) or [])
+
+
+def get_property_path(data: dict[str, Any], property_path: str) -> Optional[Any]:
+    """Get the value represented by a property path."""
+    props = property_path.split(".")
+    current_node = data
     for prop in props:
         if prop in current_node:
             current_node = current_node[prop]
         else:
-            return False
-    return admin_group in current_node
+            return None
+
+    return current_node
 
 
 @cache

lambda/repository/lambda_functions.py

@@ -23,9 +23,10 @@
 from botocore.config import Config
 from lisapy.langchain import LisaOpenAIEmbeddings
 from utilities.common_functions import api_wrapper, get_cert_path, get_id_token, retry_config
+from utilities.exceptions import HTTPException
 from utilities.file_processing import process_record
 from utilities.validation import validate_model_name, ValidationError
-from utilities.vector_store import get_vector_store_client
+from utilities.vector_store import find_repository_by_id, get_registered_repositories, get_vector_store_client
 
 logger = logging.getLogger(__name__)
 session = boto3.Session()
@@ -44,7 +45,6 @@
     ),
 )
 lisa_api_endpoint = ""
-registered_repositories: List[Dict[str, Any]] = []
 
 
 def _get_embeddings(model_name: str, id_token: str) -> LisaOpenAIEmbeddings:
@@ -186,13 +186,25 @@ def list_all(event: dict, context: dict) -> List[Dict[str, Any]]:
     Currently there is not support for dynamic repositories so only a single OpenSearch repository
     is returned.
     """
-    global registered_repositories
 
-    if not registered_repositories:
-        registered_repositories_response = ssm_client.get_parameter(Name=os.environ["REGISTERED_REPOSITORIES_PS_NAME"])
-        registered_repositories = json.loads(registered_repositories_response["Parameter"]["Value"])
+    user_groups = json.loads(event["requestContext"]["authorizer"]["groups"]) or []
+    registered_repositories = get_registered_repositories()
+
+    return list(
+        filter(lambda repository: user_has_group(user_groups, repository["allowedGroups"]), registered_repositories)
+    )
 
-    return registered_repositories
+
+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.
+
+    If allowed groups is empty this will return True.
+    """
+
+    if len(allowed_groups) > 0:
+        return len(set(user_groups).intersection(set(allowed_groups))) > 0
+    else:
+        return True
 
 
 @api_wrapper
@@ -207,13 +219,16 @@ def similarity_search(event: dict, context: dict) -> Dict[str, Any]:
     query_string_params = event["queryStringParameters"]
     model_name = query_string_params["modelName"]
     query = query_string_params["query"]
-    repository_type = query_string_params["repositoryType"]
     top_k = query_string_params.get("topK", 3)
+    repository_id = event["pathParameters"]["repositoryId"]
+
+    repository = find_repository_by_id(repository_id)
+    ensure_repository_access(event, repository)
 
     id_token = get_id_token(event)
 
     embeddings = _get_embeddings(model_name=model_name, id_token=id_token)
-    vs = get_vector_store_client(repository_type, index=model_name, embeddings=embeddings)
+    vs = get_vector_store_client(repository_id, index=model_name, embeddings=embeddings)
     docs = vs.similarity_search(
         query,
         k=top_k,
@@ -225,6 +240,13 @@ def similarity_search(event: dict, context: dict) -> Dict[str, Any]:
     return doc_return
 
 
+def ensure_repository_access(event, repository):
+    "Ensures a user has access to the repository or else raises an HTTPException"
+    user_groups = json.loads(event["requestContext"]["authorizer"]["groups"]) or []
+    if not user_has_group(user_groups, repository["allowedGroups"]):
+        raise HTTPException(status_code=403, message="User does not have permission to access this repository")
+
+
 @api_wrapper
 def purge_document(event: dict, context: dict) -> Dict[str, Any]:
     """Purge all records related to the specified document from the RAG repository."""
@@ -248,9 +270,13 @@ def ingest_documents(event: dict, context: dict) -> dict:
     model_name = embedding_model["modelName"]
 
     query_string_params = event["queryStringParameters"]
-    repository_type = query_string_params["repositoryType"]
     chunk_size = int(query_string_params["chunkSize"]) if "chunkSize" in query_string_params else None
     chunk_overlap = int(query_string_params["chunkOverlap"]) if "chunkOverlap" in query_string_params else None
+    repository_id = event["pathParameters"]["repositoryId"]
+    logger.info(f"using repository {repository_id}")
+
+    repository = find_repository_by_id(repository_id)
+    ensure_repository_access(event, repository)
 
     docs = process_record(s3_keys=body["keys"], chunk_size=chunk_size, chunk_overlap=chunk_overlap)
 
@@ -264,7 +290,7 @@ def ingest_documents(event: dict, context: dict) -> dict:
 
     id_token = get_id_token(event)
     embeddings = _get_embeddings(model_name=model_name, id_token=id_token)
-    vs = get_vector_store_client(repository_type, index=model_name, embeddings=embeddings)
+    vs = get_vector_store_client(repository_id, index=model_name, embeddings=embeddings)
     ids = vs.add_texts(texts=texts, metadatas=metadatas)
     return {"ids": ids, "count": len(ids)}
 
@@ -294,3 +320,8 @@ def presigned_url(event: dict, context: dict) -> dict:
         ExpiresIn=3600,
     )
     return {"response": response}
+
+
+def get_groups(event: Any) -> List[str]:
+    groups: List[str] = json.loads(event["requestContext"]["authorizer"]["groups"])
+    return groups

lambda/repository/pipeline_ingest_documents.py

@@ -122,7 +122,7 @@ def handle_pipeline_ingest_documents(event: Dict[str, Any], context: Any) -> Dic
 
         # Initialize vector store using model name as index, matching lambda_functions.py pattern
         vs = get_vector_store_client(
-            store=repository_type,  # Changed from repository_type to store
+            repository_id,
             index=embedding_model,  # Use model name as index to match lambda_functions.py
             embeddings=embeddings,
         )

lambda/utilities/exceptions.py

@@ -17,3 +17,10 @@
 
 class RagUploadException(Exception):
     """RAG upload error exception."""
+
+
+class HTTPException(Exception):
+    def __init__(self, status_code: int = 400, message: str = "Bad Request") -> None:
+        self.http_status_code = status_code
+        self.message = message
+        super().__init__(self.message)

lambda/utilities/vector_store.py

@@ -16,6 +16,8 @@
 import json
 import logging
 import os
+from optparse import Option
+from typing import Any, Dict, List
 
 import boto3
 import create_env_variables  # noqa: F401
@@ -32,14 +34,40 @@
 session = boto3.Session()
 ssm_client = boto3.client("ssm", region_name=os.environ["AWS_REGION"], config=retry_config)
 secretsmanager_client = boto3.client("secretsmanager", region_name=os.environ["AWS_REGION"], config=retry_config)
+registered_repositories: List[Dict[str, Any]] = []
 
 
-def get_vector_store_client(store: str, index: str, embeddings: Embeddings) -> VectorStore:
+def get_registered_repositories() -> List[Any]:
+    """Get a list of all registered RAG repositories."""
+    global registered_repositories
+    if not registered_repositories:
+        registered_repositories_response = ssm_client.get_parameter(Name=os.environ["REGISTERED_REPOSITORIES_PS_NAME"])
+        registered_repositories = json.loads(registered_repositories_response["Parameter"]["Value"])
+
+    return registered_repositories
+
+
+def find_repository_by_id(repository_id: str) -> Option[Dict]:
+    """Find a RAG repository by id."""
+    return next(
+        (repository for repository in get_registered_repositories() if repository["repositoryId"] == repository_id),
+        None,
+    )
+
+
+def get_vector_store_client(repository_id: str, index: str, embeddings: Embeddings) -> VectorStore:
     """Return Langchain VectorStore corresponding to the specified store.
 
     Creates a langchain vector store based on the specified embeddigs adapter and backing store.
     """
-    if store == "opensearch":
+
+    repository = find_repository_by_id(repository_id)
+    repository_type = repository.get("repository_type", None)
+
+    prefix = os.environ["REGISTERED_REPOSITORIES_PS_PREFIX"]
+    connection_info = ssm_client.get_parameter(Name=f"{prefix}{repository_id}")
+
+    if repository_type == "opensearch":
         service = "es"
         session = boto3.Session()
         credentials = session.get_credentials()
@@ -52,11 +80,7 @@ def get_vector_store_client(store: str, index: str, embeddings: Embeddings) -> V
             session_token=credentials.token,
         )
 
-        global opensearch_endpoint
-
-        if not opensearch_endpoint:
-            opensearch_param_response = ssm_client.get_parameter(Name=os.environ["OPENSEARCH_ENDPOINT_PS_NAME"])
-            opensearch_endpoint = f'https://{opensearch_param_response["Parameter"]["Value"]}'
+        opensearch_endpoint = f'https://{connection_info["Parameter"]["Value"]}'
 
         return OpenSearchVectorSearch(
             opensearch_url=opensearch_endpoint,
@@ -69,11 +93,8 @@ def get_vector_store_client(store: str, index: str, embeddings: Embeddings) -> V
             connection_class=RequestsHttpConnection,
         )
 
-    elif store == "pgvector":
-        connection_info = json.loads(
-            ssm_client.get_parameter(Name=os.environ["RDS_CONNECTION_INFO_PS_NAME"])["Parameter"]["Value"]
-        )
-
+    elif repository_type == "pgvector":
+        connection_info = json.loads(connection_info["Parameter"]["Value"])
         secrets_response = secretsmanager_client.get_secret_value(SecretId=connection_info["passwordSecretId"])
         password = json.loads(secrets_response["SecretString"])["password"]
 
@@ -91,4 +112,4 @@ def get_vector_store_client(store: str, index: str, embeddings: Embeddings) -> V
             embedding_function=embeddings,
         )
 
-    raise ValueError(f"Unrecognized RAG store: '{store}'")
+    raise ValueError(f"Unrecognized RAG store: '{repository_id}'")

lib/networking/vpc/security-group-factory.ts

@@ -92,13 +92,14 @@ export class SecurityGroupFactory {
         securityGroup: ISecurityGroup,
         securityGroupName: string,
         vpc: Vpc,
-        config: Config): void {
+        config: Config,
+        port: number): void {
         const subNets = config.subnets && config.vpcId ? vpc.subnetSelection?.subnets : vpc.vpc.isolatedSubnets.concat(vpc.vpc.privateSubnets);
         subNets?.forEach((subnet) => {
             securityGroup.connections.allowFrom(
                 Peer.ipv4(config.subnets ? config.subnets.filter((filteredSubnet: { subnetId: string; }) =>
                     filteredSubnet.subnetId === subnet.subnetId)?.[0]?.ipv4CidrBlock :  subnet.ipv4CidrBlock),
-                Port.tcp(config.restApiConfig.rdsConfig.dbPort),
+                Port.tcp(port),
                 `Allow REST API private subnets to communicate with ${securityGroupName}`,
             );
         });