chat.py

##########################################################################
#
# pgAdmin 4 - PostgreSQL Tools
#
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
# This software is released under the PostgreSQL Licence
#
##########################################################################

"""LLM chat functionality with database tool integration.

This module provides high-level functions for running LLM conversations
that can use database tools to query and inspect PostgreSQL databases.
"""

import json
from collections.abc import Generator
from typing import Optional, Union

from pgadmin.llm.client import get_llm_client, is_llm_available
from pgadmin.llm.models import Message, LLMResponse, StopReason
from pgadmin.llm.tools import DATABASE_TOOLS, execute_tool, DatabaseToolError
from pgadmin.llm.utils import get_max_tool_iterations


# Default system prompt for database assistant
DEFAULT_SYSTEM_PROMPT = (
    "You are a PostgreSQL database assistant integrated into pgAdmin 4. "
    "You have access to tools that allow you to query the database and "
    "inspect its schema.\n\n"
    "When helping users:\n"
    "1. First understand the database structure using get_database_schema "
    "or get_table_info\n"
    "2. Write efficient SQL queries to answer questions about the data\n"
    "3. Explain your findings clearly and concisely\n"
    "4. If a query might return many rows, consider using LIMIT or "
    "aggregations\n\n"
    "Important:\n"
    "- All queries run in READ ONLY mode - you cannot modify data\n"
    "- Results are limited to 1000 rows\n"
    "- Always validate your understanding of the schema before writing "
    "complex queries"
)


def chat_with_database(
    user_message: str,
    sid: int,
    did: int,
    conversation_history: Optional[list[Message]] = None,
    system_prompt: Optional[str] = None,
    max_tool_iterations: Optional[int] = None,
    provider: Optional[str] = None,
    model: Optional[str] = None
) -> tuple[str, list[Message]]:
    """
    Run an LLM chat conversation with database tool access.

    This function handles the full conversation loop, executing any
    tool calls the LLM makes and continuing until a final response
    is generated.

    Args:
        user_message: The user's message/question
        sid: Server ID for database connection
        did: Database ID for database connection
        conversation_history: Optional list of previous messages
        system_prompt: Optional custom system prompt (uses default if None)
        max_tool_iterations: Maximum number of tool call
            rounds. Uses preference setting if None.
        provider: Optional LLM provider override
        model: Optional model override

    Returns:
        Tuple of (final_response_text, updated_conversation_history)

    Raises:
        LLMClientError: If the LLM request fails
        RuntimeError: If LLM is not available or max iterations exceeded
    """
    if not is_llm_available():
        raise RuntimeError("LLM is not configured. Please configure an LLM "
                           "provider in Preferences > AI.")

    client = get_llm_client(provider=provider, model=model)
    if not client:
        raise RuntimeError("Failed to create LLM client")

    # Initialize conversation history
    messages = list(conversation_history) if conversation_history else []
    messages.append(Message.user(user_message))

    # Use default system prompt if none provided
    if system_prompt is None:
        system_prompt = DEFAULT_SYSTEM_PROMPT

    # Get max iterations from preferences if not specified
    if max_tool_iterations is None:
        max_tool_iterations = get_max_tool_iterations()

    iteration = 0
    while iteration < max_tool_iterations:
        iteration += 1

        # Call the LLM
        response = client.chat(
            messages=messages,
            tools=DATABASE_TOOLS,
            system_prompt=system_prompt
        )

        # Add assistant response to history
        messages.append(response.to_message())

        # Check if we're done
        if response.stop_reason != StopReason.TOOL_USE:
            return response.content, messages

        # Execute tool calls
        tool_results = []
        for tool_call in response.tool_calls:
            try:
                result = execute_tool(
                    tool_name=tool_call.name,
                    arguments=tool_call.arguments,
                    sid=sid,
                    did=did
                )
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps(result, default=str),
                    is_error=False
                ))
            except (DatabaseToolError, ValueError) as e:
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps({"error": str(e)}),
                    is_error=True
                ))
            except Exception as e:
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps({
                        "error": f"Unexpected error: {str(e)}"
                    }),
                    is_error=True
                ))

        # Add tool results to history
        messages.extend(tool_results)

    raise RuntimeError(
        f"Exceeded maximum tool iterations ({max_tool_iterations})"
    )


def chat_with_database_stream(
    user_message: str,
    sid: int,
    did: int,
    conversation_history: Optional[list[Message]] = None,
    system_prompt: Optional[str] = None,
    max_tool_iterations: Optional[int] = None,
    provider: Optional[str] = None,
    model: Optional[str] = None
) -> Generator[Union[str, tuple], None, None]:
    """
    Stream an LLM chat conversation with database tool access.

    Like chat_with_database, but yields text chunks as the final
    response streams in. During tool-use iterations, no text is
    yielded (tools are executed silently).

    Yields:
        str: Text content chunks from the final LLM response.

    The last item yielded is a 3-tuple of
    ('complete', final_response_text, updated_conversation_history).

    Raises:
        LLMClientError: If the LLM request fails.
        RuntimeError: If LLM is not available or max iterations exceeded.
    """
    if not is_llm_available():
        raise RuntimeError("LLM is not configured. Please configure an LLM "
                           "provider in Preferences > AI.")

    client = get_llm_client(provider=provider, model=model)
    if not client:
        raise RuntimeError("Failed to create LLM client")

    messages = list(conversation_history) if conversation_history else []
    messages.append(Message.user(user_message))

    if system_prompt is None:
        system_prompt = DEFAULT_SYSTEM_PROMPT

    if max_tool_iterations is None:
        max_tool_iterations = get_max_tool_iterations()

    iteration = 0
    while iteration < max_tool_iterations:
        iteration += 1

        # Stream the LLM response, yielding text chunks as they arrive
        response = None
        for item in client.chat_stream(
            messages=messages,
            tools=DATABASE_TOOLS,
            system_prompt=system_prompt
        ):
            if isinstance(item, LLMResponse):
                response = item
            elif isinstance(item, str):
                yield item

        if response is None:
            raise RuntimeError("No response received from LLM")

        messages.append(response.to_message())

        if response.stop_reason != StopReason.TOOL_USE:
            # Final response - yield a 3-tuple to distinguish from
            # the 2-tuple tool_use event
            yield ('complete', response.content, messages)
            return

        # Signal that tools are being executed so the caller can
        # reset streaming state and show a thinking indicator
        yield ('tool_use', [tc.name for tc in response.tool_calls])

        # Execute tool calls
        tool_results = []
        for tool_call in response.tool_calls:
            try:
                result = execute_tool(
                    tool_name=tool_call.name,
                    arguments=tool_call.arguments,
                    sid=sid,
                    did=did
                )
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps(result, default=str),
                    is_error=False
                ))
            except (DatabaseToolError, ValueError) as e:
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps({"error": str(e)}),
                    is_error=True
                ))
            except Exception as e:
                tool_results.append(Message.tool_result(
                    tool_call_id=tool_call.id,
                    content=json.dumps({
                        "error": f"Unexpected error: {str(e)}"
                    }),
                    is_error=True
                ))

        messages.extend(tool_results)

    raise RuntimeError(
        f"Exceeded maximum tool iterations ({max_tool_iterations})"
    )


def single_query(
    question: str,
    sid: int,
    did: int,
    provider: Optional[str] = None,
    model: Optional[str] = None
) -> str:
    """
    Ask a single question about the database.

    This is a convenience function for one-shot questions without
    maintaining conversation history.

    Args:
        question: The question to ask
        sid: Server ID
        did: Database ID
        provider: Optional LLM provider override
        model: Optional model override

    Returns:
        The LLM's response text

    Raises:
        LLMClientError: If the LLM request fails
        RuntimeError: If LLM is not available
    """
    response, _ = chat_with_database(
        user_message=question,
        sid=sid,
        did=did,
        provider=provider,
        model=model
    )
    return response