Improved structure

Author: eldadfux

src/Agents/Adapter.php

@@ -11,6 +11,20 @@ abstract class Adapter
      */
     protected ?Agent $agent = null;
 
+    /**
+     * Input tokens count
+     *
+     * @var int
+     */
+    protected int $inputTokens = 0;
+
+    /**
+     * Output tokens count
+     *
+     * @var int
+     */
+    protected int $outputTokens = 0;
+
     /**
      * Get the adapter name
      *
@@ -21,12 +35,13 @@ abstract public function getName(): string;
     /**
      * Send a message to the AI model
      *
-     * @param  Conversation  $conversation The conversation instance containing messages and tracking tokens
+     * @param  array<Message>  $messages The messages to send to the AI model
+     * @param  callable|null  $listener The listener to call when the message is sent
      * @return Message Response from the AI model
      *
      * @throws \Exception
      */
-    abstract public function send(Conversation $conversation): Message;
+    abstract public function send(array $messages, ?callable $listener = null): Message;
 
     /**
      * Get available models for this adapter
@@ -74,4 +89,61 @@ public function setAgent(Agent $agent): self
 
         return $this;
     }
+
+
+    /**
+     * Get input tokens count
+     *
+     * @return int
+     */
+    public function getInputTokens(): int
+    {
+        return $this->inputTokens;
+    }
+
+    /**
+     * Add to input tokens count
+     *
+     * @param  int  $tokens
+     * @return self
+     */
+    public function countInputTokens(int $tokens): self
+    {
+        $this->inputTokens += $tokens;
+
+        return $this;
+    }
+
+    /**
+     * Get output tokens count
+     *
+     * @return int
+     */
+    public function getOutputTokens(): int
+    {
+        return $this->outputTokens;
+    }
+
+    /**
+     * Add to output tokens count
+     *
+     * @param  int  $tokens
+     * @return self
+     */
+    public function countOutputTokens(int $tokens): self
+    {
+        $this->outputTokens += $tokens;
+
+        return $this;
+    }
+
+    /**
+     * Get total tokens count
+     *
+     * @return int
+     */
+    public function getTotalTokens(): int
+    {
+        return $this->inputTokens + $this->outputTokens;
+    }
 }

src/Agents/Adapters/Anthropic.php

@@ -6,7 +6,6 @@
 use Utopia\Agents\Conversation;
 use Utopia\Agents\Message;
 use Utopia\Agents\Messages\Text;
-use Utopia\Agents\Roles\Assistant;
 use Utopia\Fetch\Chunk;
 use Utopia\Fetch\Client;
 
@@ -77,12 +76,13 @@ public function __construct(
     /**
      * Send a message to the Anthropic API
      *
-     * @param  Conversation  $conversation
+     * @param  array<Message>  $messages
+     * @param  callable|null  $listener
      * @return Message
      *
      * @throws \Exception
      */
-    public function send(Conversation $conversation): Message
+    public function send(array $messages, ?callable $listener = null): Message
     {
         if ($this->getAgent() === null) {
             throw new \Exception('Agent not set');
@@ -94,8 +94,7 @@ public function send(Conversation $conversation): Message
             ->addHeader('anthropic-version', '2023-06-01')
             ->addHeader('content-type', 'application/json');
 
-        $messages = [];
-        foreach ($conversation->getMessages() as $message) {
+        foreach ($messages as $message) {
             $messages[] = [
                 'role' => $message['role'],
                 'content' => $message['content'],
@@ -121,8 +120,8 @@ public function send(Conversation $conversation): Message
                 'stream' => true,
             ],
             [],
-            function ($chunk) use ($conversation, &$content) {
-                $content .= $this->process($chunk, $conversation, $conversation->getListener());
+            function ($chunk) use (&$content, $listener) {
+                $content .= $this->process($chunk, $listener);
             }
         );
 
@@ -145,7 +144,7 @@ function ($chunk) use ($conversation, &$content) {
      *
      * @throws \Exception
      */
-    protected function process(Chunk $chunk, Conversation $conversation, ?callable $listener): string
+    protected function process(Chunk $chunk, ?callable $listener): string
     {
         $block = '';
         $data = $chunk->getData();
@@ -176,10 +175,10 @@ protected function process(Chunk $chunk, Conversation $conversation, ?callable $
                     if (isset($json['message']['usage'])) {
                         $usage = $json['message']['usage'];
                         if (isset($usage['input_tokens']) && is_int($usage['input_tokens'])) {
-                            $conversation->countInputTokens($usage['input_tokens']);
+                            $this->countInputTokens($usage['input_tokens']);
                         }
                         if (isset($usage['output_tokens']) && is_int($usage['output_tokens'])) {
-                            $conversation->countOutputTokens($usage['output_tokens']);
+                            $this->countOutputTokens($usage['output_tokens']);
                         }
                     }
                     break;
@@ -211,13 +210,13 @@ protected function process(Chunk $chunk, Conversation $conversation, ?callable $
                     break;
 
                 case 'message_delta':
-                    if (isset($json['message']['usage'])) {
-                        $usage = $json['message']['usage'];
+                    if (isset($json['usage'])) {
+                        $usage = $json['usage'];
                         if (isset($usage['input_tokens']) && is_int($usage['input_tokens'])) {
-                            $conversation->countInputTokens($usage['input_tokens']);
+                            $this->countInputTokens($usage['input_tokens']);
                         }
                         if (isset($usage['output_tokens']) && is_int($usage['output_tokens'])) {
-                            $conversation->countOutputTokens($usage['output_tokens']);
+                            $this->countOutputTokens($usage['output_tokens']);
                         }
                     }
                     break;

src/Agents/Conversation.php

@@ -85,7 +85,11 @@ public function message(Role $from, Message $message): self
      */
     public function send(): Message
     {
-        $message = $this->agent->getAdapter()->send($this);
+        $message = $this->agent->getAdapter()->send($this->messages, $this->listener);
+        
+        $this->countInputTokens($this->agent->getAdapter()->getInputTokens());
+        $this->countOutputTokens($this->agent->getAdapter()->getOutputTokens());
+
         $from = new Assistant($this->agent->getAdapter()->getModel(), 'Assistant');
         $this->message($from, $message);