Skip to content

Commit 32871f7

Browse files
author
Jarosław Lach
committed
[Ai Bundle][Postgres] Add dbal connection configuration
[AI Bundle][Postgres] Add dbal connection configuration Fix code style errors
1 parent 26f7453 commit 32871f7

File tree

3 files changed

+74
-14
lines changed

3 files changed

+74
-14
lines changed

‎src/ai-bundle/config/options.php‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
use Symfony\AI\Platform\Bridge\OpenAi\PlatformFactory;
1818
use Symfony\AI\Platform\Capability;
1919
use Symfony\AI\Platform\PlatformInterface;
20+
use Symfony\AI\Store\Bridge\Postgres\Distance as PostgresDistance;
2021
use Symfony\AI\Store\Bridge\Redis\Distance;
2122
use Symfony\AI\Store\Document\VectorizerInterface;
2223
use Symfony\AI\Store\StoreInterface;
@@ -678,8 +679,19 @@
678679
->stringNode('table_name')->isRequired()->end()
679680
->stringNode('vector_field')->end()
680681
->enumNode('distance')
681-
->values(['cosine', 'inner_product', 'l1', 'l2'])
682+
->info('Distance metric to use for vector similarity search')
683+
->enumFqcn(PostgresDistance::class)
684+
->defaultValue(PostgresDistance::L2)
682685
->end()
686+
->stringNode('dbal_connection')->cannotBeEmpty()->end()
687+
->end()
688+
->validate()
689+
->ifTrue(static fn ($v) => !isset($v['dsn']) && !isset($v['dbal_connection']))
690+
->thenInvalid('Either "dsn" or "dbal_connection" must be configured.')
691+
->end()
692+
->validate()
693+
->ifTrue(static fn ($v) => isset($v['dsn'], $v['dbal_connection']))
694+
->thenInvalid('Either "dsn" or "dbal_connection" can be configured, but not both.')
683695
->end()
684696
->end()
685697
->end()

‎src/ai-bundle/src/AiBundle.php‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@
6969
use Symfony\AI\Store\Bridge\MongoDb\Store as MongoDbStore;
7070
use Symfony\AI\Store\Bridge\Neo4j\Store as Neo4jStore;
7171
use Symfony\AI\Store\Bridge\Pinecone\Store as PineconeStore;
72-
use Symfony\AI\Store\Bridge\Postgres\Store;
72+
use Symfony\AI\Store\Bridge\Postgres\Store as PostgresStore;
7373
use Symfony\AI\Store\Bridge\Qdrant\Store as QdrantStore;
7474
use Symfony\AI\Store\Bridge\Redis\Store as RedisStore;
7575
use Symfony\AI\Store\Bridge\SurrealDb\Store as SurrealDbStore;
@@ -1213,17 +1213,27 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
12131213

12141214
if ('postgres' === $type) {
12151215
foreach ($stores as $name => $store) {
1216-
$pdo = new Definition(\PDO::class);
1217-
$pdo->setArguments([
1218-
$store['dsn'],
1219-
$store['username'] ?? null,
1220-
$store['password'] ?? null],
1221-
);
1222-
1223-
$arguments = [
1224-
$pdo,
1225-
$store['table_name'],
1226-
];
1216+
$definition = new Definition(PostgresStore::class);
1217+
1218+
if (\array_key_exists('dbal_connection', $store)) {
1219+
$definition->setFactory([PostgresStore::class, 'fromDbal']);
1220+
$arguments = [
1221+
new Reference($store['dbal_connection']),
1222+
$store['table_name'],
1223+
];
1224+
} else {
1225+
$pdo = new Definition(\PDO::class);
1226+
$pdo->setArguments([
1227+
$store['dsn'],
1228+
$store['username'] ?? null,
1229+
$store['password'] ?? null],
1230+
);
1231+
1232+
$arguments = [
1233+
$pdo,
1234+
$store['table_name'],
1235+
];
1236+
}
12271237

12281238
if (\array_key_exists('vector_field', $store)) {
12291239
$arguments[2] = $store['vector_field'];
@@ -1233,7 +1243,6 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
12331243
$arguments[3] = $store['distance'];
12341244
}
12351245

1236-
$definition = new Definition(Store::class);
12371246
$definition
12381247
->addTag('ai.store')
12391248
->setArguments($arguments);

‎src/ai-bundle/tests/DependencyInjection/AiBundleTest.php‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,45 @@ public function testInMemoryStoreWithCustomStrategyCanBeConfigured()
366366
$this->assertSame('ai.store.distance_calculator.my_memory_store_with_custom_strategy', (string) $definition->getArgument(0));
367367
}
368368

369+
public function testPostgresStoreWithDifferentConnectionCanBeConfigured()
370+
{
371+
$container = $this->buildContainer([
372+
'ai' => [
373+
'store' => [
374+
'postgres' => [
375+
'db' => [
376+
'dsn' => 'pgsql:host=localhost;port=5432;dbname=testdb;user=app;password=mypass',
377+
'table_name' => 'vectors',
378+
],
379+
],
380+
],
381+
],
382+
]);
383+
384+
$this->assertTrue($container->hasDefinition('ai.store.postgres.db'));
385+
386+
$definition = $container->getDefinition('ai.store.postgres.db');
387+
$this->assertCount(3, $definition->getArguments());
388+
$this->assertInstanceOf(Definition::class, $definition->getArgument(0));
389+
390+
$container = $this->buildContainer([
391+
'ai' => [
392+
'store' => [
393+
'postgres' => [
394+
'db' => [
395+
'dbal_connection' => 'my_connection',
396+
'table_name' => 'vectors',
397+
],
398+
],
399+
],
400+
],
401+
]);
402+
403+
$definition = $container->getDefinition('ai.store.postgres.db');
404+
$this->assertCount(3, $definition->getArguments());
405+
$this->assertInstanceOf(Reference::class, $definition->getArgument(0));
406+
}
407+
369408
public function testConfigurationWithUseAttributeAsKeyWorksWithoutNormalizeKeys()
370409
{
371410
// Test that configurations using useAttributeAsKey work correctly

0 commit comments

Comments
 (0)