# Build a  PostgreSQL AI Agent Using LangChain + Ollama 

## 🔥 Introduction

What if you could query your database like this:

> *"Show me top 10 users by revenue"*

…and get instant results—without writing SQL?

Welcome to the world of **AI-powered database agents**.

In this tutorial, you'll learn how to build a **secure PostgreSQL AI agent** using:

*   🧩 **LangChain** — for agent orchestration and tool chaining
    
*   🦙 **Ollama** — to run a local LLM with zero API cost
    
*   🐘 **PostgreSQL** — as the target database
    
*   🛡️ **Custom SQL safety guard** — to block destructive queries
    

By the end, you'll have a production-ready AI database assistant that understands natural language and safely executes SQL queries.

> 💻 **Source Code:** [https://github.com/icon-gaurav/postgres-agent](https://github.com/icon-gaurav/postgres-agent)

* * *

## 🤖 What is a PostgreSQL AI Agent?

A **PostgreSQL AI Agent** is an LLM-powered system that:

*   Converts natural language → SQL queries
    
*   Executes queries on PostgreSQL
    
*   Returns structured results
    

> 👉 Think of it as ChatGPT for your database, but controlled and safe.

* * *

## ⚙️ Tech Stack

| Tool | Purpose |
| --- | --- |
| [LangChain](https://python.langchain.com) | AI agent orchestration, tool use |
| [Ollama](https://ollama.com) | Local LLM inference, no API key required |
| [langchain-ollama](https://pypi.org/project/langchain-ollama/) | LangChain ↔ Ollama integration |
| [psycopg2](https://pypi.org/project/psycopg2/) | PostgreSQL database adapter for Python |
| Python | Core backend runtime |

* * *

## 🧱 Architecture

![](https://cdn.hashnode.com/uploads/covers/616e87a8f1f4c944cc6b49a1/e4f8ba23-4815-415e-8275-175b7f19216d.png align="center")

* * *

## 🔌 Step 1: PostgreSQL Connection

First, configure your database connection using `psycopg2`:

```python
import psycopg2
 
DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "database": "postgres",
    "user": "postgres",
    "password": "root",
}
 
def get_connection():
    return psycopg2.connect(**DB_CONFIG)
```

> 🔐 **Security tip:** In production, load credentials from environment variables or a secrets manager — never hardcode passwords.

* * *

## 🛠️ Step 2: Create Database Tools

LangChain agents interact with the world through **tools** — Python functions decorated with `@tool` that the LLM can invoke by name.

### 📋 List Tables Tool

```python
@tool
def list_tables() -> str:
    """List all tables in the database."""
    conn = get_connection()
    try:
        cur = conn.cursor()
        cur.execute("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'public'
        """)
        tables = [row[0] for row in cur.fetchall()]
        return f"Tables: {', '.join(tables)}" if tables else "No tables found."
    finally:
        conn.close()
```

This enables **dynamic schema discovery** — the agent doesn't need hardcoded table names.

### 📑 Get Table Schema

```python
@tool
@tool
def get_table_schema(table_name: str) -> str:
    """Get the schema (columns and types) of a specific table."""
    conn = get_connection()
    try:
        cur = conn.cursor()
        cur.execute("""
            SELECT column_name, data_type, is_nullable
            FROM information_schema.columns
            WHERE table_schema = 'public' AND table_name = %s
            ORDER BY ordinal_position
        """, (table_name,))
        columns = cur.fetchall()
        if not columns:
            return f"Table '{table_name}' not found."
        schema = "\n".join([f"  {col[0]} ({col[1]}, nullable={col[2]})" for col in columns])
        return f"Schema for '{table_name}':\n{schema}"
    finally:
        conn.close()

```

Allows the agent to understand:

*   Column names
    
*   Data types
    
*   Constraints
    

### ⚡ Execute SQL Tool

```python
@tool
def execute_sql(query: str) -> str:
    """Execute a SQL query against the PostgreSQL database and return results. Use this for SELECT queries."""
    is_safe, reason = validate_read_only_sql(query)
    if not is_safe:
        return f"Safety Guard: Blocked query. {reason}"

    conn = get_connection()
    try:
        cur = conn.cursor()
        cur.execute(query)
        if cur.description:
            columns = [desc[0] for desc in cur.description]
            rows = cur.fetchall()
            if not rows:
                return "Query returned no results."
            result = " | ".join(columns) + "\n"
            result += "\n".join([" | ".join(str(v) for v in row) for row in rows[:50]])
            if len(rows) > 50:
                result += f"\n... ({len(rows)} total rows)"
            return result
        else:
            conn.commit()
            return f"Query executed successfully. Rows affected: {cur.rowcount}"
    except Exception as e:
        conn.rollback()
        return f"SQL Error: {e}"
    finally:
        conn.close()

```

This is the **core execution layer**.

* * *

## Step 3: SQL Safety Guard — Prevent Destructive Queries

Allowing an LLM to run arbitrary SQL is a critical security risk. The **SQL safety guard** validates every query before execution.

### Read-Only Allowlist

Only these SQL statement types are permitted:

| Allowed | Blocked |
| --- | --- |
| `SELECT` | `INSERT` |
| `WITH` (CTEs) | `UPDATE` |
| `SHOW` | `DELETE` |
| `EXPLAIN` | `DROP` |
| — | `ALTER` |
| — | `TRUNCATE` |

### 🧼 Normalize Queries

We remove:

*   Comments
    
*   Strings
    
*   Hidden injections
    

> 👉 This ensures **safe AI execution** in production environments.

* * *

## 🧠 Step 4: Setup Ollama (Local LLM)

[Ollama](https://ollama.com) lets you download and run large language models entirely on your own machine — no cloud account, no API key, no usage fees.

> 📚 **Official Resources:**
> 
> *   🌐 Website: [ollama.com](https://ollama.com)
>     
> *   📖 Documentation: [docs.ollama.com](https://docs.ollama.com)
>     
> *   🗂️ Model Library: [ollama.com/library](https://ollama.com/library)
>     
> *   🐙 GitHub: [github.com/ollama/ollama](https://github.com/ollama/ollama)
>     

* * *

### 🔽 Pulling a Model

Once Ollama is running, pull the model used in this project:

```bash
ollama pull qwen2.5:7b
```

You can verify it's available with:

```bash
ollama list
```

Browse all available models at [ollama.com/library](https://ollama.com/library). Some good alternatives for SQL agents:

| Model | Command | Notes |
| --- | --- | --- |
| Qwen 2.5 7B | `ollama pull qwen2.5:7b` | Used in this tutorial |
| Llama 3.1 8B | `ollama pull llama3.1` | Strong general-purpose model |
| DeepSeek-R1 7B | `ollama pull deepseek-r1` | Good reasoning ability |
| Mistral 7B | `ollama pull mistral` | Fast and lightweight |

* * *

### 📦 Installing the LangChain Ollama Package

The LangChain integration for Ollama lives in the dedicated [`langchain-ollama`](https://pypi.org/project/langchain-ollama/) package:

```bash
pip install langchain-ollama
```

> 📖 **Package References:**
> 
> *   📦 PyPI: [pypi.org/project/langchain-ollama](https://pypi.org/project/langchain-ollama/)
>     
> *   🔗 LangChain Docs: [docs.langchain.com — ChatOllama](https://docs.langchain.com/oss/python/integrations/chat/ollama)
>     
> *   📐 API Reference: [reference.langchain.com/python/langchain-ollama](https://reference.langchain.com/python/langchain-ollama)
>     

### ⚙️ Configuring ChatOllama

Now initialize the LLM in your Python code:

```python
from langchain_ollama import ChatOllama
 
llm = ChatOllama(model="qwen2.5:7b", temperature=0)
```

Setting `temperature=0` makes the model deterministic — essential for reliable SQL generation. You can tune other key parameters as needed:

| Parameter | Default | Description |
| --- | --- | --- |
| `model` | required | Model name from `ollama list` |
| `temperature` | `0.8` | Creativity — use `0` for SQL tasks |
| `num_predict` | `128` | Max tokens to generate |
| `base_url` | `http://localhost:11434` | Ollama server URL |

**Why Ollama?**

*   ✅ No API cost
    
*   ✅ Runs fully locally — data never leaves your machine
    
*   ✅ Privacy-friendly — ideal for sensitive database workloads
    
*   ✅ Fast inference with GPU support
    
*   ✅ Supports dozens of open-source models
    

* * *

## 🔗 Step 5: Create LangChain Agent

```python
tools = [list_tables, get_table_schema, execute_sql]
agent = create_agent(llm, tools)
```

LangChain allows the AI to:

*   Decide which tool to use
    
*   Chain multiple steps
    
*   Reason dynamically
    

* * *

## 💬 Step 6: Interactive Chat Loop

```python
while True:
    user_input = input("\nYou: ").strip()
    if user_input.lower() in ("exit", "quit"):
        print("Goodbye!")
        break
    if not user_input:
        continue
```

This makes your agent:

*   Conversational
    
*   Stateful
    
*   Easy to debug
    

* * *

## 🧾Step 7: Debugging & Observability

Visibility into what the agent is doing is essential for development. This helper function prints each tool call and its result:

```python
def print_turn_details(messages: list[BaseMessage]) -> None:
    final_response = ""

    for message in messages:
        if isinstance(message, AIMessage):
            for tool_call in message.tool_calls:
                tool_name = tool_call.get("name", "unknown_tool")
                tool_args = format_tool_payload(tool_call.get("args", {}))
                print(f"\nTool call: {tool_name}({tool_args})")

            content = format_content(message.content).strip()
            if content:
                final_response = content

        elif isinstance(message, ToolMessage):
            tool_name = getattr(message, "name", None) or "tool"
            tool_output = format_content(message.content).strip() or "(no output)"
            print(f"\nTool response [{tool_name}]: {tool_output}")

    if final_response:
        print(f"\nAgent: {final_response}")
    else:
        print("\nAgent: I couldn't generate a response.")
```

Shows:

*   Tool calls
    
*   Tool outputs
    
*   Final response
    

> 👉 This is extremely useful for **debugging agent behavior**.

* * *

## 🧪 Example Queries

Try asking:

*   `"List all tables"`
    
*   `"Show schema of users table"`
    
*   `"Get top 5 users by revenue"`
    
*   `"How many orders were placed last month?"`
    

![](https://cdn.hashnode.com/uploads/covers/616e87a8f1f4c944cc6b49a1/0865e8a5-269f-4bce-a094-e963a49d5aed.png align="center")

* * *

## Real-World Use Cases for a Natural Language Database Agent

This architecture can power a wide range of applications:

**📊 AI-Powered Analytics Dashboards** — let non-technical stakeholders query live data in plain English without learning SQL.

**💬 Internal Data Chatbots** — embed in Slack or Microsoft Teams so product and ops teams can self-serve data questions.

**🧾 Automated Reporting** — schedule the agent to answer recurring questions and generate daily or weekly reports.

**🏢 SaaS Admin Panels** — give your ops team a natural language interface to your product database.

**🤖 AI Copilots for Data Teams** — speed up analyst workflows by auto-generating SQL drafts from plain-English specs.

* * *

## 🎯 Conclusion

You've built more than just a demo. This is a **secure, extensible AI database agent** that can be used in real-world applications.

**Key Takeaways:**

*   LangChain simplifies agent workflows
    
*   Ollama enables local LLM execution
    
*   SQL safety is critical
    
*   Tool-based architecture = powerful AI agents
    

The core insight: wrapping your database in typed, well-described LangChain tools gives the LLM exactly the context it needs to generate correct SQL — without ever exposing raw database access.

> [📦 Full Source Code](https://github.com/icon-gaurav/postgres-agent) — complete working implementation from this tutorial