AI-Powered Features

pip install mage-ai[ai]

GPT_MODEL = "gpt-4o"

export ENABLE_HUGGING_FACE=1

ai_config:
  mode: hugging_face
  hugging_face_config:
    model_name: your-model-name

class LLMPipelineWizard:
    def __init__(self):
        ai_config = AIConfig.load(config=get_repo_config().ai_config)
        if ENABLE_OPEN_AI and ai_config.mode == AIMode.OPEN_AI:
            self.client = OpenAIClient(ai_config.open_ai_config)
        elif ENABLE_HUGGING_FACE and ai_config.mode == AIMode.HUGGING_FACE:
            self.client = HuggingFaceClient(ai_config.hugging_face_config)

wizard = LLMPipelineWizard()

# Generate complete pipeline
pipeline = await wizard.async_generate_pipeline_from_description(
    "Load data from MySQL and Postgres, filter rows with price > 100, and save to BigQuery"
)

# Returns dictionary of blocks
# {
#   '1': { block_type: 'data_loader', ... },  # MySQL loader
#   '2': { block_type: 'data_loader', ... },  # Postgres loader
#   '3': { block_type: 'transformer', ... },  # Filter transformer
#   '4': { block_type: 'data_exporter', ... } # BigQuery exporter
# }

wizard = LLMPipelineWizard()

# Add comments to functions in a block
commented_code = await wizard.async_generate_comment_for_block(
    block_content=your_block_code
)

print(commented_code)  # Code with added docstrings

PROMPT_FOR_BLOCK = """
The {file_type} delimited by triple backticks is used to {purpose}.
Write a documentation based on the {file_type}. {add_on_prompt}
Ignore the imported libraries and the @test decorator.
```{block_content}```"""

PROMPT_FOR_CUSTOMIZED_CODE_IN_PYTHON = """
The content within the triple backticks is a code description.

Your task is to answer the following two questions.

1. Is there any filter logic mentioned in the description to remove rows or columns of the data?
If yes, write ONLY the filter logic as a if condition without "if" at beginning.
Return your response as one field in JSON format with the key "action_code".

2. Does the description mention any columns or rows to aggregrate on or group by?
If yes, list ONLY those columns in an array and return it as a field in JSON response
with the key "arguments".

<code description>: ```{code_description}```

Provide your response in JSON format.
"""

from mage_ai.ai.openai_client import OpenAIClient
from mage_ai.orchestration.ai.config import OpenAIConfig

client = OpenAIClient(OpenAIConfig(openai_api_key='sk-...'))

# Custom inference
result = await client.inference_with_prompt(
    variable_values={
        'input_data': 'your data',
        'requirement': 'your requirement'
    },
    prompt_template="Given {input_data}, generate code to {requirement}",
    is_json_response=True  # Expect JSON response
)

response = client.openai_client.chat.completions.create(
    model=GPT_MODEL,
    messages=messages,
    tools=tools,
    tool_choice={
        "type": "function", 
        "function": {"name": CLASSIFICATION_FUNCTION_NAME}
    },
)

max_retries = 2
attempt = 0
response = self.__chat_completion_request(messages)
while attempt <= max_retries and isinstance(response, Exception):
    response = self.__chat_completion_request(messages)
    attempt += 1

import asyncio
from mage_ai.ai.llm_pipeline_wizard import LLMPipelineWizard

async def create_ai_pipeline():
    wizard = LLMPipelineWizard()
    
    # 1. Generate pipeline structure
    blocks = await wizard.async_generate_pipeline_from_description(
        "Extract sales data from Snowflake, aggregate by region, and load to BigQuery"
    )
    
    # 2. Generate documentation
    for block_id, block_config in blocks.items():
        # Generate code comments
        if block_config['language'] == 'python':
            commented = await wizard.async_generate_comment_for_block(
                block_config['content']
            )
            block_config['content'] = commented
    
    # 3. Create pipeline with blocks
    # ... pipeline creation logic ...
    
    return blocks

# Run async workflow
blocks = asyncio.run(create_ai_pipeline())

if not resp.startswith('{') and not resp.endswith('}'):
    resp = f'{{{resp.strip()}}}'