Skip to content

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

@ilyachase

Description

@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! 🙂

Metadata

Metadata

Assignees

No one assigned

    Labels

    help wantedExtra attention is neededquestionFurther information is requested

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions