tweaks

Author: tymondesigns

src/LLM/AbstractLLM.php

@@ -23,6 +23,7 @@
 use Cortex\ModelInfo\Data\ModelInfo;
 use Cortex\LLM\Data\ChatStreamResult;
 use Cortex\Exceptions\PipelineException;
+use Cortex\LLM\Data\ChatGenerationChunk;
 use Cortex\ModelInfo\Enums\ModelFeature;
 use Cortex\JsonSchema\Types\ObjectSchema;
 use Cortex\LLM\Data\Messages\UserMessage;
@@ -409,18 +410,19 @@ public function shouldParseOutput(bool $shouldParseOutput = true): static
         return $this;
     }
 
-    protected function applyOutputParserIfApplicable(ChatGeneration $generation): ChatGeneration
-    {
+    protected function applyOutputParserIfApplicable(
+        ChatGeneration|ChatGenerationChunk $generationOrChunk,
+    ): ChatGeneration|ChatGenerationChunk {
         if ($this->shouldParseOutput && $this->outputParser !== null) {
             try {
-                $parsedOutput = $this->outputParser->parse($generation);
-                $generation = $generation->cloneWithParsedOutput($parsedOutput);
+                $parsedOutput = $this->outputParser->parse($generationOrChunk);
+                $generationOrChunk = $generationOrChunk->cloneWithParsedOutput($parsedOutput);
             } catch (OutputParserException $e) {
                 $this->outputParserError = $e->getMessage();
             }
         }
 
-        return $generation;
+        return $generationOrChunk;
     }
 
     /**

src/LLM/Drivers/AnthropicChat.php

@@ -35,7 +35,6 @@
 use Cortex\LLM\Data\Messages\ToolMessage;
 use Cortex\ModelInfo\Enums\ModelProvider;
 use Cortex\LLM\Data\Messages\SystemMessage;
-use Cortex\Exceptions\OutputParserException;
 use Cortex\Tasks\Enums\StructuredOutputMode;
 use Cortex\LLM\Data\Messages\AssistantMessage;
 use Cortex\LLM\Data\Messages\MessageCollection;
@@ -273,9 +272,7 @@ protected function mapStreamResponse(StreamResponse $response): ChatStreamResult
                     );
                 }
 
-                $isFinal = $finishReason !== null;
-
-                $generationChunk = new ChatGenerationChunk(
+                $chunk = new ChatGenerationChunk(
                     id: $messageId ?? 'unknown',
                     message: new AssistantMessage(
                         content: $chunkDelta,
@@ -293,25 +290,18 @@ protected function mapStreamResponse(StreamResponse $response): ChatStreamResult
                     finishReason: $finishReason,
                     usage: $usage,
                     contentSoFar: $contentSoFar,
-                    isFinal: $isFinal,
+                    isFinal: $finishReason !== null,
                 );
 
-                if ($this->outputParser !== null) {
-                    try {
-                        $parsedOutput = $this->outputParser->parse($generationChunk);
-                        $generationChunk = $generationChunk->cloneWithParsedOutput($parsedOutput);
-                    } catch (OutputParserException $e) {
-                        $this->outputParserError = $e->getMessage();
-                    }
-                }
+                $chunk = $this->applyOutputParserIfApplicable($chunk);
 
                 $this->dispatchEvent(
-                    $generationChunk->isFinal
-                        ? new ChatModelEnd($generationChunk)
-                        : new ChatModelStream($generationChunk),
+                    $chunk->isFinal
+                        ? new ChatModelEnd($chunk)
+                        : new ChatModelStream($chunk),
                 );
 
-                yield $generationChunk;
+                yield $chunk;
             }
         });
     }

src/LLM/Drivers/OpenAIChat.php

@@ -213,7 +213,6 @@ protected function mapStreamResponse(StreamResponse $response): ChatStreamResult
                 }
 
                 $finishReason = static::mapFinishReason($choice->finishReason ?? null);
-                $isFinal = $finishReason !== null; // && $usage !== null;
 
                 $chunk = new ChatGenerationChunk(
                     id: $chunk->id,
@@ -233,17 +232,10 @@ protected function mapStreamResponse(StreamResponse $response): ChatStreamResult
                     finishReason: $finishReason,
                     usage: $usage,
                     contentSoFar: $contentSoFar,
-                    isFinal: $isFinal,
+                    isFinal: $finishReason !== null,
                 );
 
-                if ($this->outputParser !== null) {
-                    try {
-                        $parsedOutput = $this->outputParser->parse($chunk);
-                        $chunk = $chunk->cloneWithParsedOutput($parsedOutput);
-                    } catch (OutputParserException $e) {
-                        $this->outputParserError = $e->getMessage();
-                    }
-                }
+                $chunk = $this->applyOutputParserIfApplicable($chunk);
 
                 $this->dispatchEvent(
                     $chunk->isFinal