Skip to content

Commit 4807521

Browse files
chr-hertelchr-hertel
committed
feature #44 feat: add Postgres store (smnandre)
This PR was squashed before being merged into the main branch. Discussion ---------- feat: add Postgres store | Q | A | ------------- | --- | Bug fix? | no | New feature? | yes | Docs? | no | Issues | Fix #... | License | MIT Here is a first/naive attempt to add a Postgres store (via pgvector extension) Very inspired by the [MariaDB store ](#39) PR (and existing examples). Trying to digg in the repository, please tell me if anything is not in the spirit/style of what you guys building here.. Commits ------- 87f2ddb feat: add Postgres store
2 parents c0122f5 + 87f2ddb commit 4807521

File tree

4 files changed

+257
-0
lines changed

4 files changed

+257
-0
lines changed

‎compose.yaml

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,3 +6,12 @@ services:
66
MARIADB_DATABASE: my_database
77
ports:
88
- "3309:3306"
9+
10+
postgres:
11+
image: ankane/pgvector
12+
environment:
13+
POSTGRES_DB: my_database
14+
POSTGRES_USER: postgres
15+
POSTGRES_PASSWORD: postgres
16+
ports:
17+
- "5432:5432"

‎examples/.env

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ MONGODB_URI=
5959
PINECONE_API_KEY=
6060
PINECONE_HOST=
6161

62+
# For using Postgres (store)
63+
POSTGRES_URI=pdo-pgsql://postgres:[email protected]:5432/my_database
64+
6265
# Some examples are expensive to run, so we disable them by default
6366
RUN_EXPENSIVE_EXAMPLES=false
6467

‎examples/store/postgres-similarity-search.php

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
use Doctrine\DBAL\DriverManager;
13+
use Doctrine\DBAL\Tools\DsnParser;
14+
use Symfony\AI\Agent\Agent;
15+
use Symfony\AI\Agent\Toolbox\AgentProcessor;
16+
use Symfony\AI\Agent\Toolbox\Tool\SimilaritySearch;
17+
use Symfony\AI\Agent\Toolbox\Toolbox;
18+
use Symfony\AI\Platform\Bridge\OpenAI\Embeddings;
19+
use Symfony\AI\Platform\Bridge\OpenAI\GPT;
20+
use Symfony\AI\Platform\Bridge\OpenAI\PlatformFactory;
21+
use Symfony\AI\Platform\Message\Message;
22+
use Symfony\AI\Platform\Message\MessageBag;
23+
use Symfony\AI\Store\Bridge\Postgres\Store;
24+
use Symfony\AI\Store\Document\Metadata;
25+
use Symfony\AI\Store\Document\TextDocument;
26+
use Symfony\AI\Store\Indexer;
27+
use Symfony\Component\Dotenv\Dotenv;
28+
use Symfony\Component\Uid\Uuid;
29+
30+
require_once dirname(__DIR__).'/vendor/autoload.php';
31+
(new Dotenv())->loadEnv(dirname(__DIR__).'/.env');
32+
33+
if (empty($_ENV['OPENAI_API_KEY']) || empty($_ENV['POSTGRES_URI'])) {
34+
echo 'Please set OPENAI_API_KEY and POSTGRES_URI environment variables.'.\PHP_EOL;
35+
exit(1);
36+
}
37+
38+
// initialize the store
39+
$store = Store::fromDbal(
40+
connection: DriverManager::getConnection((new DsnParser())->parse($_ENV['POSTGRES_URI'])),
41+
tableName: 'my_table',
42+
vectorFieldName: 'embedding',
43+
);
44+
45+
// our data
46+
$movies = [
47+
['title' => 'Inception', 'description' => 'A skilled thief is given a chance at redemption if he can successfully perform inception, the act of planting an idea in someone\'s subconscious.', 'director' => 'Christopher Nolan'],
48+
['title' => 'The Matrix', 'description' => 'A hacker discovers the world he lives in is a simulated reality and joins a rebellion to overthrow its controllers.', 'director' => 'The Wachowskis'],
49+
['title' => 'The Godfather', 'description' => 'The aging patriarch of an organized crime dynasty transfers control of his empire to his reluctant son.', 'director' => 'Francis Ford Coppola'],
50+
];
51+
52+
// create embeddings and documents
53+
$documents = [];
54+
foreach ($movies as $i => $movie) {
55+
$documents[] = new TextDocument(
56+
id: Uuid::v4(),
57+
content: 'Title: '.$movie['title'].\PHP_EOL.'Director: '.$movie['director'].\PHP_EOL.'Description: '.$movie['description'],
58+
metadata: new Metadata($movie),
59+
);
60+
}
61+
62+
// initialize the table
63+
$store->initialize();
64+
65+
// create embeddings for documents
66+
$platform = PlatformFactory::create($_ENV['OPENAI_API_KEY']);
67+
$indexer = new Indexer($platform, $embeddings = new Embeddings(), $store);
68+
$indexer->index($documents);
69+
70+
$model = new GPT(GPT::GPT_4O_MINI);
71+
72+
$similaritySearch = new SimilaritySearch($platform, $embeddings, $store);
73+
$toolbox = Toolbox::create($similaritySearch);
74+
$processor = new AgentProcessor($toolbox);
75+
$agent = new Agent($platform, $model, [$processor], [$processor]);
76+
77+
$messages = new MessageBag(
78+
Message::forSystem('Please answer all user questions only using SimilaritySearch function.'),
79+
Message::ofUser('Which movie fits the theme of technology?')
80+
);
81+
$response = $agent->call($messages);
82+
83+
echo $response->getContent().\PHP_EOL;

‎src/store/src/Bridge/Postgres/Store.php

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <[email protected]>
7+
*
8+
* For the full copyright and license information, please view the LICENSE
9+
* file that was distributed with this source code.
10+
*/
11+
12+
namespace Symfony\AI\Store\Bridge\Postgres;
13+
14+
use Doctrine\DBAL\Connection;
15+
use Symfony\AI\Platform\Vector\Vector;
16+
use Symfony\AI\Platform\Vector\VectorInterface;
17+
use Symfony\AI\Store\Document\Metadata;
18+
use Symfony\AI\Store\Document\VectorDocument;
19+
use Symfony\AI\Store\Exception\InvalidArgumentException;
20+
use Symfony\AI\Store\InitializableStoreInterface;
21+
use Symfony\AI\Store\VectorStoreInterface;
22+
use Symfony\Component\Uid\Uuid;
23+
24+
/**
25+
* Requires PostgreSQL with pgvector extension.
26+
*
27+
* @author Simon André <[email protected]>
28+
*
29+
* @see https://github.com/pgvector/pgvector
30+
*/
31+
final readonly class Store implements VectorStoreInterface, InitializableStoreInterface
32+
{
33+
public function __construct(
34+
private \PDO $connection,
35+
private string $tableName,
36+
private string $vectorFieldName = 'embedding',
37+
) {
38+
}
39+
40+
public static function fromPdo(\PDO $connection, string $tableName, string $vectorFieldName = 'embedding'): self
41+
{
42+
return new self($connection, $tableName, $vectorFieldName);
43+
}
44+
45+
public static function fromDbal(Connection $connection, string $tableName, string $vectorFieldName = 'embedding'): self
46+
{
47+
$pdo = $connection->getNativeConnection();
48+
49+
if (!$pdo instanceof \PDO) {
50+
throw new InvalidArgumentException('Only DBAL connections using PDO driver are supported.');
51+
}
52+
53+
return self::fromPdo($pdo, $tableName, $vectorFieldName);
54+
}
55+
56+
public function add(VectorDocument ...$documents): void
57+
{
58+
$statement = $this->connection->prepare(
59+
\sprintf(
60+
'INSERT INTO %1$s (id, metadata, %2$s)
61+
VALUES (:id, :metadata, :vector)
62+
ON CONFLICT (id) DO UPDATE SET metadata = EXCLUDED.metadata, %2$s = EXCLUDED.%2$s',
63+
$this->tableName,
64+
$this->vectorFieldName,
65+
),
66+
);
67+
68+
foreach ($documents as $document) {
69+
$operation = [
70+
'id' => $document->id->toRfc4122(),
71+
'metadata' => json_encode($document->metadata->getArrayCopy(), \JSON_THROW_ON_ERROR),
72+
'vector' => $this->toPgvector($document->vector),
73+
];
74+
75+
$statement->execute($operation);
76+
}
77+
}
78+
79+
/**
80+
* @param array<string, mixed> $options
81+
* @param float|null $minScore Minimum score to filter results (optional)
82+
*
83+
* @return VectorDocument[]
84+
*/
85+
public function query(Vector $vector, array $options = [], ?float $minScore = null): array
86+
{
87+
$sql = \sprintf(
88+
'SELECT id, %s AS embedding, metadata, (%s <-> :embedding) AS score
89+
FROM %s
90+
%s
91+
ORDER BY score ASC
92+
LIMIT %d',
93+
$this->vectorFieldName,
94+
$this->vectorFieldName,
95+
$this->tableName,
96+
null !== $minScore ? "WHERE ({$this->vectorFieldName} <-> :embedding) >= :minScore" : '',
97+
$options['limit'] ?? 5,
98+
);
99+
$statement = $this->connection->prepare($sql);
100+
101+
$params = [
102+
'embedding' => $this->toPgvector($vector),
103+
];
104+
if (null !== $minScore) {
105+
$params['minScore'] = $minScore;
106+
}
107+
108+
$statement->execute($params);
109+
110+
$documents = [];
111+
foreach ($statement->fetchAll(\PDO::FETCH_ASSOC) as $result) {
112+
$documents[] = new VectorDocument(
113+
id: Uuid::fromString($result['id']),
114+
vector: new Vector($this->fromPgvector($result['embedding'])),
115+
metadata: new Metadata(json_decode($result['metadata'] ?? '{}', true, 512, \JSON_THROW_ON_ERROR)),
116+
score: $result['score'],
117+
);
118+
}
119+
120+
return $documents;
121+
}
122+
123+
public function initialize(array $options = []): void
124+
{
125+
if ([] !== $options) {
126+
throw new InvalidArgumentException('No supported options');
127+
}
128+
129+
$this->connection->exec('CREATE EXTENSION IF NOT EXISTS vector');
130+
131+
$this->connection->exec(
132+
\sprintf(
133+
'CREATE TABLE IF NOT EXISTS %s (
134+
id UUID PRIMARY KEY,
135+
metadata JSONB,
136+
%s vector(1536) NOT NULL
137+
)',
138+
$this->tableName,
139+
$this->vectorFieldName,
140+
),
141+
);
142+
$this->connection->exec(
143+
\sprintf(
144+
'CREATE INDEX IF NOT EXISTS %s_%s_idx ON %s USING ivfflat (%s vector_cosine_ops)',
145+
$this->tableName,
146+
$this->vectorFieldName,
147+
$this->tableName,
148+
$this->vectorFieldName,
149+
),
150+
);
151+
}
152+
153+
private function toPgvector(VectorInterface $vector): string
154+
{
155+
return '['.implode(',', $vector->getData()).']';
156+
}
157+
158+
private function fromPgvector(string $vector): array
159+
{
160+
return json_decode($vector, true, 512, \JSON_THROW_ON_ERROR);
161+
}
162+
}

0 commit comments

Comments
 (0)