StructuredOutputParser.php

<?php

declare(strict_types=1);

namespace Cortex\OutputParsers;

use Override;
use Cortex\LLM\Data\ChatGeneration;
use Cortex\JsonSchema\Contracts\Schema;
use Cortex\LLM\Data\ChatGenerationChunk;
use Cortex\JsonSchema\Types\ObjectSchema;

class StructuredOutputParser extends AbstractOutputParser
{
    public function __construct(
        protected Schema $schema,
        protected bool $strict = true,
    ) {}

    /**
     * @return array<string|int, mixed>
     */
    public function parse(ChatGeneration|ChatGenerationChunk|string $output): array
    {
        $parser = match (true) {
            is_string($output) => new JsonOutputParser(),
            // If the message has tool calls and no text, assume we are using the schema tool
            $output->message->hasToolCalls() && in_array($output->message->text(), [null, ''], true) => new JsonOutputToolsParser(singleToolCall: true),
            default => new JsonOutputParser(),
        };

        $parsedOutput = $parser->parse($output);

        if ($this->shouldEnforceStrictOutput($output)) {
            // Filter out any keys that are not in the schema
            if ($this->schema instanceof ObjectSchema) {
                $parsedOutput = array_intersect_key($parsedOutput, array_flip($this->schema->getPropertyKeys()));
            }

            // Ensure the output is valid as per the schema
            $this->schema->validate($parsedOutput);
        }

        return $parsedOutput;
    }

    public function shouldEnforceStrictOutput(ChatGeneration|ChatGenerationChunk|string $output): bool
    {
        if ($this->strict) {
            // Only the final chunk of streaming output should be validated
            return ! ($output instanceof ChatGenerationChunk && ! $output->isFinal);
        }

        return false;
    }

    #[Override]
    public function formatInstructions(): ?string
    {
        if ($this->formatInstructions !== null) {
            return $this->formatInstructions;
        }

        $schema = $this->schema->toJson(JSON_THROW_ON_ERROR);

        return <<<FORMAT
            You MUST format your output as a JSON value that adheres to a given "JSON Schema" instance.
            Here is the JSON Schema instance your output must adhere to:
            ```json
            {$schema}
            ```
            FORMAT;
    }
}