Domain identification

This notebook illustrates how to identify AI domains based on specific use cases.¶

Import libraries¶

In [2]:

from ai_atlas_nexus.blocks.inference import (
    RITSInferenceEngine,
    WMLInferenceEngine,
    OllamaInferenceEngine,
    VLLMInferenceEngine,
)
from ai_atlas_nexus.blocks.inference.params import (
    InferenceEngineCredentials,
    RITSInferenceEngineParams,
    WMLInferenceEngineParams,
    OllamaInferenceEngineParams,
    VLLMInferenceEngineParams,
)
from ai_atlas_nexus.library import AIAtlasNexus
import os

from ai_atlas_nexus.blocks.inference import ( RITSInferenceEngine, WMLInferenceEngine, OllamaInferenceEngine, VLLMInferenceEngine, ) from ai_atlas_nexus.blocks.inference.params import ( InferenceEngineCredentials, RITSInferenceEngineParams, WMLInferenceEngineParams, OllamaInferenceEngineParams, VLLMInferenceEngineParams, ) from ai_atlas_nexus.library import AIAtlasNexus import os

/Users/dhaval/.pyenv/versions/3.14.3/envs/aan/lib/python3.14/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

AI Atlas Nexus uses Large Language Models (LLMs) to infer risks dimensions. Therefore requires access to LLMs to inference or call the model.¶

Available Inference Engines: WML, Ollama, vLLM, RITS. Please follow the Inference APIs guide before going ahead.

Note: RITS is intended solely for internal IBM use and requires TUNNELALL VPN for access.

In [ ]:

inference_engine = OllamaInferenceEngine(
    model_name_or_path="granite3.3:8b",
    credentials=InferenceEngineCredentials(api_url="http://localhost:11434"),
    parameters=OllamaInferenceEngineParams(
        num_predict=1000, num_ctx=8192, temperature=0
    ),
)

# inference_engine = WMLInferenceEngine(
#     model_name_or_path="ibm/granite-4-h-small",
#     credentials={
#         "api_key": os.getenv("WML_API_KEY"),
#         "api_url": os.getenv("WML_API_URL"),
#         "project_id": os.getenv("WML_PROJECT_ID"),
#     },
#     parameters=WMLInferenceEngineParams(
#         max_completion_tokens=1024, temperature=0, seed=99
#     ),
# )

# inference_engine = VLLMInferenceEngine(
#     model_name_or_path="ibm-granite/granite-3.3-8b-instruct",
#     credentials=InferenceEngineCredentials(
#         api_url=os.getenv("VLLM_API_URL"), api_key=os.getenv("VLLM_API_KEY")
#     ),
#     parameters=VLLMInferenceEngineParams(max_tokens=1000, temperature=0),
# )

# inference_engine = RITSInferenceEngine(
#     model_name_or_path="ibm-granite/granite-3.3-8b-instruct",
#     credentials={
#         "api_key": os.getenv("RITS_API_KEY"),
#         "api_url": os.getenv("RITS_API_URL"),
#     },
#     parameters=RITSInferenceEngineParams(max_completion_tokens=1000, temperature=0),
# )

inference_engine = OllamaInferenceEngine( model_name_or_path="granite3.3:8b", credentials=InferenceEngineCredentials(api_url="http://localhost:11434"), parameters=OllamaInferenceEngineParams( num_predict=1000, num_ctx=8192, temperature=0 ), ) # inference_engine = WMLInferenceEngine( # model_name_or_path="ibm/granite-4-h-small", # credentials={ # "api_key": os.getenv("WML_API_KEY"), # "api_url": os.getenv("WML_API_URL"), # "project_id": os.getenv("WML_PROJECT_ID"), # }, # parameters=WMLInferenceEngineParams( # max_completion_tokens=1024, temperature=0, seed=99 # ), # ) # inference_engine = VLLMInferenceEngine( # model_name_or_path="ibm-granite/granite-3.3-8b-instruct", # credentials=InferenceEngineCredentials( # api_url=os.getenv("VLLM_API_URL"), api_key=os.getenv("VLLM_API_KEY") # ), # parameters=VLLMInferenceEngineParams(max_tokens=1000, temperature=0), # ) # inference_engine = RITSInferenceEngine( # model_name_or_path="ibm-granite/granite-3.3-8b-instruct", # credentials={ # "api_key": os.getenv("RITS_API_KEY"), # "api_url": os.getenv("RITS_API_URL"), # }, # parameters=RITSInferenceEngineParams(max_completion_tokens=1000, temperature=0), # )

Model 'ibm/granite-20b-code-instruct' is not supported for this environment. Supported models: ['ibm/granite-3-1-8b-base', 'ibm/granite-4-h-small', 'ibm/granite-embedding-278m-multilingual', 'ibm/granite-ttm-1024-96-r2', 'ibm/granite-ttm-1536-96-r2', 'ibm/granite-ttm-512-96-r2', 'ibm/slate-125m-english-rtrvr-v2', 'ibm/slate-30m-english-rtrvr-v2', 'intfloat/multilingual-e5-large', 'meta-llama/llama-3-1-70b-gptq', 'meta-llama/llama-3-1-8b', 'meta-llama/llama-3-3-70b-instruct', 'meta-llama/llama-4-maverick-17b-128e-instruct-fp8', 'mistralai/mistral-small-3-1-24b-instruct-2503', 'openai/gpt-oss-120b']

---------------------------------------------------------------------------
WMLClientError                            Traceback (most recent call last)
Cell In[4], line 9
      5 #         num_predict=1000, num_ctx=8192, temperature=0
      6 #     ),
      7 # )
      8 
----> 9 inference_engine = WMLInferenceEngine(
     10     model_name_or_path="ibm/granite-20b-code-instruct",
     11     credentials={
     12         "api_key": os.getenv("WML_API_KEY"),

File ~/Projects/AI-Governance/ai-atlas-nexus/src/ai_atlas_nexus/blocks/inference/base.py:78, in InferenceEngine.__init__(self, model_name_or_path, credentials, parameters, backend, concurrency_limit, auto_download_model)
     75 self.auto_download_model = auto_download_model
     77 # Create inference client
---> 78 self.client = self.create_client()
     80 # Health check
     81 try:

File ~/Projects/AI-Governance/ai-atlas-nexus/src/ai_atlas_nexus/blocks/inference/wml.py:101, in WMLInferenceEngine.create_client(self)
     98 else:
     99     client.set.default_project(self.credentials["project_id"])
--> 101 return ModelInference(
    102     model_id=self.model_name_or_path, api_client=client, params=self.parameters
    103 )

File ~/.pyenv/versions/3.14.3/envs/aan/lib/python3.14/site-packages/ibm_watsonx_ai/foundation_models/inference/model_inference.py:232, in ModelInference.__init__(self, model_id, deployment_id, params, credentials, project_id, space_id, verify, api_client, validate, max_retries, delay_time, retry_status_codes, **kwargs)
    230 self._inference: BaseModelInference
    231 if self._model_id:
--> 232     self._inference = FMModelInference(
    233         model_id=self._model_id,
    234         api_client=self._client,
    235         params=self.params,
    236         validate=validate,
    237         max_retries=max_retries,
    238         delay_time=delay_time,
    239         retry_status_codes=retry_status_codes,
    240     )
    241 else:
    242     self._deployment_id = cast(str, self._deployment_id)

File ~/.pyenv/versions/3.14.3/envs/aan/lib/python3.14/site-packages/ibm_watsonx_ai/foundation_models/inference/fm_model_inference.py:89, in FMModelInference.__init__(self, model_id, api_client, params, validate, max_retries, delay_time, retry_status_codes)
     86             break
     88     if not self._tech_preview:
---> 89         raise WMLClientError(
     90             error_msg=f"Model '{self.model_id}' is not supported for this environment. "
     91             f"Supported models: {supported_models}"
     92         )
     94 # check if model is in constricted mode
     95 _check_model_state(
     96     self._client,
     97     self.model_id,
     98     tech_preview=self._tech_preview,
     99     model_specs=model_specs,
    100 )

WMLClientError: Model 'ibm/granite-20b-code-instruct' is not supported for this environment. Supported models: ['ibm/granite-3-1-8b-base', 'ibm/granite-4-h-small', 'ibm/granite-embedding-278m-multilingual', 'ibm/granite-ttm-1024-96-r2', 'ibm/granite-ttm-1536-96-r2', 'ibm/granite-ttm-512-96-r2', 'ibm/slate-125m-english-rtrvr-v2', 'ibm/slate-30m-english-rtrvr-v2', 'intfloat/multilingual-e5-large', 'meta-llama/llama-3-1-70b-gptq', 'meta-llama/llama-3-1-8b', 'meta-llama/llama-3-3-70b-instruct', 'meta-llama/llama-4-maverick-17b-128e-instruct-fp8', 'mistralai/mistral-small-3-1-24b-instruct-2503', 'openai/gpt-oss-120b']

Create an instance of AIAtlasNexus¶

Note: (Optional) You can specify your own directory in AIAtlasNexus(base_dir=<PATH>) to utilize custom AI ontologies. If left blank, the system will use the provided AI ontologies.

In [2]:

ai_atlas_nexus = AIAtlasNexus()

ai_atlas_nexus = AIAtlasNexus()

[2026-08-09 21:36:43:643] - INFO - AIAtlasNexus - Created AIAtlasNexus instance. Base_dir: None

AI Domain Identification API - Default backend¶

• inference directly using the provided inference engine

AIAtlasNexus.identify_domain_from_usecases()

Params:

• usecases (List[str]): A List of strings describing AI usecases
• inference_engine (InferenceEngine): An LLM inference engine to identify AI tasks from usecases.
• verbose (bool, optional): prints detailed output during the inference process. Defaults to True.

In [14]:

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers."

domains = ai_atlas_nexus.identify_domain_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

domains[0].prediction

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers." domains = ai_atlas_nexus.identify_domain_from_usecases( usecases=[usecase], inference_engine=inference_engine, ) domains[0].prediction

Inferring with OLLAMA, backend - MELLEA:   0%|          | 0/1 [00:00<?, ?it/s]

=== 22:15:16-INFO ======
SUCCESS

  0%|          | 0/3 [00:06<?, ?it/s]
Inferring with OLLAMA, backend - MELLEA: 100%|██████████| 1/1 [00:06<00:00,  6.20s/it]

Out[14]:

{'answer': 'Customer service/support',
 'explanation': "The use case involves generating personalized, relevant responses, recommendations, and summaries of claims for customers to support agents. This directly aligns with the definition of 'Customer Service/Support' AI agents that handle customer inquiries, resolve issues, provide product information, and manage support tickets across channels like chat, email, and phone."}

AI Domain Identification API - Mellea backend using Ollama¶

• Inference is performed using the Mellea backend, which utilizes the specified inference engine.
• Mellea backend currently only supports Ollama, WML and RITS inference engines.

In [15]:

inference_engine = OllamaInferenceEngine(
    model_name_or_path="granite3.3:8b",
    credentials=InferenceEngineCredentials(api_url="http://localhost:11434"),
    parameters=OllamaInferenceEngineParams(
        num_predict=1000, num_ctx=8192, temperature=0
    ),
    backend="mellea",
)

inference_engine = OllamaInferenceEngine( model_name_or_path="granite3.3:8b", credentials=InferenceEngineCredentials(api_url="http://localhost:11434"), parameters=OllamaInferenceEngineParams( num_predict=1000, num_ctx=8192, temperature=0 ), backend="mellea", )

=== 22:16:03-INFO ======
Starting Mellea session: backend=OLLAMA, model=granite3.3:8b, context=SimpleContext

[2026-03-18 22:16:03:874] - INFO - AIAtlasNexus - ✓ Created OLLAMA inference engine for model: granite3.3:8b, backend - MELLEA

In [7]:

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers."

domains = ai_atlas_nexus.identify_domain_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

domains[0].prediction

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers." domains = ai_atlas_nexus.identify_domain_from_usecases( usecases=[usecase], inference_engine=inference_engine, ) domains[0].prediction

Inferring with OLLAMA, backend - MELLEA:   0%|          | 0/1 [00:00<?, ?it/s]

=== 22:10:10-INFO ======
SUCCESS

  0%|          | 0/3 [00:08<?, ?it/s]
Inferring with OLLAMA, backend - MELLEA: 100%|██████████| 1/1 [00:08<00:00,  8.50s/it]

Out[7]:

{'answer': 'Customer service/support',
 'explanation': "The use case involves generating personalized, relevant responses, recommendations, and summaries of claims for customers to support agents. This directly aligns with the definition of 'Customer Service/Support' AI agents that handle customer inquiries, resolve issues, provide product information, and manage support tickets across channels like chat, email, and phone."}

AI Domain Identification API - Mellea backend using WML¶

• Inference is performed using the Mellea backend, which utilizes the specified inference engine.
• Mellea backend currently supports only the Ollama and WML backends.

In [3]:

inference_engine = WMLInferenceEngine(
    model_name_or_path="ibm/granite-4-h-small",
    credentials={
        "api_key": os.getenv("WML_API_KEY"),
        "api_url": os.getenv("WML_API_URL"),
        "project_id": os.getenv("WML_PROJECT_ID"),
    },
    parameters=WMLInferenceEngineParams(
        max_completion_tokens=1024, temperature=0, seed=99, repetition_penalty=1
    ),
    backend="mellea",
)

inference_engine = WMLInferenceEngine( model_name_or_path="ibm/granite-4-h-small", credentials={ "api_key": os.getenv("WML_API_KEY"), "api_url": os.getenv("WML_API_URL"), "project_id": os.getenv("WML_PROJECT_ID"), }, parameters=WMLInferenceEngineParams( max_completion_tokens=1024, temperature=0, seed=99, repetition_penalty=1 ), backend="mellea", )

=== 21:36:54-INFO ======
Starting Mellea session: backend=WML, model=ibm/granite-4-h-small, context=SimpleContext

[2026-08-09 21:36:54:477] - INFO - AIAtlasNexus - ✓ Created WML inference engine for model: ibm/granite-4-h-small, backend - MELLEA

In [4]:

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers."

domains = ai_atlas_nexus.identify_domain_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

domains[0].prediction

usecase = "Generate personalized, relevant responses, recommendations, and summaries of claims for customers to support agents to enhance their interactions with customers." domains = ai_atlas_nexus.identify_domain_from_usecases( usecases=[usecase], inference_engine=inference_engine, ) domains[0].prediction

Inferring with WML, backend - MELLEA:   0%|          | 0/1 [00:00<?, ?it/s]

=== 21:37:02-INFO ======
SUCCESS

  0%|          | 0/3 [00:06<?, ?it/s]
Inferring with WML, backend - MELLEA: 100%|██████████| 1/1 [00:06<00:00,  6.60s/it]

Out[4]:

{'answer': 'Customer service/support',
 'explanation': "The use case involves generating personalized, relevant responses, recommendations, and summaries of claims to assist support agents in their interactions with customers. This directly aligns with the definition of the 'Customer Service/Support' AI domain, which focuses on AI agents handling customer inquiries, resolving issues, providing product information, and managing support tickets across various communication channels. The described functionality is a core component of customer service and support operations."}