[Chat] Support for chat "fork"

src/chat/src/Bridge/HttpFoundation/SessionStore.php

@@ -12,6 +12,7 @@
 namespace Symfony\AI\Chat\Bridge\HttpFoundation;
 
 use Symfony\AI\Agent\Exception\RuntimeException;
+use Symfony\AI\Chat\ForkedMessageStoreInterface;
 use Symfony\AI\Chat\ManagedStoreInterface;
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Message\MessageBag;
@@ -21,12 +22,12 @@
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  */
-final readonly class SessionStore implements ManagedStoreInterface, MessageStoreInterface
+final readonly class SessionStore implements ManagedStoreInterface, MessageStoreInterface, ForkedMessageStoreInterface
 {
     private SessionInterface $session;
 
     public function __construct(
-        RequestStack $requestStack,
+        private RequestStack $requestStack,
         private string $sessionKey = 'messages',
     ) {
         if (!class_exists(RequestStack::class)) {
@@ -46,13 +47,23 @@ public function save(MessageBag $messages): void
         $this->session->set($this->sessionKey, $messages);
     }
 
-    public function load(): MessageBag
+    public function load(?string $id = null): MessageBag
     {
-        return $this->session->get($this->sessionKey, new MessageBag());
+        return $this->session->get($id ?? $this->sessionKey, new MessageBag());
     }
 
     public function drop(): void
     {
         $this->session->remove($this->sessionKey);
     }
+
+    public function fork(string $id, MessageBag $existingMessages): ForkedMessageStoreInterface
+    {
+        $fork = new self($this->requestStack, $id);
+
+        $fork->setup();
+        $fork->save($existingMessages);
+
+        return $fork;
+    }
 }

src/chat/src/Bridge/Local/CacheStore.php

@@ -13,14 +13,15 @@
 
 use Psr\Cache\CacheItemPoolInterface;
 use Symfony\AI\Agent\Exception\RuntimeException;
+use Symfony\AI\Chat\ForkedMessageStoreInterface;
 use Symfony\AI\Chat\ManagedStoreInterface;
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Message\MessageBag;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  */
-final readonly class CacheStore implements ManagedStoreInterface, MessageStoreInterface
+final readonly class CacheStore implements ManagedStoreInterface, MessageStoreInterface, ForkedMessageStoreInterface
 {
     public function __construct(
         private CacheItemPoolInterface $cache,
@@ -52,9 +53,9 @@ public function save(MessageBag $messages): void
         $this->cache->save($item);
     }
 
-    public function load(): MessageBag
+    public function load(?string $id = null): MessageBag
     {
-        $item = $this->cache->getItem($this->cacheKey);
+        $item = $this->cache->getItem($id ?? $this->cacheKey);
 
         return $item->isHit() ? $item->get() : new MessageBag();
     }
@@ -63,4 +64,14 @@ public function drop(): void
     {
         $this->cache->deleteItem($this->cacheKey);
     }
+
+    public function fork(string $id, MessageBag $existingMessages): ForkedMessageStoreInterface
+    {
+        $fork = new self($this->cache, $id, $this->ttl);
+
+        $fork->setup();
+        $fork->save($existingMessages);
+
+        return $fork;
+    }
 }

src/chat/src/Bridge/Local/InMemoryStore.php

@@ -11,14 +11,15 @@
 
 namespace Symfony\AI\Chat\Bridge\Local;
 
+use Symfony\AI\Chat\ForkedMessageStoreInterface;
 use Symfony\AI\Chat\ManagedStoreInterface;
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Message\MessageBag;
 
 /**
  * @author Christopher Hertel <mail@christopher-hertel.de>
  */
-final class InMemoryStore implements ManagedStoreInterface, MessageStoreInterface
+final class InMemoryStore implements ManagedStoreInterface, MessageStoreInterface, ForkedMessageStoreInterface
 {
     /**
      * @var MessageBag[]
@@ -40,13 +41,23 @@ public function save(MessageBag $messages): void
         $this->messages[$this->identifier] = $messages;
     }
 
-    public function load(): MessageBag
+    public function load(?string $id = null): MessageBag
     {
-        return $this->messages[$this->identifier] ?? new MessageBag();
+        return $this->messages[$id ?? $this->identifier] ?? new MessageBag();
     }
 
     public function drop(): void
     {
         $this->messages[$this->identifier] = new MessageBag();
     }
+
+    public function fork(string $id, MessageBag $existingMessages): ForkedMessageStoreInterface
+    {
+        $fork = new self($id);
+
+        $fork->setup();
+        $fork->save($existingMessages);
+
+        return $fork;
+    }
 }

src/chat/src/Bridge/Meilisearch/MessageStore.php

@@ -14,6 +14,7 @@
 use Symfony\AI\Chat\Exception\InvalidArgumentException;
 use Symfony\AI\Chat\Exception\LogicException;
 use Symfony\AI\Chat\Exception\RuntimeException;
+use Symfony\AI\Chat\ForkedMessageStoreInterface;
 use Symfony\AI\Chat\ManagedStoreInterface;
 use Symfony\AI\Chat\MessageStoreInterface;
 use Symfony\AI\Platform\Message\AssistantMessage;
@@ -37,7 +38,7 @@
 /**
  * @author Guillaume Loulier <personal@guillaumeloulier.fr>
  */
-final readonly class MessageStore implements ManagedStoreInterface, MessageStoreInterface
+final readonly class MessageStore implements ManagedStoreInterface, MessageStoreInterface, ForkedMessageStoreInterface
 {
     public function __construct(
         private HttpClientInterface $httpClient,
@@ -79,9 +80,9 @@ public function save(MessageBag $messages): void
         ));
     }
 
-    public function load(): MessageBag
+    public function load(?string $id = null): MessageBag
     {
-        $messages = $this->request('POST', \sprintf('indexes/%s/documents/fetch', $this->indexName), [
+        $messages = $this->request('POST', \sprintf('indexes/%s/documents/fetch', $id ?? $this->indexName), [
             'sort' => ['addedAt:asc'],
         ]);
 
@@ -93,6 +94,16 @@ public function drop(): void
         $this->request('DELETE', \sprintf('indexes/%s/documents', $this->indexName));
     }
 
+    public function fork(string $id, MessageBag $existingMessages): ForkedMessageStoreInterface
+    {
+        $fork = new self($this->httpClient, $this->endpointUrl, $this->apiKey, $this->clock, $id);
+
+        $fork->setup();
+        $fork->save($existingMessages);
+
+        return $fork;
+    }
+
     /**
      * @param array<string, mixed>|list<array<string, mixed>> $payload
      *

src/chat/src/Chat.php

@@ -25,7 +25,7 @@
 {
     public function __construct(
         private AgentInterface $agent,
-        private MessageStoreInterface&ManagedStoreInterface $store,
+        private MessageStoreInterface&ManagedStoreInterface&ForkedMessageStoreInterface $store,
     ) {
     }
 
@@ -51,4 +51,16 @@ public function submit(UserMessage $message): AssistantMessage
 
         return $assistantMessage;
     }
+
+    public function fork(string $id): self
+    {
+        $existingMessages = $this->store->load();
+
+        $forkedStore = $this->store->fork($id, $existingMessages);
+
+        return new self(
+            $this->agent,
+            $forkedStore,
+        );
+    }
 }

src/chat/src/ChatInterface.php

@@ -27,4 +27,6 @@ public function initiate(MessageBag $messages): void;
      * @throws ExceptionInterface When the chat submission fails due to agent errors
      */
     public function submit(UserMessage $message): AssistantMessage;
+
+    public function fork(string $id): self;
 }

src/chat/src/ForkedMessageStoreInterface.php

@@ -0,0 +1,22 @@
+<?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\Chat;
+
+use Symfony\AI\Platform\Message\MessageBag;
+
+/**
+ * @author Guillaume Loulier <personal@guillaumeloulier.fr>
+ */
+interface ForkedMessageStoreInterface
+{
+    public function fork(string $id, MessageBag $existingMessages): self;
+}

src/chat/src/MessageStoreInterface.php

@@ -20,5 +20,5 @@ interface MessageStoreInterface
 {
     public function save(MessageBag $messages): void;
 
-    public function load(): MessageBag;
+    public function load(?string $id = null): MessageBag;
 }

src/chat/tests/ChatTest.php

@@ -52,7 +52,7 @@ public function testItSubmitsUserMessageAndReturnsAssistantMessage()
 
         $this->agent->expects($this->once())
             ->method('call')
-            ->with($this->callback(function (MessageBag $messages) use ($userMessage) {
+            ->with($this->callback(static function (MessageBag $messages) use ($userMessage): bool {
                 $messagesArray = $messages->getMessages();
 
                 return end($messagesArray) === $userMessage;
@@ -100,7 +100,7 @@ public function testItHandlesEmptyMessageStore()
 
         $this->agent->expects($this->once())
             ->method('call')
-            ->with($this->callback(function (MessageBag $messages) {
+            ->with($this->callback(static function (MessageBag $messages): bool {
                 $messagesArray = $messages->getMessages();
 
                 return 1 === \count($messagesArray);
@@ -113,4 +113,21 @@ public function testItHandlesEmptyMessageStore()
         $this->assertSame($assistantContent, $result->getContent());
         $this->assertCount(2, $this->store->load());
     }
+
+    public function testItCanBeForked()
+    {
+        $agent = $this->createMock(AgentInterface::class);
+
+        $store = new InMemoryStore();
+
+        $chat = new Chat($agent, $store);
+        $chat->submit(Message::ofUser('hello world'));
+
+        $this->assertCount(1, $store->load());
+
+        $forkedChat = $chat->fork('foo');
+        $forkedChat->submit(Message::ofUser('Second hello world'));
+
+        $this->assertCount(2, $store->load('foo'));
+    }
 }