feature #765 [AiBundle] Improve support for TraceableMessageStore and TraceableChat (Guikingone)

This PR was merged into the main branch.

Discussion
----------

[AiBundle] Improve support for `TraceableMessageStore` and `TraceableChat`

| Q             | A
| ------------- | ---
| Bug fix?      | no
| New feature?  | yes
| Docs?         | yes
| Issues        | Fix none
| License       | MIT

Commits
-------

eaaaf44 feat(aibundle): improve support for traceable*

Author: chr-hertel

docs/bundles/ai-bundle.rst

@@ -950,6 +950,48 @@ The profiler panel provides insights into the agent's execution:
 .. image:: profiler.png
    :alt: Profiler Panel
 
+Message stores
+--------------
+
+Message stores are critical to store messages sent to agents in the short / long term, they can be configured
+and reused in multiple chats, providing the capacity to agents to keep previous interactions.
+
+Configuring message stores
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Message stores are defined in the ``message_store`` section of your configuration:
+
+.. code-block:: yaml
+
+    ai:
+        # ...
+        message_store:
+            youtube:
+                cache:
+                    service: 'cache.app'
+                    key: 'youtube'
+
+Chats
+-----
+
+Chats are the entrypoint when it comes to sending messages to agents and retrieving content (mostly text)
+that contains the response from the agent.
+
+Each chat requires to define an agent and a message store.
+
+Configuring Chats
+~~~~~~~~~~~~~~~~~
+
+Chats are defined in the ``chat`` section of your configuration:
+
+.. code-block:: yaml
+
+    ai:
+        # ...
+        chat:
+            youtube:
+                agent: 'ai.agent.youtube'
+                message_store: 'ai.message_store.cache.youtube'
 
 .. _`Symfony AI Agent`: https://github.com/symfony/ai-agent
 .. _`Symfony AI Chat`: https://github.com/symfony/ai-chat

src/ai-bundle/CHANGELOG.md

@@ -33,3 +33,4 @@ CHANGELOG
    - Automatic registration of token output processors for Mistral, OpenAI and Vertex AI
    - Token usage metadata in agent results including prompt, completion, total, cached, and thinking tokens
    - Rate limit information tracking for supported platforms
+ * Add support for configuring chats and message stores

src/ai-bundle/composer.json

@@ -19,6 +19,7 @@
         "symfony/ai-chat": "@dev",
         "symfony/ai-platform": "@dev",
         "symfony/ai-store": "@dev",
+        "symfony/clock": "^7.3|^8.0",
         "symfony/config": "^7.3|^8.0",
         "symfony/console": "^7.3|^8.0",
         "symfony/dependency-injection": "^7.3|^8.0",
@@ -30,7 +31,6 @@
         "phpstan/phpstan": "^2.1",
         "phpstan/phpstan-strict-rules": "^2.0",
         "phpunit/phpunit": "^11.5",
-        "symfony/clock": "^7.3|^8.0",
         "symfony/expression-language": "^7.3|^8.0",
         "symfony/security-core": "^7.3|^8.0",
         "symfony/translation": "^7.3|^8.0"

src/ai-bundle/config/options.php

@@ -814,6 +814,15 @@
                     ->end()
                 ->end()
             ->end()
+            ->arrayNode('chat')
+                ->useAttributeAsKey('name')
+                ->arrayPrototype()
+                    ->children()
+                        ->stringNode('agent')->cannotBeEmpty()->end()
+                        ->stringNode('message_store')->cannotBeEmpty()->end()
+                    ->end()
+                ->end()
+            ->end()
             ->arrayNode('vectorizer')
                 ->info('Vectorizers for converting strings to Vector objects and transforming TextDocument arrays to VectorDocument arrays')
                 ->useAttributeAsKey('name')

src/ai-bundle/config/services.php

@@ -169,6 +169,8 @@
             ->args([
                 tagged_iterator('ai.traceable_platform'),
                 tagged_iterator('ai.traceable_toolbox'),
+                tagged_iterator('ai.traceable_message_store'),
+                tagged_iterator('ai.traceable_chat'),
             ])
             ->tag('data_collector')
 

src/ai-bundle/src/AiBundle.php

@@ -31,13 +31,17 @@
 use Symfony\AI\Agent\Toolbox\ToolFactory\MemoryToolFactory;
 use Symfony\AI\AiBundle\DependencyInjection\ProcessorCompilerPass;
 use Symfony\AI\AiBundle\Exception\InvalidArgumentException;
+use Symfony\AI\AiBundle\Profiler\TraceableChat;
+use Symfony\AI\AiBundle\Profiler\TraceableMessageStore;
 use Symfony\AI\AiBundle\Profiler\TraceablePlatform;
 use Symfony\AI\AiBundle\Profiler\TraceableToolbox;
 use Symfony\AI\AiBundle\Security\Attribute\IsGrantedTool;
 use Symfony\AI\Chat\Bridge\HttpFoundation\SessionStore;
 use Symfony\AI\Chat\Bridge\Meilisearch\MessageStore as MeilisearchMessageStore;
 use Symfony\AI\Chat\Bridge\Pogocache\MessageStore as PogocacheMessageStore;
 use Symfony\AI\Chat\Bridge\Redis\MessageStore as RedisMessageStore;
+use Symfony\AI\Chat\Chat;
+use Symfony\AI\Chat\ChatInterface;
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Bridge\Anthropic\PlatformFactory as AnthropicPlatformFactory;
 use Symfony\AI\Platform\Bridge\Azure\OpenAi\PlatformFactory as AzureOpenAiPlatformFactory;
@@ -181,11 +185,49 @@ public function loadExtension(array $config, ContainerConfigurator $container, C
             $builder->setAlias(MessageStoreInterface::class, reset($messageStores));
         }
 
+        if ($builder->getParameter('kernel.debug')) {
+            foreach ($messageStores as $messageStore) {
+                $traceableMessageStoreDefinition = (new Definition(TraceableMessageStore::class))
+                    ->setDecoratedService($messageStore)
+                    ->setArguments([
+                        new Reference('.inner'),
+                        new Reference(ClockInterface::class),
+                    ])
+                    ->addTag('ai.traceable_message_store');
+                $suffix = u($messageStore)->afterLast('.')->toString();
+                $builder->setDefinition('ai.traceable_message_store.'.$suffix, $traceableMessageStoreDefinition);
+            }
+        }
+
         if ([] === $messageStores) {
             $builder->removeDefinition('ai.command.setup_message_store');
             $builder->removeDefinition('ai.command.drop_message_store');
         }
 
+        foreach ($config['chat'] ?? [] as $name => $chat) {
+            $this->processChatConfig($name, $chat, $builder);
+        }
+
+        $chats = array_keys($builder->findTaggedServiceIds('ai.chat'));
+
+        if (1 === \count($chats)) {
+            $builder->setAlias(ChatInterface::class, reset($chats));
+        }
+
+        if ($builder->getParameter('kernel.debug')) {
+            foreach ($chats as $chat) {
+                $traceableChatDefinition = (new Definition(TraceableChat::class))
+                    ->setDecoratedService($chat)
+                    ->setArguments([
+                        new Reference('.inner'),
+                        new Reference(ClockInterface::class),
+                    ])
+                    ->addTag('ai.traceable_chat');
+                $suffix = u($chat)->afterLast('.')->toString();
+                $builder->setDefinition('ai.traceable_chat.'.$suffix, $traceableChatDefinition);
+            }
+        }
+
         foreach ($config['vectorizer'] ?? [] as $vectorizerName => $vectorizer) {
             $this->processVectorizerConfig($vectorizerName, $vectorizer, $builder);
         }
@@ -1475,6 +1517,26 @@ private function processMessageStoreConfig(string $type, array $messageStores, C
         }
     }
 
+    /**
+     * @param array{
+     *     agent: string,
+     *     message_store: string,
+     * } $configuration
+     */
+    private function processChatConfig(string $name, array $configuration, ContainerBuilder $container): void
+    {
+        $definition = new Definition(Chat::class);
+        $definition
+            ->setArguments([
+                new Reference($configuration['agent']),
+                new Reference($configuration['message_store']),
+            ])
+            ->addTag('ai.chat');
+
+        $container->setDefinition('ai.chat.'.$name, $definition);
+        $container->registerAliasForArgument('ai.chat.'.$name, ChatInterface::class, $name);
+    }
+
     /**
      * @param array<string, mixed> $config
      */

src/ai-bundle/src/Profiler/DataCollector.php

@@ -23,6 +23,8 @@
  * @author Christopher Hertel <mail@christopher-hertel.de>
  *
  * @phpstan-import-type PlatformCallData from TraceablePlatform
+ * @phpstan-import-type MessageStoreData from TraceableMessageStore
+ * @phpstan-import-type ChatData from TraceableChat
  */
 final class DataCollector extends AbstractDataCollector implements LateDataCollectorInterface
 {
@@ -37,15 +39,31 @@ final class DataCollector extends AbstractDataCollector implements LateDataColle
     private readonly array $toolboxes;
 
     /**
-     * @param TraceablePlatform[] $platforms
-     * @param TraceableToolbox[]  $toolboxes
+     * @var TraceableMessageStore[]
+     */
+    private readonly array $messageStores;
+
+    /**
+     * @var TraceableChat[]
+     */
+    private readonly array $chats;
+
+    /**
+     * @param TraceablePlatform[]     $platforms
+     * @param TraceableToolbox[]      $toolboxes
+     * @param TraceableMessageStore[] $messageStores
+     * @param TraceableChat[]         $chats
      */
     public function __construct(
         iterable $platforms,
         iterable $toolboxes,
+        iterable $messageStores,
+        iterable $chats,
     ) {
         $this->platforms = $platforms instanceof \Traversable ? iterator_to_array($platforms) : $platforms;
         $this->toolboxes = $toolboxes instanceof \Traversable ? iterator_to_array($toolboxes) : $toolboxes;
+        $this->messageStores = $messageStores instanceof \Traversable ? iterator_to_array($messageStores) : $messageStores;
+        $this->chats = $chats instanceof \Traversable ? iterator_to_array($chats) : $chats;
     }
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
@@ -59,6 +77,8 @@ public function lateCollect(): void
             'tools' => $this->getAllTools(),
             'platform_calls' => array_merge(...array_map($this->awaitCallResults(...), $this->platforms)),
             'tool_calls' => array_merge(...array_map(fn (TraceableToolbox $toolbox) => $toolbox->calls, $this->toolboxes)),
+            'messages' => array_merge(...array_map(static fn (TraceableMessageStore $messageStore): array => $messageStore->calls, $this->messageStores)),
+            'chats' => array_merge(...array_map(static fn (TraceableChat $chat): array => $chat->calls, $this->chats)),
         ];
     }
 
@@ -91,6 +111,22 @@ public function getToolCalls(): array
         return $this->data['tool_calls'] ?? [];
     }
 
+    /**
+     * @return MessageStoreData[]
+     */
+    public function getMessages(): array
+    {
+        return $this->data['messages'] ?? [];
+    }
+
+    /**
+     * @return ChatData[]
+     */
+    public function getChats(): array
+    {
+        return $this->data['chats'] ?? [];
+    }
+
     /**
      * @return Tool[]
      */

src/ai-bundle/src/Profiler/TraceableChat.php

@@ -0,0 +1,69 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\AiBundle\Profiler;
+
+use Symfony\AI\Chat\ChatInterface;
+use Symfony\AI\Platform\Message\AssistantMessage;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Platform\Message\UserMessage;
+use Symfony\Component\Clock\ClockInterface;
+
+/**
+ * @author Guillaume Loulier <personal@guillaumeloulier.fr>
+ *
+ * @phpstan-type ChatData array{
+ *      action: string,
+ *      bag?: MessageBag,
+ *      message?: UserMessage,
+ *      saved_at: \DateTimeImmutable,
+ *  }
+ */
+final class TraceableChat implements ChatInterface
+{
+    /**
+     * @var array<int, array{
+     *     action: string,
+     *     bag?: MessageBag,
+     *     message?: UserMessage,
+     *     saved_at: \DateTimeImmutable,
+     * }>
+     */
+    public array $calls = [];
+
+    public function __construct(
+        private readonly ChatInterface $chat,
+        private readonly ClockInterface $clock,
+    ) {
+    }
+
+    public function initiate(MessageBag $messages): void
+    {
+        $this->calls[] = [
+            'action' => __FUNCTION__,
+            'bag' => $messages,
+            'saved_at' => $this->clock->now(),
+        ];
+
+        $this->chat->initiate($messages);
+    }
+
+    public function submit(UserMessage $message): AssistantMessage
+    {
+        $this->calls[] = [
+            'action' => __FUNCTION__,
+            'message' => $message,
+            'saved_at' => $this->clock->now(),
+        ];
+
+        return $this->chat->submit($message);
+    }
+}