Blame for utils/postgres.py

"""
PostgreSQL client for Pipecat Server.

Provides direct database access for:
- Loading agent configuration at call start
- Writing call logs after call completion
"""

import os
import json
from typing import Optional
import asyncpg
from utils.logging import log_info, log_error, log_warning


class PostgresClient:
    """Async PostgreSQL client using asyncpg connection pool."""

    _pool: Optional[asyncpg.Pool] = None

    @classmethod
    async def get_pool(cls) -> asyncpg.Pool:
        """Get or create connection pool."""
        if cls._pool is None:
            database_url = os.getenv("DATABASE_URL")
            if not database_url:
                raise ValueError("DATABASE_URL environment variable not set")

            cls._pool = await asyncpg.create_pool(
                database_url,
                min_size=2,
                max_size=10,
                command_timeout=30,
            )
            log_info("postgres", "Connection pool created")

        return cls._pool

    @classmethod
    async def close(cls) -> None:
        """Close connection pool."""
        if cls._pool:
            await cls._pool.close()
            cls._pool = None
            log_info("postgres", "Connection pool closed")

    @classmethod
    async def get_agent(cls, agent_id: str) -> Optional[dict]:
        """
        Fetch ServiceConfiguration with services and botData by ID.

        Returns dict with:
        - id, companyName, voiceId, model, maxDuration, language, botData
        - services: list of {id, name, description, keywords, order}
        """
        pool = await cls.get_pool()

        async with pool.acquire() as conn:
            # Get configuration including botData
            config_row = await conn.fetchrow(
                '''
                SELECT id, "companyName", "voiceId", model, "maxDuration", language, "botData"
                FROM "ServiceConfiguration"
                WHERE id = $1
                ''',
                agent_id,
            )

            if not config_row:
                log_error("postgres", f"Agent not found: {agent_id}")
                return None

            # Get services
            service_rows = await conn.fetch(
                '''
                SELECT id, name, description, keywords, "order"
                FROM "Service"
                WHERE "configurationId" = $1
                ORDER BY "order" ASC
                ''',
                agent_id,
            )

            config = dict(config_row)
            config["services"] = [dict(row) for row in service_rows]

            # Parse botData if it's a string
            if config.get("botData") and isinstance(config["botData"], str):
                config["botData"] = json.loads(config["botData"])

            log_info("postgres", f"Loaded agent: {config['companyName']} ({len(config['services'])} services, botData={'yes' if config.get('botData') else 'no'})")
            return config

    @classmethod
    async def get_bot_data(cls, agent_id: str) -> Optional[dict]:
        """
        Load botData JSON from ServiceConfiguration.

        Returns the botData dict if found, None otherwise.
        Falls back to legacy generation if botData is not set.
        """
        agent_config = await cls.get_agent(agent_id)
        if not agent_config:
            return None

        bot_data = agent_config.get("botData")
        if bot_data:
            log_info("postgres", f"Loaded botData for agent: {agent_id}")
            return bot_data

        # Fallback: generate from agent config (legacy mode)
        log_warning("postgres", f"No botData in database for {agent_id}, generating legacy config")
        return cls._generate_bot_data_legacy(agent_config)

    @staticmethod
    def _generate_bot_data_legacy(agent_config: dict) -> dict:
        """
        Legacy BotData generation from agent config.

        DEPRECATED: Ensure botData is generated by Next.js API.
        """
        company_name = agent_config.get("companyName", "Компанія")
        services = agent_config.get("services", [])

        services_list = "\n".join(
            f"{i + 1}. {s['name']} - {s['description']}"
            for i, s in enumerate(sorted(services, key=lambda x: x.get("order", 0)))
        )
        service_ids = [s["id"] for s in services]

        return {
            "flow": {
                "initial_node": "greeting",
                "nodes": {
                    "greeting": {
                        "role_messages": [{"role": "system", "content": f'Ти - голосовий асистент компанії "{company_name}". Говори українською мовою.'}],
                        "task_messages": [{"role": "system", "content": f'Привітай клієнта від імені компанії "{company_name}" і запропонуй обрати послугу:\n{services_list}'}],
                        "functions": [
                            {"type": "function", "function": {"name": "select_service", "description": "Клієнт обрав послугу", "parameters": {"type": "object", "properties": {"service_id": {"type": "string", "enum": service_ids}}, "required": ["service_id"]}, "transition_to": "service_details"}},
                            {"type": "function", "function": {"name": "end_call", "description": "Завершити розмову", "parameters": {"type": "object", "properties": {}}, "transition_to": "goodbye"}},
                        ],
                    },
                    "service_details": {
                        "role_messages": [{"role": "system", "content": f'Ти - голосовий асистент компанії "{company_name}". Надай інформацію про послугу.'}],
                        "task_messages": [{"role": "system", "content": "Надай інформацію про обрану послугу."}],
                        "functions": [
                            {"type": "function", "function": {"name": "back_to_menu", "description": "Повернутися до меню", "parameters": {"type": "object", "properties": {}}, "transition_to": "greeting"}},
                            {"type": "function", "function": {"name": "end_call", "description": "Завершити розмову", "parameters": {"type": "object", "properties": {}}, "transition_to": "goodbye"}},
                        ],
                    },
                    "goodbye": {
                        "role_messages": [{"role": "system", "content": f'Попрощайся з клієнтом від імені "{company_name}".'}],
                        "task_messages": [{"role": "system", "content": "Подякуй і побажай гарного дня."}],
                    },
                },
            },
            "config": {
                "voice_id": agent_config.get("voiceId") or "79a125e8-cd45-4c13-8a67-188112f4dd22",
                "model": agent_config.get("model") or "gpt-4.1-mini",
            },
            "metadata": {
                "max_duration": agent_config.get("maxDuration") or 3600,
                "environment": "production",
            },
        }

    @classmethod
    async def create_call_log(
        cls,
        configuration_id: str,
        caller_number: str,
        call_sid: Optional[str] = None,
        stream_sid: Optional[str] = None,
        pipecat_session_id: Optional[str] = None,
        duration: int = 0,
        status: str = "COMPLETED",
        transcript: Optional[dict] = None,
    ) -> str:
        """
        Create a new call log entry.

        Args:
            pipecat_session_id: Pipecat Cloud session ID for fetching logs via API

        Returns the created log ID.
        """
        pool = await cls.get_pool()

        # Generate cuid-like ID
        import time
        import random
        import string

        timestamp = hex(int(time.time() * 1000))[2:]
        random_part = "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
        log_id = f"cl{timestamp}{random_part}"

        async with pool.acquire() as conn:
            await conn.execute(
                '''
                INSERT INTO "CallLog" (
                    id, "configurationId", "callerNumber", "callSid", "streamSid",
                    "pipecatSessionId", duration, status, transcript, "startedAt", "endedAt", "createdAt"
                )
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8::\"CallStatus\", $9, NOW(), NOW(), NOW())
                ''',
                log_id,
                configuration_id,
                caller_number,
                call_sid,
                stream_sid,
                pipecat_session_id,
                duration,
                status,
                json.dumps(transcript) if transcript else None,
            )

            log_info("postgres", f"Call log created: {log_id} [status={status}, session={pipecat_session_id}]")
            return log_id

    @classmethod
    async def update_call_log(
        cls,
        log_id: str,
        duration: Optional[int] = None,
        status: Optional[str] = None,
        transcript: Optional[dict] = None,
    ) -> None:
        """Update an existing call log."""
        pool = await cls.get_pool()

        updates = []
        params = [log_id]
        param_idx = 2

        if duration is not None:
            updates.append(f"duration = ${param_idx}")
            params.append(duration)
            param_idx += 1

        if status is not None:
            updates.append(f'status = ${param_idx}::"CallStatus"')
            params.append(status)
            param_idx += 1

        if transcript is not None:
            updates.append(f"transcript = ${param_idx}")
            params.append(json.dumps(transcript))
            param_idx += 1

        updates.append('"endedAt" = NOW()')

        if not updates:
            return

        query = f'UPDATE "CallLog" SET {", ".join(updates)} WHERE id = $1'

        async with pool.acquire() as conn:
            await conn.execute(query, *params)

        log_info("postgres", f"Call log updated: {log_id}")


# Module-level instance for convenience
postgres = PostgresClient()