# Upgrading My PostgreSQL AI Agent: 3 Architecture Decisions I Made (And Why)

*This is Part 2 of an ongoing series on building a production-ready PostgreSQL AI agent using LangChain and Ollama. If you haven't read* [*Part 1*](https://gauravbytes.dev/build-a-postgresql-ai-agent-using-langchain-ollama)*, start there — it covers the foundation: connecting to PostgreSQL, setting up tools, and getting a working chat loop running locally.*

## The Prototype Was Working. That Was the Problem.

After Part 1, I had a working agent. You could type a question in plain English, it would figure out the right SQL, run it, and return results. Impressive enough to demo. Not solid enough to build on.

The more I used it, the more I noticed the cracks:

*   The database connection was hardcoded. Switching databases meant editing the source file.
    
*   Every time the agent needed to understand a table, it fetched the schema live — adding latency and unnecessary DB hits.
    
*   The conversation had no memory. Every session started from zero.
    
*   The SQL safety check lived inside the tool itself, meaning there was nothing stopping me (or a future contributor) from adding another path that bypassed it entirely.
    

None of these are showstoppers at the prototype stage. All of them become real problems the moment you want to hand the tool to someone else, open-source it, or run it against a database that actually matters.

This article covers three architectural decisions I made to address these — what I changed, what I considered, and why I landed where I did.

![](https://cdn.hashnode.com/uploads/covers/616e87a8f1f4c944cc6b49a1/595ef207-16b9-409d-9903-b982bd8aa91a.png align="center")

## Decision 1: A Single SQL Execution Gateway

### The Problem

In v1, the `execute_sql` tool handled everything inline: it validated the query, opened a connection, ran it, and returned the result — all in one function that the LangChain agent called directly.

That works. But it means the safety logic and the execution logic are coupled inside a single tool. If I ever added a second tool — say, one that runs EXPLAIN plans or fetches row counts — I'd have to either duplicate the safety check or remember to call some shared validator manually. Both options are brittle.

The deeper issue: there was no structural guarantee that every SQL statement passed through a single checkpoint. The safety net existed, but it was enforced by convention, not architecture.

### The Decision

I moved all database interaction behind a single `execute_sql()` function in `db.py`. This is the only place in the codebase that runs SQL against the database. Every tool, every command, every path goes through it — no exceptions.

Alongside this, I introduced two custom exceptions:

```python
class SQLSafetyError(Exception):
    """Raised when a query violates the read-only safety policy."""
    pass

class SQLExecutionError(Exception):
    """Raised when a valid query fails during execution."""
    pass
```

And the connection API became private — a `_get_connection()` function that nothing outside `db.py` should call directly.

```python
def _get_connection():
    """Private. Use execute_sql() — do not call this directly."""
    return psycopg2.connect(**_load_db_config())
```

### Why This Matters

The separation of `SQLSafetyError` and `SQLExecutionError` is deliberate. A blocked query and a failed query are different situations — they should produce different responses to the user and, eventually, different log entries. Collapsing them into a single generic error message loses information that's useful for debugging.

The private connection API is about making the right path the only path. A future contributor shouldn't have to know to call the safety check before running a query — the architecture should make that automatic. If `_get_connection()` is internal and `execute_sql()` is the public interface, there's no accidental bypass.

This is the kind of constraint that feels like overhead when you're the only developer. It pays off the moment the codebase has more than one pair of eyes on it.

## Decision 2: Schema Caching at Startup

### The Problem

In v1, schema discovery was lazy. When the agent needed to understand a table's structure, it called the `get_table_schema` tool, which hit the database on demand.

This created a few practical problems:

*   **Latency.** Every query that involved schema introspection added a round-trip to the database before the actual query ran.
    
*   **Inconsistency.** The LLM sometimes skipped the schema fetch and guessed column names based on the table name alone. It was usually right. When it was wrong, the error wasn't obvious.
    
*   **Context fragmentation.** Schema information arrived piecemeal, one table at a time, as the conversation progressed. The agent never had a complete picture of the database upfront.
    

### The Decision

The caching implementation is deliberately simple. A module-level variable acts as the cache, populated on the first call and reused on every subsequent one:

```python
def get_cached_tables() -> list[str]:
    """Get list of all tables in public schema (cached after first call)."""
    global _schema_cache
    if _schema_cache is not None:
        return _schema_cache
    try:
        result = execute_sql("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'public'
            ORDER BY table_name
        """)
        _schema_cache = [row[0] for row in result.rows]
    except (SQLSafetyError, SQLExecutionError):
        _schema_cache = []
    return _schema_cache
```

Two things worth noting here. First, this goes through `execute_sql()` — the single gateway from Decision 1 — so the safety check applies even to internal schema queries. Second, if the fetch fails, the cache is set to an empty list rather than left as `None`. That means a second call won't retry a failed fetch mid-session, which keeps the behavior predictable.

The cached table list feeds into a `{schema_context}` block that gets injected into the system prompt at startup:

```python
"""You are PGChat, an expert PostgreSQL database assistant. You have access to tools to inspect and query the connected database.

Current date and time: {current_time}

Guidelines:
- Use the SCHEMA SNAPSHOT below to answer questions about tables and columns — do NOT call list_tables or get_table_schema unless the user explicitly asks to refresh.
- Prefer short, precise answers. When showing data, summarize it unless the user asks for raw output.
- Never guess table or column names — verify against the schema snapshot.
- When you write SQL, explain what it does briefly.
- If a query fails, analyze the error and suggest fixes.
- Use get_table_sample to preview data when helpful.
- Use search_schema when the user asks "which table has X column?"
- Use run_query to execute SQL queries.

{schema_context}"""
```

The guideline `do NOT call list_tables or get_table_schema unless the user explicitly asks to refresh` is doing real work here. Without it, the LLM will sometimes call those tools anyway out of habit, even when the snapshot has the answer. Explicitly instructing it to trust the snapshot eliminates those redundant tool calls in practice.

A `/refresh-schema` command lets the user reset `_schema_cache` to `None` and trigger a fresh fetch if the schema has changed during a session.

### Why This Matters

The key insight is that schema information is almost never volatile within a single session. Tables don't get columns added or dropped while you're mid-conversation. Fetching once and injecting as context is both faster and more reliable than fetching on demand.

The system prompt approach has a meaningful side effect: the LLM has the full schema available before it generates its first query. It doesn't have to reason about what tools to call to understand the database — that context is already there. In practice, this reduces the number of tool calls per user question and noticeably improves query accuracy, especially on databases with multiple tables that share similar naming patterns.

The `/refresh-schema` escape hatch matters for user trust. Schema caching only works if users are confident the snapshot reflects reality. Giving them a command to invalidate it manually means they're never stuck with stale information.

One tradeoff worth acknowledging: for databases with hundreds of tables, injecting the full schema as a system message could push against context window limits. For the current use case this isn't a problem, but it's something to revisit as the tool grows. A reasonable next step would be selectively injecting only the tables relevant to the current query — something worth exploring in a later iteration.

## Decision 3: Session Memory with JSON (For Now)

### The Problem

In v1, every conversation started from scratch. The agent had no memory of what you'd asked before, what tables you'd explored, or what the last query returned. For a quick one-off query this is fine. For anything resembling a real workflow — iterating on a report, exploring a dataset across multiple questions — it's a significant limitation.

### The Decision

I added persistent session memory backed by JSON. At the end of each conversation turn, the session state is written to a `.json` file. When you start the agent again, it loads that file and resumes where you left off.

The interface is intentionally minimal:

```python
def save_session(session: Session, path: str) -> None:
    """Persist session state to disk."""
    ...

def load_session(path: str) -> Session | None:
    """Load a previous session, or return None if none exists."""
    ...
```

### Why JSON and Not SQLite or PostgreSQL?

This is the question I spent the most time on, because it seems obvious to reach for a database when you're already using one.

The case for SQLite: structured queries, better concurrency handling, easier to inspect with standard tooling. The case for PostgreSQL: the agent is already connected to one, so no additional dependency. Both are reasonable. Neither is what I chose.

JSON wins for now for a few specific reasons:

**Zero additional dependencies.** The agent already requires Python, psycopg2, LangChain, and Ollama. Adding a SQLite ORM or a second PostgreSQL connection for session storage would increase the setup surface area for something that isn't the core feature. For a CLI tool that someone installs with `pip`, simpler setup matters.

**Human-readable state.** A JSON session file is easy to inspect, copy, share, and debug without any tooling. If something goes wrong with session state, you can open the file in a text editor and understand it immediately.

**The interface abstraction is what matters.** Because `save_session()` and `load_session()` are the only way the rest of the codebase touches session state, the storage backend can be swapped later without changing anything else. If the JSON approach hits a real limitation — concurrent sessions, large history, search across sessions — migrating to SQLite or PostgreSQL is a one-file change. Building that migration path in now, by keeping the interface clean, costs nothing.

The honest answer is also that JSON is good enough for the current use case. The agent runs as a single-user CLI tool. There's no concurrency to manage, no performance requirement that JSON can't meet, and no schema to maintain. Reaching for a database here would be solving a problem I don't have yet.

![](https://cdn.hashnode.com/uploads/covers/616e87a8f1f4c944cc6b49a1/3f06f693-f985-4bfc-bdc6-c6da51f9d02d.png align="center")

## What's Next

These three changes move the agent from a prototype to something I'd actually hand to another developer. The architecture is more intentional, the failure modes are clearer, and the session experience is meaningfully better.

The next step is packaging — making the agent installable as a proper CLI tool via `pip install postgres-agent`, with a clean entry point and configuration handled through environment variables or a config file rather than source edits.

After that, the session storage decision will get revisited. As the tool supports multiple named sessions or shared use cases, JSON will start to show its limits. That's the right time to migrate — not before.

The source code is on GitHub: [github.com/icon-gaurav/pgchat](https://github.com/icon-gaurav/pgchat). If you're building something similar or have thoughts on any of these decisions, I'd like to hear them.