AI Tasks identification

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

Import libraries¶

In [1]:

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 [7]:

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), # )

[2026-08-09 21:40:22:680] - INFO - AIAtlasNexus - ✓ Created OLLAMA inference engine for model: granite3.3:8b, backend - DEFAULT

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 [3]:

ai_atlas_nexus = AIAtlasNexus()

ai_atlas_nexus = AIAtlasNexus()

[2026-08-09 21:38:34:430] - INFO - AIAtlasNexus - Created AIAtlasNexus instance. Base_dir: None

AI Tasks Identification API - Default backend using OLLAMA¶

AIAtlasNexus.identify_ai_tasks_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 [8]:

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

risks = ai_atlas_nexus.identify_ai_tasks_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

risks[0].prediction

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

Inferring with OLLAMA, backend - DEFAULT: 100%|██████████| 1/1 [00:14<00:00, 14.17s/it]

Out[8]:

{'ai_tasks': [{'ai_task': 'Text-to-Text',
   'explanation': 'This task involves generating text based on given input, which aligns with the requirement to create personalized and relevant responses for customers.'},
  {'ai_task': 'Summarization',
   'explanation': 'Summarizing claims or information to provide concise and relevant summaries for customer interactions falls under this task.'}]}

AI Tasks 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 [6]:

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:27:32-INFO ======
Starting Mellea session: backend=OLLAMA, model=granite3.3:8b, context=SimpleContext

[2026-03-18 22:27:32:506] - 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."

risks = ai_atlas_nexus.identify_ai_tasks_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

risks[0].prediction

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

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

=== 22:27:53-INFO ======
SUCCESS

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

Out[7]:

{'ai_tasks': [{'ai_task': 'Any-to-Any',
   'explanation': "The use case involves generating personalized, relevant responses and summaries of claims, which requires understanding various inputs (customer data, context, etc.) and producing diverse outputs (text, recommendations). This aligns with the 'Any-to-any' task description."},
  {'ai_task': 'Text Generation',
   'explanation': 'The use case specifically involves generating text for responses and summaries. Text generation models can produce new text based on given inputs, which fits this requirement.'}]}

AI Tasks 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 [5]:

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:39:22-INFO ======
Starting Mellea session: backend=WML, model=ibm/granite-4-h-small, context=SimpleContext

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

In [6]:

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

risks = ai_atlas_nexus.identify_ai_tasks_from_usecases(
    usecases=[usecase],
    inference_engine=inference_engine,
)

risks[0].prediction

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

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

=== 21:39:44-INFO ======
FAILED. Valid: 3/4. Failed: 
	 - Provide a brief, plausible explanation for your choice

=== 21:39:55-INFO ======
SUCCESS

 33%|███▎      | 1/3 [00:25<00:51, 25.74s/it]
Inferring with WML, backend - MELLEA: 100%|██████████| 1/1 [00:25<00:00, 25.75s/it]

Out[6]:

{'ai_tasks': [{'ai_task': 'Text Classification',
   'explanation': "The use case involves generating personalized, relevant responses, recommendations, and summaries of claims for customers. This requires understanding the customer's query or claim, which can be achieved through text classification. By classifying the customer's input into predefined categories, the AI can determine the appropriate response or action to take."},
  {'ai_task': 'Text Generation',
   'explanation': "Generating personalized responses and summaries of claims requires the AI to create new text based on the customer's input. Text generation models can take the customer's query or claim as input and generate relevant, coherent, and contextually appropriate responses or summaries."},
  {'ai_task': 'Question Answering',
   'explanation': 'The AI may need to retrieve specific information from a knowledge base or document to provide accurate responses to customer queries or claims. Question answering models can take a question as input and retrieve the most relevant answer from a given text or document, ensuring the AI provides accurate and helpful information to the customer.'}]}