Client for PostgreSQL vector database operations with pgvector extension. Provides a unified interface for indexing, storing, and retrieving vector-based search results. Requirements: - PostgreSQL database with pgvector extension installed - A configured table with vector type
| 29 | logger = get_configured_logger("postgres_client") |
| 30 | |
| 31 | class PgVectorClient(RetrievalClientBase): |
| 32 | """ |
| 33 | Client for PostgreSQL vector database operations with pgvector extension. |
| 34 | Provides a unified interface for indexing, storing, and retrieving vector-based search results. |
| 35 | |
| 36 | Requirements: |
| 37 | - PostgreSQL database with pgvector extension installed |
| 38 | - A configured table with vector type column for embeddings |
| 39 | """ |
| 40 | |
| 41 | def __init__(self, endpoint_name: str | None = None): |
| 42 | super().__init__() # Initialize the base class with caching |
| 43 | """ |
| 44 | Initialize the PostgreSQL vector database client. |
| 45 | |
| 46 | Args: |
| 47 | endpoint_name: Name of the endpoint to use (defaults to preferred endpoint in CONFIG) |
| 48 | """ |
| 49 | self.endpoint_name = endpoint_name or CONFIG.write_endpoint |
| 50 | self._conn_lock = asyncio.Lock() |
| 51 | self._pool = None |
| 52 | self._pool_init_lock = asyncio.Lock() |
| 53 | |
| 54 | logger.info(f"Initializing PgVectorClient for endpoint: {self.endpoint_name}") |
| 55 | |
| 56 | # Get endpoint configuration |
| 57 | self.endpoint_config = self._get_endpoint_config() |
| 58 | self.api_endpoint = self.endpoint_config.api_endpoint |
| 59 | self.api_key = self.endpoint_config.api_key |
| 60 | self.database_path = self.endpoint_config.database_path |
| 61 | self.default_collection_name = self.endpoint_config.index_name or "nlweb_collection" |
| 62 | |
| 63 | self.pg_raw_config = self._get_config_from_postgres_connection_string(self.api_endpoint) |
| 64 | |
| 65 | self.host = self.pg_raw_config.get("host") |
| 66 | self.port = self.pg_raw_config.get("port", 5432) # Default PostgreSQL port |
| 67 | self.dbname = self.pg_raw_config.get("database") |
| 68 | self.username = self.pg_raw_config.get("username") |
| 69 | self.password = self.api_key or self.pg_raw_config.get("password") |
| 70 | self.table_name = self.default_collection_name or "documents" |
| 71 | |
| 72 | # Validate critical configuration |
| 73 | if not self.host: |
| 74 | error_msg = f"Missing 'host' in PostgreSQL configuration for endpoint '{self.endpoint_name}'" |
| 75 | logger.error(error_msg) |
| 76 | logger.error(f"Available configuration keys: {list(self.pg_raw_config.keys())}") |
| 77 | raise ValueError(error_msg) |
| 78 | if not self.dbname: |
| 79 | error_msg = f"Missing 'database_name' in PostgreSQL configuration for endpoint '{self.endpoint_name}'" |
| 80 | logger.error(error_msg) |
| 81 | raise ValueError(error_msg) |
| 82 | |
| 83 | logger.info(f"PostgreSQL client configured for table: {self.default_collection_name}") |
| 84 | |
| 85 | def _get_config_from_postgres_connection_string(self, connection_string: str) -> dict[str, Any]: |
| 86 | """ |
| 87 | Parse the PostgreSQL connection string and return a dictionary of configuration parameters. |
| 88 |
no outgoing calls
no test coverage detected