SWAT Code Generation and Execution in SAS Viya with Azure OpenAI and LangChain: Behind the Scenes

2 Likes

The SAS Scripting Wrapper for Analytics Transfer (SWAT) custom agent is designed to be an interactive AI assistant. It streamlines SAS code generation and execution for users, regardless of their expertise in SAS, Python, or CAS actions. In this post, we'll look into the agent's code, outline the prerequisites, and guide you through replicating the agent in your own environment.

Previously

Check out the video in SWAT Code Generation and Execution in SAS Viya with Azure OpenAI and LangChain for an insider's look at how this agent functions. I'll walk you through its operations, potential uses, and some of the intriguing behaviors that the custom agent displays.

Select any image to see a larger version.
Mobile users: To view the images, select the "Full" version at the bottom of the page.

Now it's time to reveal the code and walk you through the details.

Custom Agent Code

The core Python script, named sasgenexswat.py, establishes the SWAT LangChain custom agent that interfaces with SAS Viya

# LangChain custom SAS Viya agent with memory and only two tools
# Resource: https://python.langchain.com/docs/modules/agents/how_to/custom_agent

import os
import json
import urllib3
urllib3.disable_warnings

# Load the Large Language Model - Azure OpenAI
from langchain_openai import AzureChatOpenAI
from openai import AzureOpenAI

# Get Azure OpenAI environment variables from .env
from dotenv import load_dotenv

load_dotenv()
azure_oai_endpoint = os.getenv("AZURE_OAI_ENDPOINT")
azure_oai_key = os.getenv("AZURE_OAI_KEY")
azure_oai_deployment = os.getenv("AZURE_OAI_DEPLOYMENT")
azure_oai_model = os.getenv("AZURE_OAI_DEPLOYMENT")

os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["AZURE_OPENAI_API_KEY"] = azure_oai_key
os.environ["AZURE_OPENAI_ENDPOINT"] = azure_oai_endpoint
os.environ["OPENAI_API_VERSION"] = "2024-05-01-preview"

## We'll use GPT-4 as deployed model
llm = AzureChatOpenAI(deployment_name=azure_oai_model, temperature=0, max_tokens=2000)
print("We'll be using the following LLM:\n", llm, "\n")

# LLM Client
from openai import AzureOpenAI
client = AzureOpenAI(
    azure_endpoint = azure_oai_endpoint,
    api_key = azure_oai_key,
    api_version= "2024-05-01-preview"
    )


# Arguments
import sys
# Check if at least one argument is provided
if len(sys.argv) > 1:
    # sys.argv[1] is the first parameter passed to the script
    sas_viya_url = sys.argv[1]
    print(f"""Passing prompts to Azure OpenAI.
        \nAzure OpenAI will generate the SWAT code and execute it on: {sas_viya_url}.
          \nAssuming you obtained an access token first by running get_token_2023.py.""")
    print('\nRetrieving the saved token from api/access_token.txt')
    with open("api/access_token.txt", "r", encoding="UTF-8") as f:
        token = f.read()
else:
    print("No SAS Viya URL provided.")


import swat # SWAT session

# Define custom tools - SWAT GenEx
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool, StructuredTool, tool

import re # code parsing
from io import StringIO # exec output

@tool
def get_word_length(word: str) -> int:
    """Returns the length of a word."""
    return len(word)

@tool
def generate_swat_code(prompt: str) -> str:
    """Returns the generated SWAT Python code, based on the user prompt.
    The code should be aimed at printing the execution output, as it would be used further.
    Args:
        prompt: The prompt which will be sent to Azure OpenAI.
    """
    response = client.chat.completions.create(
        model=azure_oai_model,
        messages=[
                {"role": "system", "content": "You are a SAS Python SWAT package assistant that helps people write SWAT code. Return only the Python SWAT code, ready to be executed. This is the code in between ```python and ```. Assume there is already a connection named 'conn' defined."},
                # ... other system and user messages ...
                {"role": "user", "content": prompt},
            ]
        )
    sas_response = response.choices[0].message.content + '\n'
    # Strip code of characters
    pattern = re.compile(r"py\n(.*?)", re.DOTALL)
    match = pattern.search(sas_response)
    if match:
        swat_code = match.group(1).strip()
        print('\n',swat_code, '\n')
    else:
        swat_code = sas_response
    # Clean code block
    code_block = swat_code
    swat_code = code_block.replace("```python\n", "").replace("\n```", "")
    # Save the SWAT code to a file
    filename = "generated_swat_code.py"
    # Open the file first
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            current_content = file.read()
    except FileNotFoundError:
        current_content = ""
    # Prepend the new SWAT code to the existing content
    new_content = '# Prompt -> ' + prompt + '\n' '# Generated code -> ' + '\n' + swat_code + '\n' + current_content
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(new_content)
    print(f"Generated SWAT code has been saved to {filename}")
    return swat_code


@tool
def execute_swat_code(swat_code: str) -> str:
    """Useful when you need to execute generated Python SWAT code.
        Input to this tool is correct Python SWAT code, ready to be executed in a SAS Viya environment CAS session.
        Assume an already established SWAT connection, called 'conn'.
        Identify the results, summarize and provide a short explanation.
        If an error is returned, you may need to rewrite the SWAT code, using the generate_SWAT_code tool.
        Provide a Final Result: Summarize and provide a short explanation of what was done and what the outcome was.
        If you encounter an issue, detail what the issue is.
    Args:
        swat_code: The SWAT Python code which will be submitted to a CAS session in SAS Viya.
    """

    # Python code as a string stored in a variable
    print (f'\nI will submit the following code to {sas_viya_url}: \n',swat_code,'\n...End of code...\n')
    try:
        # Create SWAT connection using the url and a token obtained earlier
        port='443'
        conn_string= sas_viya_url + ':' + port + '/cas-shared-default-http'
        conn=swat.CAS(conn_string, password=token)
        # Create a dictionary to hold the execution context
        namespace = {'conn': conn}  # Add 'conn' to the namespace
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        redirected_output = sys.stdout = StringIO()
        redirected_error = sys.stderr = StringIO()
        try:
            # execute the SWAT code
            exec(swat_code, namespace)
        except Exception:
            # Optionally, handle the exception in some way here
            pass
        # Reset stdout and stderr to their original values
        sys.stdout = old_stdout
        sys.stderr = old_stderr

        # Get the captured output and errors
        output = redirected_output.getvalue()
        errors = redirected_error.getvalue()

        # Print the captured output and errors
        print("Captured Output:")
        print(output)
        print("Captured Errors:")
        print(errors)
        # Terminate SWAT connection
        conn.terminate()
        return output + errors
    except Exception as e:
        print(f'An error occurred: {e}')
        errors = e
        return errors


# Describe the tools. The description is important for the LangChain chain.
tools = [get_word_length, generate_swat_code, execute_swat_code]


# LangChain Agent with Memory

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

MEMORY_KEY = "chat_history"
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a world class AI assistant who helps people generate and execute SAS code. Interpret the results if asked.",
        ),
        MessagesPlaceholder(variable_name=MEMORY_KEY),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)


llm_with_tools = llm.bind_tools(tools)

from langchain_core.messages import AIMessage, HumanMessage
from langchain.agents.format_scratchpad.openai_tools import (
    format_to_openai_tool_messages,
)
from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser

chat_history = []

agent = (
    {
        "input": lambda x: x["input"],
        "agent_scratchpad": lambda x: format_to_openai_tool_messages(
            x["intermediate_steps"]
        ),
        "chat_history": lambda x: x["chat_history"],
    }
    | prompt
    | llm_with_tools
    | OpenAIToolsAgentOutputParser()
)

from langchain.agents import AgentExecutor

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


# Main function
def chat_with_open_ai():
    # ... as before
    while True:
        print("SAS and Azure OpenAI are listening. Write 'Stop' or press Ctrl-Z to end the conversation.")
        try:
            # SWAT start session
            #conn = start_or_retrieve_swat_session(sas_viya_url)
            # Get text from the keyboard
            q = input("Enter your question: \n")
            if q == "Stop":
                print("Conversation ended.")
                #term = conn.terminate()
                #print('SWAT session ended successfully.', term)
                break
            else:
                print("\nI will ask: ", q)
                result = agent_executor.invoke({"input": q, "chat_history": chat_history})
                chat_history.extend(
                    [
                        HumanMessage(content=q),
                        AIMessage(content=result["output"]),
                    ]
                )


        except EOFError:
            break

# main
try:
    chat_with_open_ai()
except Exception as err:
    print("Encountered exception. {}".format(err))

Custom Agent Code Explanation

The provided Python code integrates several components to create a custom agent that leverages Azure's OpenAI Large Language Model (LLM), specifically GPT-4, to generate and execute SAS code via the SAS Python SWAT (SAS Scripting Wrapper for Analytics Transfer) package. This agent is designed to facilitate interactions with SAS Viya, a cloud-enabled, in-memory analytics engine.

Here's a step-by-step breakdown of the code's functionality:

Imports and Environment Setup:
- The script starts by importing necessary Python modules.
- It loads environment variables from a .env file, which includes credentials and endpoints for Azure OpenAI services.
- The AzureChatOpenAI class from the langchain_openai package is instantiated with parameters set for the LLM, such as the model's deployment name and the maximum number of tokens to generate.
Argument Handling:
- The script checks for command-line arguments, expecting at least one argument: the SAS Viya URL. If provided, it reads an access token from a local file to use for authentication with the SAS Viya server.

Custom Tools Definition:

Several functions are decorated with @tool from the langchain package, indicating that they are custom tools that can be used within the agent's execution flow.
generate_swat_code tool sends the user prompt to Azure OpenAI, which returns the generated SWAT Python code. Letting an LLM to generate the code is the innovative aspect of the custom agent.

@tool
def generate_swat_code(prompt: str) -> str:
    """Returns the generated SWAT Python code, based on the user prompt.
    The code should be aimed at printing the execution output, as it would be used further.
    Args:
        prompt: The prompt which will be sent to Azure OpenAI.
    """
    response = client.chat.completions.create(
        model=azure_oai_model,
        messages=[
                {"role": "system", "content": """You are a SAS Python SWAT package assistant that helps people write SWAT code.
        Return only the Python SWAT code, ready to be executed. This is the code in between ```python and ```.
        Assume there is already a connection named 'conn' defined."},
                # ... other system and user messages ...
                {"role": "user", "content": prompt},
            ]
        )

The code is then cleaned and saved to a file.
execute_swat_code retrieves the code from the generate_swat_code, establishes a SWAT connection, and then tries to run the generated SWAT code on a SAS Viya server using the SWAT package. It captures the output and errors, returning them to the agent. The SWAT connection is then terminated to liberate CAS compute resources on the SAS Viya server.

@tool
def execute_swat_code(swat_code: str) -> str:
    """Useful when you need to execute generated Python SWAT code.
        Input to this tool is correct Python SWAT code, ready to be executed in a SAS Viya environment CAS session.
        Assume an already established SWAT connection, called 'conn'.
        Identify the results, summarize and provide a short explanation.
        If an error is returned, you may need to rewrite the SWAT code, using the generate_SWAT_code tool.
        Provide a Final Result: Summarize and provide a short explanation of what was done and what the outcome was.
        If you encounter an issue, detail what the issue is.
    Args:
        swat_code: The SWAT Python code which will be submitted to a CAS session in SAS Viya.
    """
    try:
        # Create SWAT connection using the url and a token obtained earlier
        port='443'
        conn_string= sas_viya_url + ':' + port + '/cas-shared-default-http'
        conn=swat.CAS(conn_string, password=token)
        # Create a dictionary to hold the execution context
        namespace = {'conn': conn}  # Add 'conn' to the namespace
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        redirected_output = sys.stdout = StringIO()
        redirected_error = sys.stderr = StringIO()
        try:
            # execute the SWAT code
            exec(swat_code, namespace)
        except Exception:
            # Optionally, handle the exception in some way here
            pass
        #

The SWAT code runs inside the exec(swat_code, namespace) function, after establishing a connection within the namespace. However, using exec means entrusting code execution to a custom agent without scrutiny, which isn't the most secure method for submitting SWAT code to SAS Viya. If you're aware of safer alternatives, I'm eager to hear them.

LangChain Agent with Memory:

A chat prompt template is defined using the LangChain package, which describes the behavior of an AI assistant capable of interpreting and acting on user input.
The llm_with_tools object binds the custom tools to the LLM.

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
MEMORY_KEY = "chat_history"
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a world class AI assistant who helps people generate and execute SAS code. Interpret the results if asked.",
        ),
        MessagesPlaceholder(variable_name=MEMORY_KEY),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)
llm_with_tools = llm.bind_tools(tools)

The agent object is created by composing several functions and components, including the prompt, LLM, and an output parser.
An AgentExecutor is instantiated, which handles the execution of the agent with the custom tools.

Main Function (chat_with_open_ai):
- The main loop of the function listens for user input and uses the agent_executor to process the input, generating and executing SAS code accordingly. The chat history is updated with each interaction.
- The user can type "Stop" to end the conversation.
- The script is wrapped in a try-except block to handle any exceptions that occur during the chat loop.
Execution:
- Finally, the chat_with_open_ai function is called to start the interactive session.

Now let’s look at what you need to build this custom agent.

SAS Viya Environment

To build this custom agent, you'll require a SAS Viya deployment. From your SAS Viya environment, you'll need to gather the following details:

The SAS Viya URL
A SAS username
The corresponding SAS user's password

Additionally, ensure that CAS (Cloud Analytic Services) access is enabled for external connections to SAS Viya.

Azure Deployment

If your SAS Viya is hosted on Azure, it's essential to configure an inbound security rule. This rule must include the IP address of the computer running Visual Studio Code and specify port 443. The target for this rule should be the public IP address of your SAS Viya's load balancer.

Azure OpenAI

To get started, you'll need to create a resource and deploy a model. If you're unfamiliar with the process, the Azure OpenAI QuickStart provides step-by-step guidance. To interact with Azure OpenAI, you require an endpoint, a key, and a deployment.

For this project, I selected the 1106-Preview GPT-4 model version, which was trained on data from December 2023. I've found that this version excels at generating SWAT code. Remember, choosing the right model is as crucial as the code that powers your application

Visual Studio Code

We will use Visual Studio Code with the Python extension installed to run the custom agent. If you’re a regular user of Visual Studio Code, great.

Python Requirements

Create a requirements.txt file that includes these Python packages:

python-dotenv
pandas
urllib3
swat
openai
langchain
langchain_openai
langchain-core
langchain-community

Azure OpenAI Configuration

In the directory containing your main Python script, create a .env file to store your Azure OpenAI configuration parameters. The file should include the following lines:

AZURE_OAI_ENDPOINT='https://myuser.openai.azure.com/'
AZURE_OAI_KEY='key_here' #Azure Open AI key here
AZURE_OAI_DEPLOYMENT='gpt-4'

The sasgenexswat.py program is configured to read these environment variables from your .env file.

SWAT Setup

Certificates

Firstly, you'll need to obtain the SAS Viya certificates, which are in the .PEM file format, and download them to your local machine.

For detailed instructions, please refer to the provided documentation.

Get a SAS Viya Access Token

The sasgenexswat.py program is designed to read an access token from a file called api/access_token.txt.

For SWAT to establish a connection within the custom agent, you'll need an OAuth token. You have two options for obtaining this token:

Use any method you're comfortable with to acquire an access token and save it to api/access_token.txt.
Follow the instructions in the post Discover Your Data with SAS Information Catalog APIs from Python – Access to run a program called get_token_2023.py. The program requires several variables.

Variables

Once you have your certificates, open Visual Studio Code, launch a Bash terminal, and proceed with the setup.

SASVIYAURL=my_sas_viya_url
SASUSER=fill_in_here # SAS user
SASPASS=pass_here # SAS user password
CLIENT=fill_in_here # client id
CLIENTSECRET=fill_in_here # client secret
PEMPATH='fill_here'
# for example PEMPATH='C:\Users\myuser\Downloads\trustedcerts.pem'
export CAS_CLIENT_SSL_CA_LIST=${PEMPATH}

Get a Token

Run a program called get_token_2023.py. You'll need to provide several environment variables.

python get_token_2023.py $SASVIYAURL $CLIENT $CLIENTSECRET $SASUSER $SASPASS $PEMPATH

Install Python Packages

To install the Python packages listed in the requirements.txt file, execute the following command:

pip install -U -r requirements.txt

For additional context, if you're interested in the specific versions of LangChain, OpenAI, and SWAT packages used in the demo, you can retrieve them using these commands:

pip freeze | grep lang
pip freeze | grep openai
pip freeze | grep swat
python --version

At the time of writing this post, the versions of the Python packages I am using are as follows:

langchain==0.2.14
langchain-community==0.2.12
langchain-core==0.2.33
langchain-openai==0.1.22
langchain-text-splitters==0.2.2
langsmith==0.1.93
langchain-openai==0.1.22
openai==1.41.0
swat==1.14.0
Python 3.11.9

Run the Program

To execute the program, you'll need to provide your SAS Viya URL, which you've previously set as an environmental variable. Run the program by typing the following command:

python sasgenexswat.py $SASVIYAURL

Prompts

Check out the video in SWAT Code Generation and Execution in SAS Viya with Azure OpenAI and LangChain for prompts and responses.

Start asking questions such as:

List all files in the SAMPLES caslib
List all CASLIBS.
List all available in-memory tables in the 'Public' caslib.
Load a CSV file from a URL into the Public caslib and promote the table. URL is
'https://support.sas.com/documentation/onlinedoc/viya/exampledatasets/cars.csv'
List all available in-memory tables in the 'Public' caslib. Confirm CARS has been loaded.
Print the columns of the CARS table.
Request simple summary statistics for the CARS table from Public caslib.
Describe the table CARS in the Public caslib.
How many rows in the CARS CAS table from Public have the Make Toyota or Acura?
From CARS, create a new calculated column NewCol2 = round(((MSRP - Invoice) / Invoice) * 100 and print three records to confirm the column creation.
Find the top 10 MPG_City consumption values ordered ascending. Print the car make and model.
Group the CAS table CARS from Public by Origin. Aggregate by column Invoice. Use means as a measure. Print the aggregation.
Stop or CTRL+Z when you want to exit the program.

Conclusions

The SWAT code Generator-Executor (GenEx) custom agent approach is innovative, as it uses a large language model, GPT-4, to generate Python SWAT code based on user prompts. The custom agent is relying on just two main components: a SWAT code generator and a SWAT code executor. These tools work in tandem to run the generated SAS code within your chosen SAS Viya environment, managing both the output and any potential errors through the agent, which then crafts a final response.

SWAT code execution is the crucial step that makes large language model (LLM) code generation practical for SAS Viya users; it determines whether the code functions correctly in their environment.

Should an error arise, the agent is programmed to interpret the issue and attempt to regenerate the code, aiming for a successful execution. This can be a double-edged sword; while it's a source of strength, it can also be a weakness, as the custom agent may occasionally become stuck in a loop.

It's remarkable how just two hundred fifty lines of code can integrate enterprise analytics and data management platforms such as SAS Viya with AI platforms like Azure OpenAI. The linchpin that unites these diverse platforms remains open source.

I hope you found this article insightful. I encourage you to give it a try! Please feel free to reach out with feedback or suggestions for enhancing the agent or taking its capabilities to the next level.

Acknowledgements

Special thanks to Grace Liao from SAS Research & Development for inspiring me to try the approach.

Thanks to Peter Styliadis for his great SWAT Series.

Additional Resources

Thank you for your time reading this post. If you liked the post, give it a thumbs up! Please comment and tell us what you think about having conversations with your data. If you wish to get more information, please write me an email.

Find more articles from SAS Global Enablement and Learning here.