feat(LAB-3088): add LLM models and project models configuration (#1774)

Co-authored-by: paulruelle <[email protected]>

Author: RuellePaul

src/kili/adapters/kili_api_gateway/kili_api_gateway.py

@@ -6,7 +6,7 @@
 from kili.adapters.kili_api_gateway.cloud_storage import CloudStorageOperationMixin
 from kili.adapters.kili_api_gateway.issue import IssueOperationMixin
 from kili.adapters.kili_api_gateway.label.operations_mixin import LabelOperationMixin
-from kili.adapters.kili_api_gateway.model_configuration.operations_mixin import (
+from kili.adapters.kili_api_gateway.llm.operations_mixin import (
     ModelConfigurationOperationMixin,
 )
 from kili.adapters.kili_api_gateway.notification.operations_mixin import (

src/kili/adapters/kili_api_gateway/model_configuration/__init__.py renamed to src/kili/adapters/kili_api_gateway/llm/__init__.py

@@ -0,0 +1,111 @@
+"""GraphQL payload data mappers for api keys operations."""
+
+from typing import Dict
+
+from kili.domain.llm import (
+    AzureOpenAICredentials,
+    ModelToCreateInput,
+    ModelToUpdateInput,
+    ModelType,
+    OpenAISDKCredentials,
+    OrganizationModelFilters,
+    ProjectModelFilters,
+    ProjectModelToCreateInput,
+    ProjectModelToUpdateInput,
+)
+
+
+def model_where_wrapper(filter: OrganizationModelFilters) -> Dict:
+    """Build the GraphQL ProjectMapperWhere variable to be sent in an operation."""
+    return {
+        "organizationId": filter.organization_id,
+    }
+
+
+def project_model_where_mapper(filter: ProjectModelFilters) -> Dict:
+    """Build the GraphQL ProjectMapperWhere variable to be sent in an operation."""
+    return {
+        "projectId": filter.project_id,
+        "modelId": filter.model_id,
+    }
+
+
+def map_create_model_input(data: ModelToCreateInput) -> Dict:
+    """Build the GraphQL ModelInput variable to be sent in an operation."""
+    if data.type == ModelType.AZURE_OPEN_AI and isinstance(
+        data.credentials, AzureOpenAICredentials
+    ):
+        credentials = {
+            "apiKey": data.credentials.api_key,
+            "deploymentId": data.credentials.deployment_id,
+            "endpoint": data.credentials.endpoint,
+        }
+    elif data.type == ModelType.OPEN_AI_SDK and isinstance(data.credentials, OpenAISDKCredentials):
+        credentials = {"apiKey": data.credentials.api_key, "endpoint": data.credentials.endpoint}
+    else:
+        raise ValueError(
+            f"Unsupported model type or credentials: {data.type}, {type(data.credentials)}"
+        )
+
+    return {
+        "credentials": credentials,
+        "name": data.name,
+        "type": data.type.value,
+        "organizationId": data.organization_id,
+    }
+
+
+def map_update_model_input(data: ModelToUpdateInput) -> Dict:
+    """Build the GraphQL UpdateModelInput variable to be sent in an operation."""
+    input_dict = {}
+    if data.name is not None:
+        input_dict["name"] = data.name
+
+    if data.credentials is not None:
+        if isinstance(data.credentials, AzureOpenAICredentials):
+            credentials = {
+                "apiKey": data.credentials.api_key,
+                "deploymentId": data.credentials.deployment_id,
+                "endpoint": data.credentials.endpoint,
+            }
+        elif isinstance(data.credentials, OpenAISDKCredentials):
+            credentials = {
+                "apiKey": data.credentials.api_key,
+                "endpoint": data.credentials.endpoint,
+            }
+        else:
+            raise ValueError(f"Unsupported credentials type: {type(data.credentials)}")
+        input_dict["credentials"] = credentials
+
+    return input_dict
+
+
+def map_create_project_model_input(data: ProjectModelToCreateInput) -> Dict:
+    """Build the GraphQL ModelInput variable to be sent in an operation."""
+    return {
+        "projectId": data.project_id,
+        "modelId": data.model_id,
+        "configuration": data.configuration,
+    }
+
+
+def map_update_project_model_input(data: ProjectModelToUpdateInput) -> Dict:
+    """Build the GraphQL UpdateProjectModelInput variable to be sent in an operation."""
+    input_dict = {}
+    if data.configuration is not None:
+        input_dict["configuration"] = data.configuration
+    return input_dict
+
+
+def map_delete_model_input(model_id: str) -> Dict:
+    """Map the input for the GraphQL deleteModel mutation."""
+    return {
+        "deleteModelId": model_id,
+    }
+
+
+def map_delete_project_model_input(project_model_id: str) -> Dict:
+    """Map the input for the GraphQL deleteProjectModel mutation."""
+    return {
+        "deleteProjectModelId": project_model_id,
+    }

src/kili/adapters/kili_api_gateway/llm/mappers.py

@@ -0,0 +1,96 @@
+"""GraphQL Asset operations."""
+
+
+def get_models_query(fragment: str) -> str:
+    """Return the GraphQL projectModels query."""
+    return f"""
+          query Models($where: ModelWhere!, $first: PageSize!, $skip: Int!) {{
+            data: models(where: $where, first: $first, skip: $skip) {{
+                {fragment}
+              }}
+            }}
+        """
+
+
+def get_model_query(fragment: str) -> str:
+    """Return the GraphQL model query by ID."""
+    return f"""
+        query Model($modelId: ID!) {{
+            model(id: $modelId) {{
+                {fragment}
+            }}
+        }}
+    """
+
+
+def get_create_model_mutation(fragment: str) -> str:
+    """Return the GraphQL createProjectModel mutation."""
+    return f"""
+          mutation CreateModel($input: CreateModelInput!) {{
+            createModel(input: $input) {{
+              {fragment}
+            }}
+          }}
+        """
+
+
+def get_update_model_mutation(fragment: str) -> str:
+    """Return the GraphQL updateModel mutation."""
+    return f"""
+        mutation UpdateModel($id: ID!, $input: UpdateModelInput!) {{
+            updateModel(id: $id, input: $input) {{
+                {fragment}
+            }}
+        }}
+    """
+
+
+def get_delete_model_mutation() -> str:
+    """Return the GraphQL deleteOrganizationModel mutation."""
+    return """
+    mutation DeleteModel($deleteModelId: ID!) {
+        deleteModel(id: $deleteModelId)
+    }
+    """
+
+
+def get_create_project_model_mutation(fragment: str) -> str:
+    """Return the GraphQL createProjectModel mutation."""
+    return f"""
+          mutation CreateProjectModel($input: CreateProjectModelInput!) {{
+            createProjectModel(input: $input) {{
+              {fragment}
+            }}
+          }}
+        """
+
+
+def get_update_project_model_mutation(fragment: str) -> str:
+    """Return the GraphQL updateProjectModel mutation."""
+    return f"""
+        mutation UpdateProjectModel($updateProjectModelId: ID!, $input: UpdateProjectModelInput!) {{
+            updateProjectModel(id: $updateProjectModelId, input: $input) {{
+                {fragment}
+            }}
+        }}
+    """
+
+
+def get_delete_project_model_mutation() -> str:
+    """Return the GraphQL deleteProjectModel mutation."""
+    return """
+    mutation DeleteProjectModel($deleteProjectModelId: ID!) {
+       deleteProjectModel(id: $deleteProjectModelId)
+    }
+    """
+
+
+def get_project_models_query(fragment: str) -> str:
+    """Return the GraphQL projectModels query."""
+    return f"""
+        query ProjectModels($where: ProjectModelWhere!, $first: PageSize!, $skip: Int!) {{
+          data: projectModels(where: $where, first: $first, skip: $skip) {{
+            {fragment}
+          }}
+        }}
+        """

src/kili/adapters/kili_api_gateway/llm/operations.py

@@ -0,0 +1,139 @@
+"""Mixin extending Kili API Gateway class with Api Keys related operations."""
+
+from typing import Dict, Generator, Optional
+
+from kili.adapters.kili_api_gateway.base import BaseOperationMixin
+from kili.adapters.kili_api_gateway.helpers.queries import (
+    PaginatedGraphQLQuery,
+    QueryOptions,
+    fragment_builder,
+)
+from kili.adapters.kili_api_gateway.llm.mappers import (
+    map_create_model_input,
+    map_create_project_model_input,
+    map_delete_model_input,
+    map_delete_project_model_input,
+    map_update_model_input,
+    map_update_project_model_input,
+    model_where_wrapper,
+    project_model_where_mapper,
+)
+from kili.adapters.kili_api_gateway.llm.operations import (
+    get_create_model_mutation,
+    get_create_project_model_mutation,
+    get_delete_model_mutation,
+    get_delete_project_model_mutation,
+    get_model_query,
+    get_models_query,
+    get_project_models_query,
+    get_update_model_mutation,
+    get_update_project_model_mutation,
+)
+from kili.domain.llm import (
+    ModelToCreateInput,
+    ModelToUpdateInput,
+    OrganizationModelFilters,
+    ProjectModelFilters,
+    ProjectModelToCreateInput,
+    ProjectModelToUpdateInput,
+)
+from kili.domain.types import ListOrTuple
+
+
+class ModelConfigurationOperationMixin(BaseOperationMixin):
+    """Mixin extending Kili API Gateway class with model configuration related operations."""
+
+    def list_models(
+        self,
+        filters: OrganizationModelFilters,
+        fields: ListOrTuple[str],
+        options: Optional[QueryOptions] = None,
+    ) -> Generator[Dict, None, None]:
+        """List models with given options."""
+        fragment = fragment_builder(fields)
+        query = get_models_query(fragment)
+        where = model_where_wrapper(filters)
+        return PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
+            query,
+            where,
+            options if options else QueryOptions(disable_tqdm=False),
+            "Retrieving organization models",
+            None,
+        )
+
+    def get_model(self, model_id: str, fields: ListOrTuple[str]) -> Dict:
+        """Get a model by ID."""
+        fragment = fragment_builder(fields)
+        query = get_model_query(fragment)
+        variables = {"modelId": model_id}
+        result = self.graphql_client.execute(query, variables)
+        return result["model"]
+
+    def create_model(self, model: ModelToCreateInput) -> Dict:
+        """Send a GraphQL request calling createModel resolver."""
+        payload = {"input": map_create_model_input(model)}
+        fragment = fragment_builder(["id"])
+        mutation = get_create_model_mutation(fragment)
+        result = self.graphql_client.execute(mutation, payload)
+        return result["createModel"]
+
+    def update_properties_in_model(self, model_id: str, model: ModelToUpdateInput) -> Dict:
+        """Send a GraphQL request calling updateModel resolver."""
+        payload = {"id": model_id, "input": map_update_model_input(model)}
+        fragment = fragment_builder(["id"])
+        mutation = get_update_model_mutation(fragment)
+        result = self.graphql_client.execute(mutation, payload)
+        return result["updateModel"]
+
+    def delete_model(self, model_id: str) -> Dict:
+        """Send a GraphQL request to delete an organization model."""
+        payload = map_delete_model_input(model_id)
+        mutation = get_delete_model_mutation()
+        result = self.graphql_client.execute(mutation, payload)
+        return result["deleteModel"]
+
+    def create_project_model(self, project_model: ProjectModelToCreateInput) -> Dict:
+        """Send a GraphQL request calling createModel resolver."""
+        payload = {"input": map_create_project_model_input(project_model)}
+        fragment = fragment_builder(["id"])
+        mutation = get_create_project_model_mutation(fragment)
+        result = self.graphql_client.execute(mutation, payload)
+        return result["createProjectModel"]
+
+    def update_project_model(
+        self, project_model_id: str, project_model: ProjectModelToUpdateInput
+    ) -> Dict:
+        """Send a GraphQL request calling updateProjectModel resolver."""
+        payload = {
+            "updateProjectModelId": project_model_id,
+            "input": map_update_project_model_input(project_model),
+        }
+        fragment = fragment_builder(["id", "configuration"])
+        mutation = get_update_project_model_mutation(fragment)
+        result = self.graphql_client.execute(mutation, payload)
+        return result["updateProjectModel"]
+
+    def delete_project_model(self, project_model_id: str) -> Dict:
+        """Send a GraphQL request to delete a project model."""
+        payload = map_delete_project_model_input(project_model_id)
+        mutation = get_delete_project_model_mutation()
+        result = self.graphql_client.execute(mutation, payload)
+        return result["deleteProjectModel"]
+
+    def list_project_models(
+        self,
+        filters: ProjectModelFilters,
+        fields: ListOrTuple[str],
+        options: Optional[QueryOptions] = None,
+    ) -> Generator[Dict, None, None]:
+        """List project models with given options."""
+        fragment = fragment_builder(fields)
+        query = get_project_models_query(fragment)
+        where = project_model_where_mapper(filters)
+        return PaginatedGraphQLQuery(self.graphql_client).execute_query_from_paginated_call(
+            query,
+            where,
+            options if options else QueryOptions(disable_tqdm=False),
+            "Retrieving project models",
+            None,
+        )