AI Tasks identification

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

Import libraries¶

In [ ]:

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

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_new_tokens=1000, decoding_method="greedy"
#     ),
# )

# 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_new_tokens=1000, decoding_method="greedy" # ), # ) # 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-03-18 22:30:16:754] - INFO - AIAtlasNexus - ✓ Created RITS inference engine for model: ibm-granite/granite-3.3-8b-instruct, 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-03-18 22:26:13:220] - 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 [4]:

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:10<00:00, 10.56s/it]

Out[4]:

{'ai_tasks': [{'ai_task': 'Text-to-Text',
   'explanation': 'The task involves generating text based on given input, which aligns with creating personalized responses and recommendations.'},
  {'ai_task': 'Summarization',
   'explanation': 'This task focuses on condensing information into a shorter form while preserving key details, useful for summarizing claims.'}]}

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

inference_engine = WMLInferenceEngine(
    model_name_or_path="ibm/granite-4-h-small",
    credentials={
        "api_key": "WML_API_KEY",
        "api_url": "WML_API_URL",
        "project_id": "WML_PROJECT_ID",
    },
    parameters=WMLInferenceEngineParams(decoding_method="greedy", repetition_penalty=1),
    backend="mellea",
)

inference_engine = WMLInferenceEngine( model_name_or_path="ibm/granite-4-h-small", credentials={ "api_key": "WML_API_KEY", "api_url": "WML_API_URL", "project_id": "WML_PROJECT_ID", }, parameters=WMLInferenceEngineParams(decoding_method="greedy", repetition_penalty=1), backend="mellea", )

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

[2026-03-18 22:28:21:220] - INFO - AIAtlasNexus - ✓ Created WML inference engine for model: ibm/granite-4-h-small, backend - MELLEA

In [9]:

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]

=== 22:28:29-INFO ======
SUCCESS

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

Out[9]:

{'ai_tasks': [{'ai_task': 'Text Classification',
   'explanation': "The use case involves generating personalized, relevant responses, recommendations, and summaries of claims for customers. This requires analyzing and classifying the customer's inquiries or claims into specific categories to provide appropriate responses. Text Classification is suitable for this as it assigns labels or classes to the given text based on its content, enabling the model to understand the context and intent of the customer's input and generate relevant output."},
  {'ai_task': 'Summarization',
   'explanation': 'Generating summaries of claims for customers is another aspect of the use case. Summarization is the task of producing a concise version of a document while preserving its important information. In this context, the AI model would take the detailed claim text as input and generate a shorter, more manageable summary that captures the key points. This helps support agents quickly grasp the essence of the claim and respond efficiently, enhancing the overall customer interaction experience.'},
  {'ai_task': 'Question Answering',
   'explanation': "Providing personalized responses to customer inquiries may involve answering specific questions related to their issues or claims. Question Answering models are designed to retrieve answers to questions from a given text. They can search for the most relevant information within a document or knowledge base to provide accurate and targeted answers to the customer's queries. By leveraging Question Answering, the AI model can assist support agents in delivering precise and satisfactory responses, improving customer satisfaction."}]}