fix(agents): get exact token calculation (#22)

Author: ijusttookadnatest

README.md

@@ -44,7 +44,6 @@ Cairo Coder is an intelligent code generation service that makes writing Cairo s
 - **Multiple LLM Support**: Works with OpenAI, Anthropic, and Google models
 - **Source-Informed Generation**: Code is generated based on Cairo documentation, ensuring correctness
 
-
 ## Installation
 
 There are mainly 2 ways of installing Cairo Coder - With Docker, Without Docker. Using Docker is highly recommended.
@@ -68,7 +67,6 @@ There are mainly 2 ways of installing Cairo Coder - With Docker, Without Docker.
    pnpm install
    ```
 
-
 5. Inside the packages/agents package, copy the `sample.config.toml` file to a `config.toml`. For development setups, you need only fill in the following fields:
 
    - `OPENAI`: Your OpenAI API key. **You only need to fill this if you wish to use OpenAI's models**.
@@ -115,45 +113,45 @@ There are mainly 2 ways of installing Cairo Coder - With Docker, Without Docker.
    ```
 
    This configuration is used by the backend and ingester services to connect to the database.
-   Note that `POSTGRES_HOST` is set to ```"postgres"``` and `POSTGRES_PORT` to ```"5432"```, which are the container's name and port in docker-compose.yml.
+   Note that `POSTGRES_HOST` is set to `"postgres"` and `POSTGRES_PORT` to `"5432"`, which are the container's name and port in docker-compose.yml.
 
    **Important:** Make sure to use the same password, username and db's name in both files. The first file initializes the database, while the second is used by your application to connect to it.
 
-
 7. **Configure LangSmith (Optional)**
 
    Cairo Coder can use LangSmith to record and monitor LLM calls. This step is optional but recommended for development and debugging.
-   
+
    - Create an account at [LangSmith](https://smith.langchain.com/)
    - Create a new project in the LangSmith dashboard
    - Retrieve your API credentials
    - Create a `.env` file in the `packages/backend` directory with the following variables:
+
    ```
-   LANGSMITH_TRACING=true
-   LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
-   LANGSMITH_API_KEY="<your-api-key>"
+   LANGCHAIN_TRACING=true
+   LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
+   LANGCHAIN_API_KEY="<your-api-key>"
    LANGCHAIN_PROJECT="<your-project-name>"
    ```
-   - Add the `.env` in an env_file section in the backend service of the docker-compose.yml 
 
-   With this configuration, all LLM calls and chain executions will be logged to your LangSmith project, allowing you to debug, analyze, and improve the system's performance.
+   - Add the `packages/backend/.env` in an env_file section in the backend service of the docker-compose.yml
 
+   With this configuration, all LLM calls and chain executions will be logged to your LangSmith project, allowing you to debug, analyze, and improve the system's performance.
 
-9. Run the application using one of the following methods:
+8. Run the application using one of the following methods:
 
    ```bash
    docker compose up postgres backend
    ```
 
-8. The API will be available at http://localhost:3001/v1/chat/completions
+9. The API will be available at http://localhost:3001/v1/chat/completions
 
 ## Running the Ingester
 
 After you have the main application running, you might need to run the ingester to process and embed documentation from various sources. The ingester is configured as a separate profile in the docker-compose file and can be executed as follows:
 
-   ```bash
-   docker compose up ingester
-   ```
+```bash
+docker compose up ingester
+```
 
 Once the ingester completes its task, the vector database will be populated with embeddings from all the supported documentation sources, making them available for RAG-based code generation requests to the API.
 
@@ -188,6 +186,7 @@ curl -X POST http://localhost:3001/v1/chat/completions \
 The API accepts all standard OpenAI Chat Completions parameters.
 
 **Supported Parameters:**
+
 - `model`: Model identifier (string)
 - `messages`: Array of message objects with `role` and `content`
 - `temperature`: Controls randomness (0-2, default: 0.7)
@@ -202,7 +201,6 @@ The API accepts all standard OpenAI Chat Completions parameters.
 - `user`: User identifier
 - `response_format`: Response format specification
 
-
 ### Response Format
 
 #### Standard Mode Response

packages/agents/__tests__/answerGenerator.test.ts

@@ -10,8 +10,9 @@ import {
 import { Document } from '@langchain/core/documents';
 import { mockDeep, MockProxy } from 'jest-mock-extended';
 import { IterableReadableStream } from '@langchain/core/utils/stream';
-import { BaseMessage, BaseMessageChunk } from '@langchain/core/messages';
+import { StreamEvent } from '@langchain/core/tracers/log_stream';
 import { BaseLanguageModelInput } from '@langchain/core/language_models/base';
+import { AIMessageChunk } from '@langchain/core/messages'; // ✅ AJOUTÉ
 
 // Mock the formatChatHistoryAsString utility
 jest.mock('../src/utils/index', () => ({
@@ -61,23 +62,21 @@ describe('AnswerGenerator', () => {
       testTemplate: '<test_template>Example test template</test_template>',
     };
 
-    // Mock the LLM stream method to return simulated stream
-    mockLLM.stream.mockImplementation(
-      (
-        input: BaseLanguageModelInput,
-      ): Promise<IterableReadableStream<BaseMessageChunk>> => {
-        return Promise.resolve({
-          [Symbol.asyncIterator]: async function* () {
-            yield {
-              content: 'This is a test response about Cairo.',
-              type: 'ai',
-              name: 'AI',
-              additional_kwargs: {},
-            };
+    // ✅ CHANGÉ: streamEvents au lieu de stream
+    (mockLLM.streamEvents as any) = jest.fn().mockReturnValue({
+      [Symbol.asyncIterator]: async function* () {
+        yield {
+          event: 'on_llm_stream',
+          data: {
+            chunk: { content: 'This is a test response about Cairo.' },
           },
-        } as unknown as IterableReadableStream<BaseMessageChunk>);
+          run_id: 'test-run-123',
+          name: 'TestLLM',
+          tags: [],
+          metadata: {},
+        } as StreamEvent;
       },
-    );
+    } as IterableReadableStream<StreamEvent>);
 
     // Create the AnswerGenerator instance
     answerGenerator = new AnswerGenerator(mockLLM, mockConfig);
@@ -122,17 +121,20 @@ describe('AnswerGenerator', () => {
       const result = await answerGenerator.generate(input, retrievedDocs);
 
       // Assert
-      expect(mockLLM.stream).toHaveBeenCalled();
+      expect(mockLLM.streamEvents).toHaveBeenCalled(); // ✅ CHANGÉ: streamEvents
 
       // Collect the stream results
-      const messages: any[] = [];
-      for await (const message of result) {
-        messages.push(message);
+      const events: StreamEvent[] = [];
+      for await (const event of result) {
+        events.push(event);
       }
 
-      // Check that we got the expected message
-      expect(messages.length).toBe(1);
-      expect(messages[0].content).toBe('This is a test response about Cairo.');
+      // Check that we got the expected events
+      expect(events.length).toBe(1);
+      expect(events[0].event).toBe('on_llm_stream');
+      expect(events[0].data.chunk.content).toBe(
+        'This is a test response about Cairo.',
+      );
     });
 
     it('should include contract template when query is contract-related', async () => {
@@ -169,27 +171,39 @@ describe('AnswerGenerator', () => {
         processedQuery,
       };
 
-      // Create a spy on the LLM stream method to capture the actual prompt
+      // Create a spy on the LLM streamEvents method to capture the actual prompt
       let capturedPrompt: string | null = null;
-      mockLLM.stream.mockImplementation((prompt) => {
-        capturedPrompt = prompt as string;
-        return Promise.resolve({
-          [Symbol.asyncIterator]: async function* () {
-            yield {
-              content: 'This is a test response about Cairo contracts.',
-              type: 'ai',
-              name: 'AI',
-              additional_kwargs: {},
-            };
-          },
-        } as unknown as IterableReadableStream<BaseMessageChunk>);
-      });
+      const mockReturnValue = {
+        [Symbol.asyncIterator]: async function* () {
+          yield {
+            event: 'on_llm_stream',
+            data: {
+              chunk: {
+                content: 'This is a test response about Cairo contracts.',
+              },
+            },
+            run_id: 'test-run-123',
+            name: 'TestLLM',
+            tags: [],
+            metadata: {},
+          } as StreamEvent;
+        },
+      } as IterableReadableStream<StreamEvent>;
+
+      (mockLLM.streamEvents as any) = jest
+        .fn()
+        .mockImplementation((...args) => {
+          capturedPrompt = args[0] as string;
+          return mockReturnValue;
+        });
 
       // Act
       await answerGenerator.generate(input, retrievedDocs);
 
       // Assert
-      expect(capturedPrompt).toContain(
+      expect(capturedPrompt).toBeDefined();
+      const promptString = JSON.stringify(capturedPrompt);
+      expect(promptString).toContain(
         '<contract_template>Example template</contract_template>',
       );
     });
@@ -229,27 +243,40 @@ describe('AnswerGenerator', () => {
         processedQuery,
       };
 
-      // Create a spy on the LLM stream method to capture the actual prompt
+      // Create a spy on the LLM streamEvents method to capture the actual prompt
       let capturedPrompt: string | null = null;
-      mockLLM.stream.mockImplementation((prompt) => {
-        capturedPrompt = prompt as string;
-        return Promise.resolve({
-          [Symbol.asyncIterator]: async function* () {
-            yield {
-              content: 'This is a test response about testing Cairo contracts.',
-              type: 'ai',
-              name: 'AI',
-              additional_kwargs: {},
-            };
-          },
-        } as unknown as IterableReadableStream<BaseMessageChunk>);
-      });
+      const mockReturnValue = {
+        [Symbol.asyncIterator]: async function* () {
+          yield {
+            event: 'on_llm_stream',
+            data: {
+              chunk: {
+                content:
+                  'This is a test response about testing Cairo contracts.',
+              },
+            },
+            run_id: 'test-run-123',
+            name: 'TestLLM',
+            tags: [],
+            metadata: {},
+          } as StreamEvent;
+        },
+      } as IterableReadableStream<StreamEvent>;
+
+      (mockLLM.streamEvents as any) = jest
+        .fn()
+        .mockImplementation((...args) => {
+          capturedPrompt = args[0] as string;
+          return mockReturnValue;
+        });
 
       // Act
       await answerGenerator.generate(input, retrievedDocs);
 
       // Assert
-      expect(capturedPrompt).toContain(
+      expect(capturedPrompt).toBeDefined();
+      const promptString = JSON.stringify(capturedPrompt);
+      expect(promptString).toContain(
         '<test_template>Example test template</test_template>',
       );
     });
@@ -273,27 +300,37 @@ describe('AnswerGenerator', () => {
         processedQuery,
       };
 
-      // Create a spy on the LLM stream method to capture the actual prompt
+      // Create a spy on the LLM streamEvents method to capture the actual prompt
       let capturedPrompt: string | null = null;
-      mockLLM.stream.mockImplementation((prompt) => {
-        capturedPrompt = prompt as string;
-        return Promise.resolve({
-          [Symbol.asyncIterator]: async function* () {
-            yield {
-              content: 'I cannot find any relevant information.',
-              type: 'ai',
-              name: 'AI',
-              additional_kwargs: {},
-            };
-          },
-        } as unknown as IterableReadableStream<BaseMessageChunk>);
-      });
+      const mockReturnValue = {
+        [Symbol.asyncIterator]: async function* () {
+          yield {
+            event: 'on_llm_stream',
+            data: {
+              chunk: { content: 'I cannot find any relevant information.' },
+            },
+            run_id: 'test-run-123',
+            name: 'TestLLM',
+            tags: [],
+            metadata: {},
+          } as StreamEvent;
+        },
+      } as IterableReadableStream<StreamEvent>;
+
+      (mockLLM.streamEvents as any) = jest
+        .fn()
+        .mockImplementation((...args) => {
+          capturedPrompt = args[0] as string;
+          return mockReturnValue;
+        });
 
       // Act
       await answerGenerator.generate(input, retrievedDocs);
 
       // Assert
-      expect(capturedPrompt).toContain(
+      expect(capturedPrompt).toBeDefined();
+      const promptString = JSON.stringify(capturedPrompt);
+      expect(promptString).toContain(
         'Sorry, no relevant information was found.',
       );
     });

packages/agents/__tests__/ragPipeline.test.ts

@@ -1,3 +1,4 @@
+import { StreamEvent } from '@langchain/core/tracers/log_stream';
 import { RagPipeline } from '../src/core/pipeline/ragPipeline';
 import { QueryProcessor } from '../src/core/pipeline/queryProcessor';
 import { DocumentRetriever } from '../src/core/pipeline/documentRetriever';
@@ -13,8 +14,8 @@ import {
 import { Document } from '@langchain/core/documents';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { IterableReadableStream } from '@langchain/core/utils/stream';
-import { BaseMessage, AIMessage } from '@langchain/core/messages';
 import { mockDeep, MockProxy } from 'jest-mock-extended';
+import { AIMessage } from '@langchain/core/messages';
 import EventEmitter from 'events';
 
 // Mock the dependencies at the module level
@@ -133,10 +134,20 @@ describe('RagPipeline', () => {
       // Create a mock response stream
       const mockStream = {
         [Symbol.asyncIterator]: async function* () {
-          yield new AIMessage('This is a test answer about Cairo contracts.');
+          yield {
+            event: 'on_llm_stream',
+            data: {
+              chunk: {
+                content: 'This is a test answer about Cairo contracts.',
+              },
+            },
+            run_id: 'test-run-123',
+            name: 'TestLLM',
+            tags: [],
+            metadata: {},
+          } as StreamEvent;
         },
-      } as IterableReadableStream<BaseMessage>;
-
+      } as IterableReadableStream<StreamEvent>;
       // Setup mock behavior
       mockQueryProcessor.process.mockResolvedValue(processedQuery);
       mockDocumentRetriever.retrieve.mockResolvedValue(retrievedDocs);