anthropic support

package.json

@@ -24,7 +24,6 @@
   "author": "Doku Labs",
   "license": "ISC",
   "dependencies": {
-    "@anthropic-ai/tokenizer": "^0.0.4",
     "stream": "^0.0.2"
   },
   "devDependencies": {

src/anthropic.js

@@ -1,5 +1,5 @@
 import {sendData} from './helpers.js';
-import {countTokens} from '@anthropic-ai/tokenizer';
+import { Readable } from 'stream';
 
 /**
  * Initializes Anthropic functionality with performance tracking.
@@ -31,33 +31,138 @@ import {countTokens} from '@anthropic-ai/tokenizer';
  * }
  */
 export default function initAnthropic({ llm, dokuUrl, apiKey, environment, applicationName, skipResp }) {
-  const originalCompletionsCreate = llm.completions.create;
+  const originalMessagesCreate = llm.messages.create;
 
   // Define wrapped method
-  llm.completions.create = async function(params) {
-    const start = performance.now();
-    const response = await originalCompletionsCreate.apply(this, params);
-    const end = performance.now();
-    const duration = (end - start) / 1000;
-
-    const data = {
-      llmReqId: response.id,
-      environment: environment,
-      applicationName: applicationName,
-      sourceLanguage: 'Javascript',
-      endpoint: 'anthropic.completions',
-      skipResp: skipResp,
-      completionTokens: countTokens(response.completion),
-      promptTokens: countTokens(prompt),
-      requestDuration: duration,
-      model: params.model,
-      prompt: params.prompt,
-      finishReason: response.stop_reason,
-      response: response.completion,
-    };
-
-    await sendData(data, dokuUrl, apiKey);
-
-    return response;
+  llm.messages.create = async function(params) {
+    let streaming = params.stream || false;
+    if (streaming) {
+      // Call original method
+      const start = performance.now();
+      const originalResponseStream = await originalMessagesCreate.call(this, params);
+
+      // Create a pass-through stream
+      const passThroughStream = new Readable({
+        read() {},
+        objectMode: true // Set to true because the chunks are objects
+      });
+
+      let dataResponse = '';
+      var responseId = '';
+      var promptTokens = 0;
+      var completionTokens = 0;
+
+      // Immediately-invoked async function to handle streaming
+      (async () => {
+        for await (const chunk of originalResponseStream) {
+          if (chunk.type === 'message_start') {
+            responseId = chunk.message.id;
+            promptTokens = chunk.message.usage.input_tokens;
+            passThroughStream.push(chunk); // Push chunk to the pass-through stream
+          } 
+          else if (chunk.type === 'content_block_delta') {
+            dataResponse += chunk.delta.text;
+            passThroughStream.push(chunk); // Push chunk to the pass-through stream
+          } 
+          else if (chunk.type === 'message_delta') {
+            completionTokens = chunk.usage.output_tokens;
+            passThroughStream.push(chunk); // Push chunk to the pass-through stream
+          }
+        }
+        passThroughStream.push(null); // Signal end of the pass-through stream
+
+        // Process response data after stream has ended
+        const end = performance.now();
+        const duration = (end - start) / 1000;
+
+        let formattedMessages = [];
+        for (let message of params.messages) {
+          let role = message.role;
+          let content = message.content;
+
+          if (Array.isArray(content)) {
+            let contentStr = content.map(item => {
+              if (item.type) {
+                return `${item.type}: ${item.text || item.image_url}`;
+              } else {
+                return `text: ${item.text}`;
+              }
+            }).join(", ");
+            formattedMessages.push(`${role}: ${contentStr}`);
+          } else {
+            formattedMessages.push(`${role}: ${content}`);
+          }
+        }
+        let prompt = formattedMessages.join("\n");
+
+        // Prepare the data object for Doku
+        const data = {
+          llmReqId: responseId,
+          environment: environment,
+          applicationName: applicationName,
+          sourceLanguage: 'Javascript',
+          endpoint: 'anthropic.messages',
+          skipResp: skipResp,
+          requestDuration: duration,
+          model: params.model,
+          prompt: prompt,
+          response: dataResponse,
+          promptTokens: promptTokens,
+          completionTokens: completionTokens,
+        };
+        data.totalTokens = data.promptTokens + data.completionTokens;
+
+        await sendData(data, dokuUrl, apiKey);
+      })();
+
+      // Return the pass-through stream to the original caller
+      return passThroughStream;
+    }
+    else{
+      const start = performance.now();
+      const response = await originalMessagesCreate.call(this, params);
+      const end = performance.now();
+      const duration = (end - start) / 1000;
+      let formattedMessages = [];
+      for (let message of params.messages) {
+        let role = message.role;
+        let content = message.content;
+
+        if (Array.isArray(content)) {
+          let contentStr = content.map(item => {
+            if (item.type) {
+              return `${item.type}: ${item.text || item.image_url}`;
+            } else {
+              return `text: ${item.text}`;
+            }
+          }).join(", ");
+          formattedMessages.push(`${role}: ${contentStr}`);
+        } else {
+          formattedMessages.push(`${role}: ${content}`);
+        }
+      }
+      let prompt = formattedMessages.join("\n");
+
+      const data = {
+        llmReqId: response.id,
+        environment: environment,
+        applicationName: applicationName,
+        sourceLanguage: 'Javascript',
+        endpoint: 'anthropic.messages',
+        skipResp: skipResp,
+        completionTokens: response.usage.input_tokens,
+        promptTokens: response.usage.output_tokens,
+        requestDuration: duration,
+        model: params.model,
+        prompt: prompt,
+        finishReason: response.stop_reason,
+        response: response.content[0].text,
+      };
+      data.totalTokens = data.promptTokens + data.completionTokens; 
+
+      await sendData(data, dokuUrl, apiKey);
+
+      return response;
+    }
   };
 }

src/index.js

@@ -56,7 +56,7 @@ function init({ llm, dokuUrl, apiKey, environment="default", applicationName="de
     initOpenAI({ llm, dokuUrl, apiKey, environment, applicationName, skipResp });
   } else if (llm.generate && typeof llm.rerank === 'function') {
     initCohere({ llm, dokuUrl, apiKey, environment, applicationName, skipResp });
-  } else if (typeof llm.summarize=== 'function') {
+  } else if (llm.messages && typeof llm.messages.create === 'function') {
     initAnthropic({ llm, dokuUrl, apiKey, environment, applicationName, skipResp });
   }
 }