EnumOutputParser.php

<?php

declare(strict_types=1);

namespace Cortex\OutputParsers;

use BackedEnum;
use Cortex\LLM\Data\ChatGeneration;
use Cortex\LLM\Data\ChatGenerationChunk;
use Cortex\Exceptions\OutputParserException;

class EnumOutputParser extends AbstractOutputParser
{
    /**
     * @param class-string<\BackedEnum> $enum
     */
    public function __construct(
        protected string $enum,
    ) {}

    public function parse(ChatGeneration|ChatGenerationChunk|string $output): BackedEnum
    {
        if (! enum_exists($this->enum) || ! is_subclass_of($this->enum, BackedEnum::class)) {
            throw OutputParserException::failed(sprintf('Invalid enum: %s', $this->enum));
        }

        $enumName = class_basename($this->enum);

        // If the output has a tool call with the enum name, then assume we are using the schema tool.
        if (! is_string($output) && $output->message->hasToolCall($enumName)) {
            $schemaTool = $output->message->getToolCall($enumName);
            $parsed = (new JsonOutputToolsParser(key: $schemaTool->function->name, singleToolCall: true))
                ->parse($output);

            if (array_key_exists($enumName, $parsed)) {
                $enum = $this->enum::tryFrom($parsed[$enumName]);

                if ($enum !== null) {
                    return $enum;
                }
            }
        }

        $textOutput = match (true) {
            $output instanceof ChatGenerationChunk => $output->contentSoFar,
            $output instanceof ChatGeneration => $output->message->text() ?? '',
            default => $output,
        };

        // First try to get the enum from the string output
        $enum = $this->enum::tryFrom($textOutput);

        // Then use the json output parser to get the enum from the json output
        if ($enum === null) {
            try {
                $data = (new JsonOutputParser())->parse($output);
            } catch (OutputParserException $e) {
                throw OutputParserException::failed(
                    sprintf('Enum %s not found in output', $this->enum),
                    previous: $e,
                );
            }

            if (array_key_exists($enumName, $data)) {
                $enum = $this->enum::from($data[$enumName]);
            }
        }

        if ($enum === null) {
            throw OutputParserException::failed(sprintf('Enum %s not found in output', $this->enum));
        }

        return $enum;
    }
}