Enhance Token Counting to Include Chat History and System Instructions

src/ChatSession.php

@@ -7,6 +7,7 @@
 use GeminiAPI\Enums\Role;
 use GeminiAPI\Resources\Content;
 use GeminiAPI\Resources\Parts\PartInterface;
+use GeminiAPI\Responses\CountTokensResponse;
 use GeminiAPI\Responses\GenerateContentResponse;
 use GeminiAPI\Traits\ArrayTypeValidator;
 use InvalidArgumentException;
@@ -98,4 +99,19 @@ public function withHistory(array $history): self
 
         return $clone;
     }
+
+    /**
+     * @throws ClientExceptionInterface
+     */
+    public function countTokens(PartInterface ...$parts): CountTokensResponse
+    {
+        // Concatenate the parts of the history if it exists
+        if (!empty($this->history)) {
+            foreach ($this->history as $content) {
+                $parts = array_merge($content->parts, $parts);
+            }
+        }
+
+        return $this->model->countTokens(...$parts);
+    }
 }

src/GenerativeModel.php

@@ -115,12 +115,17 @@ public function startChat(): ChatSession
      */
     public function countTokens(PartInterface ...$parts): CountTokensResponse
     {
-        $content = new Content($parts, Role::User);
+        $temporaryContent = new Content($parts, Role::User);
+
+        // Add system instruction if it exists
+        if ($this->systemInstruction !== null) {
+            $temporaryContent = $temporaryContent->addParts(...$this->systemInstruction->parts);
+        }
+
         $request = new CountTokensRequest(
             $this->modelName,
-            [$content],
+            [$temporaryContent]
         );
-
         return $this->client->countTokens($request);
     }
 

src/Resources/Content.php

@@ -48,6 +48,17 @@ public function addFile(MimeType $mimeType, string $file): self
         return $this;
     }
 
+    public function addParts(PartInterface ...$parts): self
+    {
+        /** @var array<int, PartInterface> $typedParts */
+        $typedParts = $parts;
+
+        $this->ensureArrayOfType($typedParts, PartInterface::class);
+        $this->parts = array_merge($this->parts, $typedParts);
+
+        return $this;
+    }
+
     public static function text(
         string $text,
         Role $role = Role::User,