Improve library structure to allow easier message parsing and request preparation for LLMs

[ilyachase]

I wrote a very basic version of the Claude API client and needed to implement prompt caching. However, I found it quite challenging to do so, considering the structure of the library - it makes it hard to pass additional information along with messages.

I explored a couple of options:

1. I tried to introduce a custom option, something like CacheMessagesOption, where I could refer to the key attribute of the SolutionMetadata, but the metadata is transformed into a MessagePrompt and then later into ChatMessage and the key is lost along the way.
2. Therefore, I had to write quite a lot of custom code to pass caching attributes. Specifically, I needed to override AgentMemoryInjector and introduce CacheableSolutionMetadata, CacheableMessagePrompt, and CacheableChatMessage. I'm going to provide some code below so it's easier to understand.

readonly class CacheableSolutionMetadata extends SolutionMetadata
{
    public function __construct(MetadataType $type, string $key, mixed $content, private readonly bool $isCached = false)
    {
        parent::__construct($type, $key, $content);
    }

    public function isCached() : bool
    {
        return $this->isCached;
    }
}

// later in the agent:

$aggregate->addMetadata(
    new CacheableSolutionMetadata(
        type: MetadataType::Memory,
        key: 'foo',
        content: 'bar'
        isCached: true,
    )
);

Then, in order to transform it into a proper request message:

class AgentMemoryInjector implements PromptInterceptorInterface
{
    public function generate(
        PromptGeneratorInput $input,
        InterceptorHandler $next,
    ) : PromptInterface {
        foreach ($input->agent->getMemory() as $metadata) {
            if ($metadata instanceof CacheableSolutionMetadata) {
                $prompt = CacheableMessagePrompt::system(prompt: $metadata->content, isCached: $metadata->isCached());
            } else {
                $prompt = MessagePrompt::system(prompt: $metadata->content);
            }

            $input = $input->withPrompt($input->prompt->withAddedMessage($prompt));
        }

        return $next(input: $input);
    }
}

final class CacheableMessagePrompt implements StringPromptInterface, HasRoleInterface, SerializableInterface
{
    public function __construct(
        private StringPromptInterface $prompt,
        public Role $role = Role::User,
        private array $with = [],
        private readonly bool $isCached = false
    ) {
    }

    public static function system(
        StringPromptInterface|string|\Stringable $prompt,
        array $values = [],
        array $with = [],
        bool $isCached = false
    ) : self {
        if (\is_string($prompt)) {
            $prompt = new StringPrompt($prompt);
        }

        return new self($prompt->withValues($values), Role::System, $with, $isCached);
    }

 // copy of the rest of the class with $isCached added

    public function toChatMessage(array $parameters = []) : ?ChatMessage
    {
        $prompt = $this->prompt;

        foreach ($this->with as $var) {
            if (!isset($parameters[$var]) || empty($parameters[$var])) {
                // condition failed
                return null;
            }
        }

        return new CacheableChatMessage(
            $prompt instanceof DataPrompt ? $prompt->toArray() : $prompt->format($parameters),
            $this->role,
            $this->isCached
        );
    }
}

class CacheableChatMessage extends ChatMessage
{
    public function __construct(array|string $content, Role $role = Role::User, private readonly bool $isCached = false)
    {
        parent::__construct($content, $role);
    }

    public function isCached(): bool
    {
        return $this->isCached;
    }
}

class MessageMapper
{
    public function map(object $message) : array
    {
        if ($message instanceof CacheableMessagePrompt) {
            $message = $message->toChatMessage();
        }

        if ($message instanceof ChatMessage) {
            $cacheControl = [];

            if ($message instanceof CacheableChatMessage && $message->isCached()) {
                $cacheControl = [
                    'cache_control' => ['type' => 'ephemeral'],
                ];
            }

            if ($message->role === Role::System) {
                return [
                    'text' => $message->content,
                    'type' => 'text',
                ] + $cacheControl;
            }

            return [
                'content' => $message->content,
                'role' => $this->mapRole($message->role),
            ] + $cacheControl;
        }
    // rest of the code
    }
}

I see a couple of issues:

1. Too much boilerplate code and overrides to pass custom information along with messages.
2. Even though the library leverages the Injectors pattern, in reality, it is not very easy to customize in cases like the above
3. Having final classes, such as final class MessagePrompt, makes it harder to override the behavior

Maybe I'm missing some better and more clever way of achieving this, so I'm really open to hearing it! 🙂

[butschster]

Hey! My telegram is @butschster We can talk there

[ilyachase]

@butschster after working some more with the library, I realized what the pitfall is - instead of having final classes which are then checked using instanceof, we should turn them into interfaces.
Example:

agents/src/AgentExecutor/Interceptor/ToolExecutorInterceptor.php

Line 42 in b295a2e

if ($result instanceof ToolCalledResponse) {

ToolCalledResponse should be an interface. Then end users can define their own response classes and extend LLM clients much more easily, whereas it's impossible to extend ToolCalledResponse currently because it's final.