Skip to content

Commit 5f3b28d

Browse files
valtzuchr-hertel
valtzu
authored and
chr-hertel
committed
feat: add Gemini Embeddings (#347)
Add Embeddings support for Google Gemini. Only batch embedding implemented. Related to #28
1 parent c6e32b6 commit 5f3b28d

File tree

8 files changed

+343
-3
lines changed

8 files changed

+343
-3
lines changed

‎examples/google/embeddings.php‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
<?php
2+
3+
/*
4+
* This file is part of the Symfony package.
5+
*
6+
* (c) Fabien Potencier <fabien@symfony.com>
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 Symfony\AI\Platform\Bridge\Google\Embeddings;
13+
use Symfony\AI\Platform\Bridge\Google\PlatformFactory;
14+
use Symfony\AI\Platform\Response\VectorResponse;
15+
use Symfony\Component\Dotenv\Dotenv;
16+
17+
require_once dirname(__DIR__).'/vendor/autoload.php';
18+
(new Dotenv())->loadEnv(dirname(__DIR__).'/.env');
19+
20+
if (empty($_ENV['GOOGLE_API_KEY'])) {
21+
echo 'Please set the GOOGLE_API_KEY environment variable.'.\PHP_EOL;
22+
exit(1);
23+
}
24+
25+
$platform = PlatformFactory::create($_ENV['GOOGLE_API_KEY']);
26+
$embeddings = new Embeddings();
27+
28+
$response = $platform->request($embeddings, <<<TEXT
29+
Once upon a time, there was a country called Japan. It was a beautiful country with a lot of mountains and rivers.
30+
The people of Japan were very kind and hardworking. They loved their country very much and took care of it. The
31+
country was very peaceful and prosperous. The people lived happily ever after.
32+
TEXT);
33+
34+
assert($response instanceof VectorResponse);
35+
36+
echo 'Dimensions: '.$response->getContent()[0]->getDimensions().\PHP_EOL;

‎examples/store/mariadb-similarity-search-gemini.php‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
<?php
2+
3+
use Doctrine\DBAL\DriverManager;
4+
use Doctrine\DBAL\Tools\DsnParser;
5+
use PhpLlm\LlmChain\Chain\Chain;
6+
use PhpLlm\LlmChain\Chain\Toolbox\ChainProcessor;
7+
use PhpLlm\LlmChain\Chain\Toolbox\Tool\SimilaritySearch;
8+
use PhpLlm\LlmChain\Chain\Toolbox\Toolbox;
9+
use PhpLlm\LlmChain\Platform\Bridge\Google\Embeddings;
10+
use PhpLlm\LlmChain\Platform\Bridge\Google\Embeddings\TaskType;
11+
use PhpLlm\LlmChain\Platform\Bridge\Google\Gemini;
12+
use PhpLlm\LlmChain\Platform\Bridge\Google\PlatformFactory;
13+
use PhpLlm\LlmChain\Platform\Message\Message;
14+
use PhpLlm\LlmChain\Platform\Message\MessageBag;
15+
use PhpLlm\LlmChain\Store\Bridge\MariaDB\Store;
16+
use PhpLlm\LlmChain\Store\Document\Metadata;
17+
use PhpLlm\LlmChain\Store\Document\TextDocument;
18+
use PhpLlm\LlmChain\Store\Indexer;
19+
use Symfony\Component\Dotenv\Dotenv;
20+
use Symfony\Component\Uid\Uuid;
21+
22+
require_once dirname(__DIR__, 2).'/vendor/autoload.php';
23+
(new Dotenv())->loadEnv(dirname(__DIR__, 2).'/.env');
24+
25+
if (empty($_ENV['GOOGLE_API_KEY']) || empty($_ENV['MARIADB_URI'])) {
26+
echo 'Please set GOOGLE_API_KEY and MARIADB_URI environment variables.'.\PHP_EOL;
27+
exit(1);
28+
}
29+
30+
// initialize the store
31+
$store = Store::fromDbal(
32+
connection: DriverManager::getConnection((new DsnParser())->parse($_ENV['MARIADB_URI'])),
33+
tableName: 'my_table',
34+
indexName: 'my_index',
35+
vectorFieldName: 'embedding',
36+
);
37+
38+
// our data
39+
$movies = [
40+
['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'],
41+
['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'],
42+
['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'],
43+
];
44+
45+
// create embeddings and documents
46+
foreach ($movies as $i => $movie) {
47+
$documents[] = new TextDocument(
48+
id: Uuid::v4(),
49+
content: 'Title: '.$movie['title'].\PHP_EOL.'Director: '.$movie['director'].\PHP_EOL.'Description: '.$movie['description'],
50+
metadata: new Metadata($movie),
51+
);
52+
}
53+
54+
// initialize the table
55+
$store->initialize(['dimensions' => 768]);
56+
57+
// create embeddings for documents
58+
$platform = PlatformFactory::create($_ENV['GOOGLE_API_KEY']);
59+
$embeddings = new Embeddings(options: ['dimensions' => 768, 'task_type' => TaskType::SemanticSimilarity]);
60+
$indexer = new Indexer($platform, $embeddings, $store);
61+
$indexer->index($documents);
62+
63+
$model = new Gemini(Gemini::GEMINI_2_FLASH_LITE);
64+
65+
$similaritySearch = new SimilaritySearch($platform, $embeddings, $store);
66+
$toolbox = Toolbox::create($similaritySearch);
67+
$processor = new ChainProcessor($toolbox);
68+
$chain = new Chain($platform, $model, [$processor], [$processor]);
69+
70+
$messages = new MessageBag(
71+
Message::forSystem('Please answer all user questions only using SimilaritySearch function.'),
72+
Message::ofUser('Which movie fits the theme of the mafia?')
73+
);
74+
$response = $chain->call($messages);
75+
76+
echo $response->getContent().\PHP_EOL;

‎src/platform/src/Bridge/Google/Embeddings.php‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
<?php
2+
3+
namespace Symfony\AI\Platform\Bridge\Google;
4+
5+
use Symfony\AI\Platform\Bridge\Google\Embeddings\TaskType;
6+
use Symfony\AI\Platform\Capability;
7+
use Symfony\AI\Platform\Model;
8+
9+
/**
10+
* @author Valtteri R <valtzu@gmail.com>
11+
*/
12+
class Embeddings extends Model
13+
{
14+
/** Supported dimensions: 3072, 1536, or 768 */
15+
public const GEMINI_EMBEDDING_EXP_03_07 = 'gemini-embedding-exp-03-07';
16+
/** Fixed 768 dimensions */
17+
public const TEXT_EMBEDDING_004 = 'text-embedding-004';
18+
/** Fixed 768 dimensions */
19+
public const EMBEDDING_001 = 'embedding-001';
20+
21+
/**
22+
* @param array{dimensions?: int, task_type?: TaskType|string} $options
23+
*/
24+
public function __construct(string $name = self::GEMINI_EMBEDDING_EXP_03_07, array $options = [])
25+
{
26+
parent::__construct($name, [Capability::INPUT_MULTIPLE], $options);
27+
}
28+
}

‎src/platform/src/Bridge/Google/Embeddings/ModelClient.php‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
<?php
2+
3+
namespace Symfony\AI\Platform\Bridge\Google\Embeddings;
4+
5+
use Symfony\AI\Platform\Bridge\Google\Embeddings;
6+
use Symfony\AI\Platform\Exception\RuntimeException;
7+
use Symfony\AI\Platform\Model;
8+
use Symfony\AI\Platform\ModelClientInterface;
9+
use Symfony\AI\Platform\Response\VectorResponse;
10+
use Symfony\AI\Platform\ResponseConverterInterface;
11+
use Symfony\AI\Platform\Vector\Vector;
12+
use Symfony\Contracts\HttpClient\HttpClientInterface;
13+
use Symfony\Contracts\HttpClient\ResponseInterface;
14+
15+
/**
16+
* @author Valtteri R <valtzu@gmail.com>
17+
*/
18+
final readonly class ModelClient implements ModelClientInterface, ResponseConverterInterface
19+
{
20+
public function __construct(
21+
private HttpClientInterface $httpClient,
22+
#[\SensitiveParameter]
23+
private string $apiKey,
24+
) {
25+
}
26+
27+
public function supports(Model $model): bool
28+
{
29+
return $model instanceof Embeddings;
30+
}
31+
32+
public function request(Model $model, array|string $payload, array $options = []): ResponseInterface
33+
{
34+
$url = \sprintf('https://generativelanguage.googleapis.com/v1beta/models/%s:%s', $model->getName(), 'batchEmbedContents');
35+
$modelOptions = $model->getOptions();
36+
37+
return $this->httpClient->request('POST', $url, [
38+
'headers' => [
39+
'x-goog-api-key' => $this->apiKey,
40+
],
41+
'json' => [
42+
'requests' => array_map(
43+
static fn (string $text) => array_filter([
44+
'model' => 'models/'.$model->getName(),
45+
'content' => ['parts' => [['text' => $text]]],
46+
'outputDimensionality' => $modelOptions['dimensions'] ?? null,
47+
'taskType' => $modelOptions['task_type'] ?? null,
48+
'title' => $options['title'] ?? null,
49+
]),
50+
\is_array($payload) ? $payload : [$payload],
51+
),
52+
],
53+
]);
54+
}
55+
56+
public function convert(ResponseInterface $response, array $options = []): VectorResponse
57+
{
58+
$data = $response->toArray();
59+
60+
if (!isset($data['embeddings'])) {
61+
throw new RuntimeException('Response does not contain data');
62+
}
63+
64+
return new VectorResponse(
65+
...array_map(
66+
static fn (array $item): Vector => new Vector($item['values']),
67+
$data['embeddings'],
68+
),
69+
);
70+
}
71+
}

‎src/platform/src/Bridge/Google/Embeddings/TaskType.php‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
<?php
2+
3+
namespace Symfony\AI\Platform\Bridge\Google\Embeddings;
4+
5+
enum TaskType: string
6+
{
7+
/** Unset value, which will default to one of the other enum values. */
8+
public const TaskTypeUnspecified = 'TASK_TYPE_UNSPECIFIED';
9+
/** Specifies the given text is a query in a search/retrieval setting. */
10+
public const RetrievalQuery = 'RETRIEVAL_QUERY';
11+
/** Specifies the given text is a document from the corpus being searched. */
12+
public const RetrievalDocument = 'RETRIEVAL_DOCUMENT';
13+
/** Specifies the given text will be used for STS. */
14+
public const SemanticSimilarity = 'SEMANTIC_SIMILARITY';
15+
/** Specifies that the given text will be classified. */
16+
public const Classification = 'CLASSIFICATION';
17+
/** Specifies that the embeddings will be used for clustering. */
18+
public const Clustering = 'CLUSTERING';
19+
/** Specifies that the given text will be used for question answering. */
20+
public const QuestionAnswering = 'QUESTION_ANSWERING';
21+
/** Specifies that the given text will be used for fact verification. */
22+
public const FactVerification = 'FACT_VERIFICATION';
23+
/** Specifies that the given text will be used for code retrieval. */
24+
public const CodeRetrievalQuery = 'CODE_RETRIEVAL_QUERY';
25+
}

‎src/platform/src/Bridge/Google/PlatformFactory.php‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
use Symfony\AI\Platform\Bridge\Google\Contract\ToolCallMessageNormalizer;
1717
use Symfony\AI\Platform\Bridge\Google\Contract\ToolNormalizer;
1818
use Symfony\AI\Platform\Bridge\Google\Contract\UserMessageNormalizer;
19+
use Symfony\AI\Platform\Bridge\Google\Embeddings\ModelClient;
1920
use Symfony\AI\Platform\Contract;
2021
use Symfony\AI\Platform\Platform;
2122
use Symfony\Component\HttpClient\EventSourceHttpClient;
@@ -33,8 +34,9 @@ public static function create(
3334
): Platform {
3435
$httpClient = $httpClient instanceof EventSourceHttpClient ? $httpClient : new EventSourceHttpClient($httpClient);
3536
$responseHandler = new ModelHandler($httpClient, $apiKey);
37+
$embeddings = new ModelClient($httpClient, $apiKey);
3638

37-
return new Platform([$responseHandler], [$responseHandler], Contract::create(
39+
return new Platform([$responseHandler, $embeddings], [$responseHandler, $embeddings], Contract::create(
3840
new AssistantMessageNormalizer(),
3941
new MessageBagNormalizer(),
4042
new ToolNormalizer(),

‎src/platform/src/ModelClientInterface.php‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ interface ModelClientInterface
2121
public function supports(Model $model): bool;
2222

2323
/**
24-
* @param array<string, mixed> $payload
25-
* @param array<string, mixed> $options
24+
* @param array<string|int, mixed> $payload
25+
* @param array<string, mixed> $options
2626
*/
2727
public function request(Model $model, array|string $payload, array $options = []): ResponseInterface;
2828
}

‎tests/Platform/Bridge/Google/Embeddings/EmbeddingsModelClientTest.php‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
<?php
2+
3+
declare(strict_types=1);
4+
5+
namespace PhpLlm\LlmChain\Tests\Platform\Bridge\Google\Embeddings;
6+
7+
use PhpLlm\LlmChain\Platform\Bridge\Google\Embeddings;
8+
use PhpLlm\LlmChain\Platform\Bridge\Google\Embeddings\ModelClient;
9+
use PhpLlm\LlmChain\Platform\Response\VectorResponse;
10+
use PhpLlm\LlmChain\Platform\Vector\Vector;
11+
use PHPUnit\Framework\Attributes\CoversClass;
12+
use PHPUnit\Framework\Attributes\Small;
13+
use PHPUnit\Framework\Attributes\Test;
14+
use PHPUnit\Framework\Attributes\UsesClass;
15+
use PHPUnit\Framework\TestCase;
16+
use Symfony\Contracts\HttpClient\HttpClientInterface;
17+
use Symfony\Contracts\HttpClient\ResponseInterface;
18+
19+
#[CoversClass(ModelClient::class)]
20+
#[Small]
21+
#[UsesClass(Vector::class)]
22+
#[UsesClass(VectorResponse::class)]
23+
#[UsesClass(Embeddings::class)]
24+
final class EmbeddingsModelClientTest extends TestCase
25+
{
26+
#[Test]
27+
public function itMakesARequestWithCorrectPayload(): void
28+
{
29+
$response = $this->createStub(ResponseInterface::class);
30+
$response
31+
->method('toArray')
32+
->willReturn(json_decode($this->getEmbeddingStub(), true));
33+
34+
$httpClient = self::createMock(HttpClientInterface::class);
35+
$httpClient->expects(self::once())
36+
->method('request')
37+
->with(
38+
'POST',
39+
'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-exp-03-07:batchEmbedContents',
40+
[
41+
'headers' => ['x-goog-api-key' => 'test'],
42+
'json' => [
43+
'requests' => [
44+
[
45+
'model' => 'models/gemini-embedding-exp-03-07',
46+
'content' => ['parts' => [['text' => 'payload1']]],
47+
'outputDimensionality' => 1536,
48+
'taskType' => 'CLASSIFICATION',
49+
],
50+
[
51+
'model' => 'models/gemini-embedding-exp-03-07',
52+
'content' => ['parts' => [['text' => 'payload2']]],
53+
'outputDimensionality' => 1536,
54+
'taskType' => 'CLASSIFICATION',
55+
],
56+
],
57+
],
58+
],
59+
)
60+
->willReturn($response);
61+
62+
$model = new Embeddings(Embeddings::GEMINI_EMBEDDING_EXP_03_07, ['dimensions' => 1536, 'task_type' => 'CLASSIFICATION']);
63+
64+
$httpResponse = (new ModelClient($httpClient, 'test'))->request($model, ['payload1', 'payload2']);
65+
self::assertSame(json_decode($this->getEmbeddingStub(), true), $httpResponse->toArray());
66+
}
67+
68+
#[Test]
69+
public function itConvertsAResponseToAVectorResponse(): void
70+
{
71+
$response = $this->createStub(ResponseInterface::class);
72+
$response
73+
->method('toArray')
74+
->willReturn(json_decode($this->getEmbeddingStub(), true));
75+
76+
$httpClient = self::createMock(HttpClientInterface::class);
77+
78+
$vectorResponse = (new ModelClient($httpClient, 'test'))->convert($response);
79+
$convertedContent = $vectorResponse->getContent();
80+
81+
self::assertCount(2, $convertedContent);
82+
83+
self::assertSame([0.3, 0.4, 0.4], $convertedContent[0]->getData());
84+
self::assertSame([0.0, 0.0, 0.2], $convertedContent[1]->getData());
85+
}
86+
87+
private function getEmbeddingStub(): string
88+
{
89+
return <<<'JSON'
90+
{
91+
"embeddings": [
92+
{
93+
"values": [0.3, 0.4, 0.4]
94+
},
95+
{
96+
"values": [0.0, 0.0, 0.2]
97+
}
98+
]
99+
}
100+
JSON;
101+
}
102+
}

0 commit comments

Comments
 (0)