AbstractLLM.php

<?php

declare(strict_types=1);

namespace Cortex\LLM;

use Closure;
use BackedEnum;
use Cortex\Pipeline;
use Cortex\Support\Utils;
use Cortex\Tools\SchemaTool;
use Cortex\LLM\Contracts\LLM;
use Cortex\LLM\Contracts\Tool;
use Cortex\LLM\Data\ToolConfig;
use Cortex\LLM\Enums\ToolChoice;
use Cortex\LLM\Contracts\Message;
use Cortex\LLM\Enums\MessageRole;
use Cortex\Contracts\OutputParser;
use Cortex\Support\Traits\CanPipe;
use Cortex\Exceptions\LLMException;
use Cortex\JsonSchema\SchemaFactory;
use Cortex\ModelInfo\Data\ModelInfo;
use Cortex\LLM\Data\ChatStreamResult;
use Cortex\Exceptions\PipelineException;
use Cortex\ModelInfo\Enums\ModelFeature;
use Cortex\JsonSchema\Types\ObjectSchema;
use Cortex\LLM\Data\Messages\UserMessage;
use Cortex\ModelInfo\Enums\ModelProvider;
use Cortex\OutputParsers\EnumOutputParser;
use Cortex\LLM\Data\Messages\SystemMessage;
use Cortex\LLM\Data\StructuredOutputConfig;
use Cortex\OutputParsers\ClassOutputParser;
use Cortex\Support\Traits\DispatchesEvents;
use Cortex\Exceptions\OutputParserException;
use Cortex\Tasks\Enums\StructuredOutputMode;
use Cortex\LLM\Data\Messages\MessageCollection;
use Cortex\OutputParsers\StructuredOutputParser;

abstract class AbstractLLM implements LLM
{
    use CanPipe;
    use DispatchesEvents;

    /**
     * @var array<string, mixed>
     */
    protected array $parameters = [];

    protected ?ToolConfig $toolConfig = null;

    protected ?StructuredOutputConfig $structuredOutputConfig = null;

    protected ?OutputParser $outputParser = null;

    protected ?string $outputParserError = null;

    protected bool $forceJsonOutput = false;

    protected bool $useCache = false;

    protected bool $streaming = false;

    protected ?ModelInfo $modelInfo = null;

    protected bool $shouldApplyFormatInstructions = false;

    /**
     * @var array<\Cortex\ModelInfo\Enums\ModelFeature>
     */
    protected array $features = [];

    protected bool $ignoreModelFeatures = false;

    public function __construct(
        protected string $model,
        protected ModelProvider $modelProvider,
    ) {
        $this->modelInfo = $modelProvider->info($model);
        $this->features = $this->modelInfo->features ?? [];
    }

    // IDEA:
    // Could have a third param which is StageMetadata
    // where it includes the class path of the next/previous stage

    public function handlePipeable(mixed $payload, Closure $next): mixed
    {
        // Invoke the LLM with the given input
        $result = match (true) {
            $payload instanceof MessageCollection, $payload instanceof Message, is_array($payload) => $this->invoke($payload),
            is_string($payload) => $this->invoke(new UserMessage($payload)),
            default => throw new PipelineException('Invalid input'),
        };

        // If the result is a stream, we attempt to yield the chunks to the next pipeable
        // And if that happens to be an output parser, we ignore any parsing errors and continue.
        // Otherwise, we return the message as is.
        return $result instanceof ChatStreamResult
            ? new ChatStreamResult(function () use ($result, $next) {
                foreach ($result as $chunk) {
                    try {
                        yield $next($chunk);
                    } catch (OutputParserException) {
                        // Ignore any parsing errors and continue
                    }
                }
            })
            : $next($result);
    }

    public function output(OutputParser $parser): Pipeline
    {
        return $this->pipe($parser);
    }

    /**
     * @param array<int, \Cortex\LLM\Contracts\Tool|\Cortex\JsonSchema\Contracts\Schema|\Closure|string> $tools
     */
    public function withTools(array $tools, ToolChoice|string $toolChoice = ToolChoice::Auto): static
    {
        $this->supportsFeatureOrFail(ModelFeature::ToolCalling);

        $this->toolConfig = $tools === []
            ? null
            : new ToolConfig(
                Utils::toToolCollection($tools)->all(),
                $toolChoice,
            );

        return $this;
    }

    /**
     * Add a tool to the LLM.
     */
    public function addTool(Tool|Closure|string $tool, ToolChoice|string $toolChoice = ToolChoice::Auto): static
    {
        return $this->withTools([...($this->toolConfig->tools ?? []), $tool], $toolChoice);
    }

    /**
     * TODO: should change to protected once Tasks are refactored to use `withStructuredOutput()`
     */
    public function withStructuredOutputConfig(
        ObjectSchema|StructuredOutputConfig $schema,
        ?string $name = null,
        ?string $description = null,
        bool $strict = true,
    ): static {
        $this->supportsFeatureOrFail(ModelFeature::StructuredOutput);

        $this->structuredOutputConfig = $schema instanceof StructuredOutputConfig
            ? $schema
            : new StructuredOutputConfig($schema, $name, $description, $strict);

        return $this;
    }

    /**
     * @param class-string|\Cortex\JsonSchema\Types\ObjectSchema $output
     */
    public function withStructuredOutput(
        ObjectSchema|string $output,
        ?string $name = null,
        ?string $description = null,
        bool $strict = true,
        StructuredOutputMode $outputMode = StructuredOutputMode::Auto,
    ): static {
        [$schema, $outputParser] = $this->resolveSchemaAndOutputParser($output, $strict);

        $this->withOutputParser($outputParser);

        if ($outputMode === StructuredOutputMode::Tool) {
            $this->supportsFeatureOrFail(ModelFeature::ToolCalling);

            return $this->withTools([new SchemaTool($schema, $name, $description)], ToolChoice::Required);
        }

        if ($outputMode === StructuredOutputMode::Auto) {
            if ($this->supportsFeature(ModelFeature::StructuredOutput)) {
                return $this->withStructuredOutputConfig(
                    $schema,
                    $name,
                    $description,
                    $strict,
                );
            }

            if ($this->supportsFeature(ModelFeature::ToolCalling)) {
                return $this->withTools([new SchemaTool($schema, $name, $description)], ToolChoice::Required);
            }
        }

        return $this->forceJsonOutput()->shouldApplyFormatInstructions(true);
    }

    public function forceJsonOutput(): static
    {
        $this->supportsFeatureOrFail(ModelFeature::JsonOutput);

        $this->forceJsonOutput = true;

        return $this;
    }

    public function withModel(string $model): static
    {
        $this->model = $model;

        return $this;
    }

    public function withTemperature(?float $temperature): static
    {
        $this->parameters['temperature'] = $temperature;

        return $this;
    }

    public function withMaxTokens(?int $maxTokens): static
    {
        $this->parameters['max_tokens'] = $maxTokens;

        return $this;
    }

    public function withStreaming(bool $streaming = true): static
    {
        $this->streaming = $streaming;

        return $this;
    }

    public function withCaching(bool $useCache = true): static
    {
        $this->useCache = $useCache;

        return $this;
    }

    public function withParameters(array $parameters): static
    {
        $this->parameters = $parameters;

        return $this;
    }

    public function withOutputParser(OutputParser $outputParser): static
    {
        $this->outputParser = $outputParser;

        return $this;
    }

    public function getModel(): string
    {
        return $this->model;
    }

    /**
     * @return array<string, mixed>
     */
    public function getParameters(): array
    {
        return $this->parameters;
    }

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

    /**
     * Determine if the LLM response should be cached.
     */
    public function shouldCache(): bool
    {
        return $this->useCache;
    }

    /**
     * Determine if the LLM supports a specific feature.
     */
    public function supportsFeature(ModelFeature $feature): bool
    {
        if ($this->ignoreModelFeatures) {
            return true;
        }

        return in_array($feature, $this->features, true);
    }

    public function supportsFeatureOrFail(ModelFeature $feature): void
    {
        if (! $this->supportsFeature($feature)) {
            throw new LLMException('LLM does not support ' . $feature->value . '.');
        }
    }

    public function getModelInfo(): ?ModelInfo
    {
        return $this->modelInfo;
    }

    public function getModelProvider(): ModelProvider
    {
        return $this->modelProvider;
    }

    public function getOutputParserError(): ?string
    {
        return $this->outputParserError;
    }

    /**
     * @return array<\Cortex\ModelInfo\Enums\ModelFeature>
     */
    public function getFeatures(): array
    {
        return $this->features;
    }

    /**
     * Explicitly set the model info for the LLM.
     * This will override the values from the ModelProvider instance.
     */
    public function withModelInfo(ModelInfo $modelInfo): static
    {
        $this->modelInfo = $modelInfo;

        return $this;
    }

    public function withFeatures(ModelFeature ...$features): static
    {
        $this->features = $features;

        return $this;
    }

    public function addFeature(ModelFeature $feature): static
    {
        $this->features[] = $feature;

        return $this;
    }

    public function ignoreFeatures(bool $ignoreModelFeatures = true): static
    {
        $this->ignoreModelFeatures = $ignoreModelFeatures;

        return $this;
    }

    /**
     * Set whether the output parser format instructions should be applied to the messages.
     */
    public function shouldApplyFormatInstructions(bool $applyFormatInstructions = true): static
    {
        $this->shouldApplyFormatInstructions = $applyFormatInstructions;

        return $this;
    }

    /**
     * Apply the given format instructions to the messages.
     *
     * @return \Cortex\LLM\Data\Messages\MessageCollection<int, \Cortex\LLM\Contracts\Message>
     */
    protected static function applyFormatInstructions(
        MessageCollection $messages,
        string $formatInstructions,
    ): MessageCollection {
        if (! $messages->hasMessageByRole(MessageRole::System)) {
            $messages->prepend(new SystemMessage($formatInstructions));
        } else {
            $messages = $messages->map(function (Message $message) use ($formatInstructions): Message {
                if ($message->role() === MessageRole::System) {
                    return $message->cloneWithContent($message->text() . "\n\n" . $formatInstructions);
                }

                return $message;
            });
        }

        return $messages;
    }

    /**
     * Resolve the schema and output parser from the given output type.
     *
     * @param class-string|\Cortex\JsonSchema\Types\ObjectSchema $outputType
     *
     * @return array{\Cortex\JsonSchema\Types\ObjectSchema, \Cortex\Contracts\OutputParser}
     */
    protected function resolveSchemaAndOutputParser(ObjectSchema|string $outputType, bool $strict = true): array
    {
        if ($outputType instanceof ObjectSchema) {
            return [$outputType, new StructuredOutputParser($outputType, $strict)];
        }

        if (enum_exists($outputType) && is_subclass_of($outputType, BackedEnum::class)) {
            $enumSchema = SchemaFactory::fromEnum($outputType);
            $title = $enumSchema->getTitle() ?? class_basename($outputType);

            $schema = new ObjectSchema($title);
            $schema->properties($enumSchema);

            return [$schema, new EnumOutputParser($outputType)];
        }

        if (class_exists($outputType)) {
            $schema = SchemaFactory::fromClass($outputType);
            $schema->title($schema->getTitle() ?? class_basename($outputType));

            return [$schema, new ClassOutputParser($outputType, $strict)];
        }

        throw new LLMException('Unsupported output type: ' . $outputType);
    }
}