Python code to get data from Db2 database tables, using AI framework
This project demonstrates how to connect an IBM Db2 database to a large language model (LLM) using LangChain and OpenAI—allowing you to run natural language queries and retrieve full SQL results with rich formatting.
- Connect to IBM Db2 using SQLAlchemy and ibm_db
- Patch LangChain's SQLDatabase to avoid row truncation
- Use OpenAI's GPT-4 via LangChain's ChatOpenAI wrapper
- Generate SQL queries from natural language prompts
- Display full query results in a clean tabulated format
- Includes Windows-specific DLL path setup for DB2 clients
| Module | Purpose |
|---|---|
| ibm_db, ibm_db_dbi | Low-level and DB-API 2.0 compliant interfaces for IBM Db2 |
| SQLAlchemy | Database abstraction and engine creation |
| LangChain | Framework for LLM-powered applications |
| ChatOpenAI | LangChain wrapper for OpenAI chat models |
| tabulate | Pretty-printing SQL results in table format |
-
Install Dependencies pip install ibm_db sqlalchemy langchain langchain-community langchain-openai tabulate
-
Configure Environment Set your Db2 credentials and OpenAI API key:
| Variable | Value |
|---|---|
| DB_USERNAM | "your-user" |
| DB_PASSWOR | "your-pw" |
| DB_HOST | "localhost" |
| DB_PORT | "25000" |
| DB_NAME | "SAMPLE" |
export OPENAI_API_KEY=your-openai-key
- Windows DLL Path Setup (Optional) If you're on Windows, the script automatically adds common DB2 client paths to the DLL search path. You can customize these paths if needed.
Run the script:
- python db2_langchain_agent.py
- Edit the query_input variable to test different natural language queries, e.g.: query_input += "List all departments and the number of employees in each."
- Interpret your query
- Generate and execute SQL
- Print tabulated results
- Provide LLM commentary
SQL Query: SELECT department_name, COUNT(*) FROM employees GROUP BY department_name;```
| department_name | count |
|---|---|
| Sales | 12 |
| Engineering | 25 |
| HR | 5 |
=== LLM Commentary (free text) === There are 3 departments with varying employee counts. Engineering has the most.
- SQLDatabaseNoLimit: Prevents LangChain from injecting LIMIT clauses
- ChatOpenAINoStop: Removes unsupported stop parameter from OpenAI calls
- AgentExecutor: Configured with return_intermediate_steps=True to extract raw SQL
- Requires a working IBM Db2 instance and valid credentials
- Assumes GPT-4 access via OpenAI