Skip to content

Commit b051882

Browse files
tymondesignstymondesigns
committed
wip
1 parent 62d0638 commit b051882

File tree

8 files changed

+120
-31
lines changed

8 files changed

+120
-31
lines changed

‎src/Agents/Agent.php‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,13 @@
2626
use Cortex\Memory\Contracts\Store;
2727
use Cortex\Pipeline\RuntimeConfig;
2828
use Cortex\Support\Traits\CanPipe;
29+
use Cortex\LLM\Data\ChatGeneration;
2930
use Cortex\Agents\Stages\AppendUsage;
3031
use Cortex\LLM\Data\ChatStreamResult;
3132
use Cortex\Events\Contracts\AgentEvent;
3233
use Cortex\Exceptions\GenericException;
3334
use Cortex\Memory\Stores\InMemoryStore;
35+
use Cortex\LLM\Data\ChatGenerationChunk;
3436
use Cortex\Agents\Stages\HandleToolCalls;
3537
use Cortex\JsonSchema\Types\ObjectSchema;
3638
use Cortex\LLM\Enums\StructuredOutputMode;
@@ -57,8 +59,6 @@ class Agent implements Pipeable
5759

5860
protected ChatMemoryContract $memory;
5961

60-
protected Usage $usage;
61-
6262
protected ObjectSchema|string|null $output = null;
6363

6464
protected Pipeline $pipeline;
@@ -97,7 +97,6 @@ public function __construct(
9797
$this->outputMode,
9898
$this->strict,
9999
);
100-
$this->usage = Usage::empty();
101100
$this->pipeline = $this->pipeline();
102101
}
103102

@@ -126,7 +125,7 @@ public function executionPipeline(bool $shouldParseOutput = true): Pipeline
126125
->pipe(function (mixed $payload, RuntimeConfig $config, Closure $next): mixed {
127126
// Only add step 1 if no steps exist yet (first invocation as a stage)
128127
// Subsequent invocations from HandleToolCalls will have already added the step
129-
if ($config->context->getSteps()->isEmpty()) {
128+
if (! $config->context->hasSteps()) {
130129
$config->context->addStep(new Step(number: 1));
131130
}
132131

@@ -136,8 +135,21 @@ public function executionPipeline(bool $shouldParseOutput = true): Pipeline
136135
})
137136
->pipe($this->llm->shouldParseOutput($shouldParseOutput))
138137
->pipe(new AddMessageToMemory($this->memory))
139-
->pipe(new AppendUsage($this->usage))
138+
->pipe(new AppendUsage())
140139
->pipe(function (mixed $payload, RuntimeConfig $config, Closure $next): mixed {
140+
// Extract generation from payload to check for tool calls
141+
$generation = match (true) {
142+
$payload instanceof ChatGeneration => $payload,
143+
$payload instanceof ChatGenerationChunk && $payload->isFinal => $payload,
144+
$payload instanceof ChatResult => $payload->generation,
145+
default => null,
146+
};
147+
148+
// Set tool calls on the current step if the generation has them
149+
if ($generation !== null && $generation->message->hasToolCalls()) {
150+
$config->context->getCurrentStep()->setToolCalls($generation->message->toolCalls);
151+
}
152+
141153
$this->dispatchEvent(new AgentStepEnd($this, $config));
142154

143155
return $next($payload, $config);
@@ -236,7 +248,7 @@ public function getMemory(): ChatMemoryContract
236248

237249
public function getUsage(): Usage
238250
{
239-
return $this->usage;
251+
return $this->runtimeConfig?->context?->getUsageSoFar() ?? Usage::empty();
240252
}
241253

242254
public function getRuntimeConfig(): ?RuntimeConfig

‎src/Agents/Stages/AddMessageToMemory.php‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $n
3232

3333
if ($message !== null) {
3434
$this->memory->addMessage($message);
35-
$config->context->set('message_history', $this->memory->getMessages());
35+
$config->context->setMessageHistory($this->memory->getMessages());
3636
}
3737

3838
return $next($payload, $config);

‎src/Agents/Stages/AppendUsage.php‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,6 @@ class AppendUsage implements Pipeable
1616
{
1717
use CanPipe;
1818

19-
public function __construct(
20-
protected Usage $usage,
21-
) {}
22-
2319
public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $next): mixed
2420
{
2521
$usage = match (true) {
@@ -28,10 +24,11 @@ public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $n
2824
};
2925

3026
if ($usage !== null) {
31-
$this->usage->add($usage);
32-
$config->context->getCurrentStep()->setUsage($this->usage);
33-
// $config->context->set('usage', $this->usage->toArray());
34-
// $config->context->set('usage_total', $this->usage->add($usage)->toArray());
27+
// Set the usage for the current step
28+
$config->context->getCurrentStep()->setUsage($usage);
29+
30+
// Append the usage to the context so we can track usage as we move through the steps.
31+
$config->context->appendUsage($usage);
3532
}
3633

3734
return $next($payload, $config);

‎src/Agents/Stages/HandleToolCalls.php‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ public function handlePipeable(mixed $payload, RuntimeConfig $config, Closure $n
5454
// Track the next step before making the LLM call
5555
$config->context->addStep(
5656
new Step(
57-
number: $config->context->getCurrentStep()->number + 1,
57+
number: $config->context->getCurrentStepNumber() + 1,
5858
toolCalls: new ToolCallCollection(),
5959
),
6060
);

‎src/Http/Controllers/AgentsController.php‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
use Illuminate\Http\Request;
1111
use Cortex\Events\AgentStepEnd;
1212
use Cortex\Events\AgentStepError;
13-
use Cortex\Events\AgentStepStart;
1413
use Illuminate\Http\JsonResponse;
1514
use Illuminate\Routing\Controller;
1615
use Symfony\Component\HttpFoundation\StreamedResponse;
@@ -21,20 +20,20 @@ public function invoke(string $agent, Request $request): JsonResponse
2120
{
2221
try {
2322
$agent = Cortex::agent($agent);
24-
$agent->onStepStart(function (AgentStepStart $event): void {
25-
dump(
26-
sprintf('step start: %d', $event->config?->context->getCurrentStep()->number),
27-
$event->config?->context->toArray(),
28-
);
29-
});
23+
// $agent->onStepStart(function (AgentStepStart $event): void {
24+
// dump(
25+
// sprintf('step start: %d', $event->config?->getCurrentStepNumber()),
26+
// $event->config?->context->toArray(),
27+
// );
28+
// });
3029
$agent->onStepEnd(function (AgentStepEnd $event): void {
3130
dump(
32-
sprintf('step end: %d', $event->config?->context->getCurrentStep()->number),
33-
$event->config?->context->toArray(),
31+
sprintf('step end: %d', $event->config?->context?->getCurrentStepNumber()),
32+
$event->config?->toArray(),
3433
);
3534
});
3635
$agent->onStepError(function (AgentStepError $event): void {
37-
dump(sprintf('step error: %d', $event->config?->context->getCurrentStep()->number));
36+
dump(sprintf('step error: %d', $event->config?->context?->getCurrentStepNumber()));
3837
dump($event->exception->getMessage());
3938
dump($event->exception->getTraceAsString());
4039
});

‎src/Pipeline/Context.php‎

Lines changed: 83 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,45 +4,124 @@
44

55
namespace Cortex\Pipeline;
66

7+
use Cortex\LLM\Data\Usage;
78
use Cortex\Agents\Data\Step;
89
use Illuminate\Support\Fluent;
910
use Illuminate\Support\Collection;
11+
use Cortex\LLM\Data\Messages\MessageCollection;
1012

1113
/**
1214
* @extends Fluent<string, mixed>
1315
*/
1416
class Context extends Fluent
1517
{
18+
public const string STEPS_KEY = 'steps';
19+
public const string CURRENT_STEP_KEY = 'current_step';
20+
public const string USAGE_SO_FAR_KEY = 'usage_so_far';
21+
public const string MESSAGE_HISTORY_KEY = 'message_history';
22+
1623
/**
24+
* Get the steps from the context.
25+
*
1726
* @return \Illuminate\Support\Collection<int, \Cortex\Agents\Data\Step>
1827
*/
1928
public function getSteps(): Collection
2029
{
21-
return $this->get('steps', new Collection());
30+
return $this->get(self::STEPS_KEY, new Collection());
31+
}
32+
33+
/**
34+
* Determines if the context has any steps.
35+
*/
36+
public function hasSteps(): bool
37+
{
38+
return $this->getSteps()->isNotEmpty();
2239
}
2340

2441
/**
42+
* Set the steps in the context.
43+
*
2544
* @param \Illuminate\Support\Collection<int, \Cortex\Agents\Data\Step> $steps
2645
*/
2746
public function setSteps(Collection $steps): void
2847
{
29-
$this->set('steps', $steps);
48+
$this->set(self::STEPS_KEY, $steps);
3049
}
3150

51+
/**
52+
* Get the current step from the context.
53+
*/
3254
public function getCurrentStep(): Step
3355
{
34-
return $this->get('current_step', new Step(number: 1));
56+
return $this->get(self::CURRENT_STEP_KEY, new Step(number: 1));
57+
}
58+
59+
/**
60+
* Get the current step number from the context.
61+
*/
62+
public function getCurrentStepNumber(): int
63+
{
64+
return $this->getCurrentStep()->number;
3565
}
3666

67+
/**
68+
* Set the current step in the context.
69+
*/
3770
public function setCurrentStep(Step $step): void
3871
{
39-
$this->set('current_step', $step);
72+
$this->set(self::CURRENT_STEP_KEY, $step);
4073
}
4174

75+
/**
76+
* Add a step to the context.
77+
*/
4278
public function addStep(Step $step): void
4379
{
4480
$steps = $this->getSteps()->push($step);
4581
$this->setCurrentStep($step);
4682
$this->setSteps($steps);
4783
}
84+
85+
/**
86+
* Get the usage so far from the context.
87+
*/
88+
public function getUsageSoFar(): Usage
89+
{
90+
return $this->get(self::USAGE_SO_FAR_KEY, Usage::empty());
91+
}
92+
93+
/**
94+
* Set the usage so far in the context.
95+
*/
96+
public function setUsageSoFar(Usage $usage): void
97+
{
98+
$this->set(self::USAGE_SO_FAR_KEY, $usage);
99+
}
100+
101+
/**
102+
* Append usage to the context.
103+
*/
104+
public function appendUsage(Usage $usage): static
105+
{
106+
$usageSoFar = $this->getUsageSoFar()->add($usage);
107+
$this->setUsageSoFar($usageSoFar);
108+
109+
return $this;
110+
}
111+
112+
/**
113+
* Get the message history from the context.
114+
*/
115+
public function getMessageHistory(): MessageCollection
116+
{
117+
return $this->get(self::MESSAGE_HISTORY_KEY, new MessageCollection());
118+
}
119+
120+
/**
121+
* Set the message history in the context.
122+
*/
123+
public function setMessageHistory(MessageCollection $messages): void
124+
{
125+
$this->set(self::MESSAGE_HISTORY_KEY, $messages);
126+
}
48127
}

‎src/Pipeline/RuntimeConfig.php‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
namespace Cortex\Pipeline;
66

77
use Illuminate\Support\Str;
8+
use Cortex\Agents\Data\Step;
9+
use Illuminate\Support\Collection;
810
use Illuminate\Contracts\Support\Arrayable;
911

1012
/**

‎tests/Unit/Agents/AgentTest.php‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -264,5 +264,5 @@ function (int $x, int $y): int {
264264
->and($step1->toolCalls)->toHaveCount(0);
265265

266266
// Verify current_step is set
267-
expect($runtimeConfig->context->getCurrentStep()->number)->toBe(1);
267+
expect($runtimeConfig->context->getCurrentStepNumber())->toBe(1);
268268
});

0 commit comments

Comments
 (0)