improvements

Author: tymondesigns

src/Agents/Agent.php

@@ -10,6 +10,7 @@
 use Cortex\LLM\Data\Usage;
 use Cortex\Prompts\Prompt;
 use Cortex\Events\AgentEnd;
+use Illuminate\Support\Str;
 use Cortex\Contracts\ToolKit;
 use Cortex\Events\AgentStart;
 use Cortex\JsonSchema\Schema;
@@ -32,7 +33,6 @@
 use Illuminate\Support\Collection;
 use Cortex\Events\AgentStreamChunk;
 use Cortex\LLM\Data\ChatStreamResult;
-use Cortex\Agents\Contracts\Middleware;
 use Cortex\Events\Contracts\AgentEvent;
 use Cortex\Exceptions\GenericException;
 use Cortex\Memory\Stores\InMemoryStore;
@@ -48,15 +48,13 @@
 use Cortex\Support\Traits\DispatchesEvents;
 use Illuminate\Contracts\Support\Arrayable;
 use Cortex\LLM\Contracts\LLM as LLMContract;
+use Cortex\Agents\Concerns\HandlesMiddleware;
 use Cortex\Agents\Stages\TrackAgentStepStart;
 use Cortex\Prompts\Builders\ChatPromptBuilder;
-use Cortex\Agents\Middleware\AfterModelWrapper;
 use Cortex\LLM\Data\Messages\MessageCollection;
-use Cortex\Agents\Middleware\BeforeModelWrapper;
 use Cortex\LLM\Data\Messages\MessagePlaceholder;
 use Cortex\Prompts\Templates\ChatPromptTemplate;
 use Cortex\Agents\Contracts\AfterModelMiddleware;
-use Cortex\Agents\Middleware\BeforePromptWrapper;
 use Cortex\Agents\Contracts\BeforeModelMiddleware;
 use Cortex\Agents\Contracts\BeforePromptMiddleware;
 use Cortex\Contracts\ChatMemory as ChatMemoryContract;
@@ -67,6 +65,7 @@ class Agent implements Pipeable
 {
     use CanPipe;
     use DispatchesEvents;
+    use HandlesMiddleware;
 
     protected LLMContract $llm;
 
@@ -383,57 +382,6 @@ protected function executionStages(): array
         ];
     }
 
-    /**
-     * Get the middleware of a specific type.
-     * If middleware implements multiple middleware interfaces, wrap it appropriately
-     * to delegate to the correct hook method (beforePrompt, beforeModel, or afterModel).
-     *
-     * @param class-string<\Cortex\Agents\Contracts\Middleware> $type
-     *
-     * @return array<int, \Cortex\Agents\Contracts\Middleware>
-     */
-    protected function getMiddleware(string $type): array
-    {
-        return array_map(
-            function (Middleware $middleware) use ($type): Middleware {
-                // Wrap all hook-based middleware to ensure hook methods are called
-                if (! $this->isHookMiddlewareType($type)) {
-                    return $middleware;
-                }
-
-                // If middleware implements multiple interfaces, wrap to delegate to correct hook
-                // If it only implements one interface, still wrap to ensure hook method is called
-                return $this->wrapMiddleware($middleware, $type);
-            },
-            array_filter($this->middleware, fn(Middleware $middleware): bool => $middleware instanceof $type),
-        );
-    }
-
-    /**
-     * Check if the given type is a hook-based middleware interface.
-     */
-    protected function isHookMiddlewareType(string $type): bool
-    {
-        return in_array($type, [
-            BeforePromptMiddleware::class,
-            BeforeModelMiddleware::class,
-            AfterModelMiddleware::class,
-        ], true);
-    }
-
-    /**
-     * Wrap middleware to delegate to the appropriate hook method.
-     */
-    protected function wrapMiddleware(Middleware $middleware, string $type): Middleware
-    {
-        return match ($type) {
-            BeforePromptMiddleware::class => new BeforePromptWrapper($middleware),
-            BeforeModelMiddleware::class => new BeforeModelWrapper($middleware),
-            AfterModelMiddleware::class => new AfterModelWrapper($middleware),
-            default => $middleware,
-        };
-    }
-
     /**
      * @param array<int, \Cortex\LLM\Contracts\Message> $messages
      * @param array<string, mixed> $input
@@ -546,8 +494,10 @@ protected static function buildPromptTemplate(
      */
     protected static function buildMemory(ChatPromptTemplate $prompt, ?Store $memoryStore = null): ChatMemoryContract
     {
-        $memoryStore ??= new InMemoryStore();
-        $memoryStore->setMessages($prompt->messages->withoutPlaceholders());
+        $memoryStore ??= new InMemoryStore(
+            threadId: Str::uuid7()->toString(),
+            messages: $prompt->messages->withoutPlaceholders(),
+        );
 
         return new ChatMemory($memoryStore);
     }

src/Agents/Concerns/HandlesMiddleware.php

@@ -0,0 +1,69 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Cortex\Agents\Concerns;
+
+use Cortex\Agents\Contracts\Middleware;
+use Cortex\Agents\Middleware\AfterModelWrapper;
+use Cortex\Agents\Middleware\BeforeModelWrapper;
+use Cortex\Agents\Contracts\AfterModelMiddleware;
+use Cortex\Agents\Middleware\BeforePromptWrapper;
+use Cortex\Agents\Contracts\BeforeModelMiddleware;
+use Cortex\Agents\Contracts\BeforePromptMiddleware;
+
+trait HandlesMiddleware
+{
+    /**
+     * Get the middleware of a specific type.
+     * If middleware implements multiple middleware interfaces, wrap it appropriately
+     * to delegate to the correct hook method (beforePrompt, beforeModel, or afterModel).
+     *
+     * @param class-string<\Cortex\Agents\Contracts\Middleware> $type
+     *
+     * @return array<int, \Cortex\Agents\Contracts\Middleware>
+     */
+    protected function getMiddleware(string $type): array
+    {
+        return array_map(
+            function (Middleware $middleware) use ($type): Middleware {
+                // Wrap all hook-based middleware to ensure hook methods are called
+                if (! $this->isHookMiddlewareType($type)) {
+                    return $middleware;
+                }
+
+                // If middleware implements multiple interfaces, wrap to delegate to correct hook
+                // If it only implements one interface, still wrap to ensure hook method is called
+                return $this->wrapMiddleware($middleware, $type);
+            },
+            array_filter($this->middleware, fn(Middleware $middleware): bool => $middleware instanceof $type),
+        );
+    }
+
+    /**
+     * Check if the given type is a hook-based middleware interface.
+     *
+     * @param class-string<\Cortex\Agents\Contracts\Middleware> $type
+     */
+    protected function isHookMiddlewareType(string $type): bool
+    {
+        return in_array($type, [
+            BeforePromptMiddleware::class,
+            BeforeModelMiddleware::class,
+            AfterModelMiddleware::class,
+        ], true);
+    }
+
+    /**
+     * Wrap middleware to delegate to the appropriate hook method.
+     */
+    protected function wrapMiddleware(Middleware $middleware, string $type): Middleware
+    {
+        return match ($type) {
+            BeforePromptMiddleware::class => new BeforePromptWrapper($middleware), // @phpstan-ignore argument.type
+            BeforeModelMiddleware::class => new BeforeModelWrapper($middleware), // @phpstan-ignore argument.type
+            AfterModelMiddleware::class => new AfterModelWrapper($middleware), // @phpstan-ignore argument.type
+            default => $middleware,
+        };
+    }
+}

src/Agents/Contracts/AgentBuilder.php

@@ -8,7 +8,6 @@
 use Cortex\Contracts\ToolKit;
 use Cortex\LLM\Contracts\LLM;
 use Cortex\LLM\Enums\ToolChoice;
-use Cortex\Memory\Contracts\Store;
 use Cortex\JsonSchema\Types\ObjectSchema;
 use Cortex\LLM\Enums\StructuredOutputMode;
 use Cortex\Prompts\Builders\ChatPromptBuilder;
@@ -72,11 +71,6 @@ public function strict(): bool;
      */
     public function initialPromptVariables(): array;
 
-    /**
-     * Specify the memory store for the agent.
-     */
-    public function memoryStore(): ?Store;
-
     /**
      * Specify the middleware for the agent.
      *

src/Agents/Middleware/AfterModelWrapper.php

@@ -7,7 +7,6 @@
 use Closure;
 use Cortex\Pipeline\RuntimeConfig;
 use Cortex\Support\Traits\CanPipe;
-use Cortex\Agents\Contracts\Middleware;
 use Cortex\Agents\Contracts\AfterModelMiddleware;
 
 /**
@@ -19,15 +18,12 @@ class AfterModelWrapper implements AfterModelMiddleware
     use CanPipe;
 
     public function __construct(
-        protected Middleware $middleware,
+        protected AfterModelMiddleware $middleware,
     ) {}
 
     public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $next): mixed
     {
-        /** @var AfterModelMiddleware $middleware */
-        $middleware = $this->middleware;
-
-        return $middleware->afterModel($payload, $config, $next);
+        return $this->middleware->afterModel($payload, $config, $next);
     }
 
     public function afterModel(mixed $payload, RuntimeConfig $config, Closure $next): mixed

src/Agents/Middleware/BeforeModelWrapper.php

@@ -7,7 +7,6 @@
 use Closure;
 use Cortex\Pipeline\RuntimeConfig;
 use Cortex\Support\Traits\CanPipe;
-use Cortex\Agents\Contracts\Middleware;
 use Cortex\Agents\Contracts\BeforeModelMiddleware;
 
 /**
@@ -19,15 +18,12 @@ class BeforeModelWrapper implements BeforeModelMiddleware
     use CanPipe;
 
     public function __construct(
-        protected Middleware $middleware,
+        protected BeforeModelMiddleware $middleware,
     ) {}
 
     public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $next): mixed
     {
-        /** @var BeforeModelMiddleware $middleware */
-        $middleware = $this->middleware;
-
-        return $middleware->beforeModel($payload, $config, $next);
+        return $this->middleware->beforeModel($payload, $config, $next);
     }
 
     public function beforeModel(mixed $payload, RuntimeConfig $config, Closure $next): mixed

src/Agents/Middleware/BeforePromptWrapper.php

@@ -7,7 +7,6 @@
 use Closure;
 use Cortex\Pipeline\RuntimeConfig;
 use Cortex\Support\Traits\CanPipe;
-use Cortex\Agents\Contracts\Middleware;
 use Cortex\Agents\Contracts\BeforePromptMiddleware;
 
 /**
@@ -19,15 +18,12 @@ class BeforePromptWrapper implements BeforePromptMiddleware
     use CanPipe;
 
     public function __construct(
-        protected Middleware $middleware,
+        protected BeforePromptMiddleware $middleware,
     ) {}
 
     public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $next): mixed
     {
-        /** @var BeforePromptMiddleware $middleware */
-        $middleware = $this->middleware;
-
-        return $middleware->beforePrompt($payload, $config, $next);
+        return $this->middleware->beforePrompt($payload, $config, $next);
     }
 
     public function beforePrompt(mixed $payload, RuntimeConfig $config, Closure $next): mixed

src/Contracts/ChatMemory.php

@@ -44,4 +44,10 @@ public function setVariables(array $variables): static;
      * @return array<string, mixed>
      */
     public function getVariables(): array;
+
+    /**
+     * Get the thread ID for this memory instance.
+     * Delegates to the underlying store - the store is the source of truth for threadId.
+     */
+    public function getThreadId(): string;
 }

src/Memory/ChatMemory.php

@@ -4,6 +4,7 @@
 
 namespace Cortex\Memory;
 
+use Illuminate\Support\Str;
 use Cortex\LLM\Contracts\Message;
 use Cortex\Memory\Contracts\Store;
 use Cortex\Memory\Stores\InMemoryStore;
@@ -17,14 +18,23 @@ class ChatMemory implements ChatMemoryContract
      */
     protected array $variables = [];
 
+    protected Store $store;
+
     /**
-     * @param Store $store The store to use for message persistence
+     * @param Store|null $store The store to use for message persistence (must have threadId)
      * @param int|null $limit The maximum number of messages to return
      */
     public function __construct(
-        protected Store $store = new InMemoryStore(),
+        ?Store $store = null,
         protected ?int $limit = null,
-    ) {}
+    ) {
+        // Create store if not provided (with auto-generated threadId)
+        if ($store === null) {
+            $store = new InMemoryStore(threadId: Str::uuid7()->toString());
+        }
+
+        $this->store = $store;
+    }
 
     public function addMessage(Message $message): void
     {
@@ -84,4 +94,9 @@ public function reset(): void
         $this->store->reset();
         $this->setVariables([]);
     }
+
+    public function getThreadId(): string
+    {
+        return $this->store->getThreadId();
+    }
 }