Merge pull request #4 from mrsteele/feature/configFix

feat: adding the options field instead of limits

README.md

@@ -42,6 +42,10 @@ const truncatedByNumber = truncateMessage(str, 100)
 // enforce truncation around all messages
 const truncatedBody = truncateWrapper({
   model: 'gpt-4', // auto-detects token limits 🙌
+  //optionally, you can supply your own limit (surpressed in output)
+  opts: {
+    limit: 1000
+  },
   messages: [{ role: 'user', content: str }]
 })
 ```

src/truncate.js

@@ -47,12 +47,10 @@ const truncateCompletion = (body = {}, limit) => {
   // if its good, just send it off
   // console.log('forceLimit', getTokens(body.messages[0].content))
   // return forceLimit
-  if (runningTotal < forceLimit) {
+  if (runningTotal <= forceLimit) {
     return body
   }
 
-  // ...otherwise
-  // find last culprit, everything else will be removed...
   const bigIndex = newMessages.findIndex(m => m.runningTotal > forceLimit)
   const newLimit = forceLimit - newMessages.slice(0, bigIndex).reduce((total, current) => total + current.tokens, 0)
   const { role, content } = body.messages[bigIndex]
@@ -69,13 +67,23 @@ const truncateCompletion = (body = {}, limit) => {
   })
 }
 
-const truncateWrapper = (body = {}, limit) => {
-  const isEmbedding = !!body.input
-  if (isEmbedding) {
-    return truncateEmbedding(body, limit)
-  } else {
-    return truncateCompletion(body, limit)
+/**
+ * Used to truncate a request
+ * @param {JSON} body - The entire body of the message
+ * @param {String} body.model - The model to pass to OpenAI
+ * @param {JSON} opts - (option) Additional options for truncation.
+ * @param {Int} opts.limit - (optional) Overrides the model limit for stricter rules.
+ * @param {JSON[]} body.messages - (optional) Used for completion messages.
+ * @param {String|String[]} body.input - (optional) Used for embeddings.
+ * @returns {JSON} The resulting object
+ */
+const truncateWrapper = (originalBody = {}, limit) => {
+  if (limit) {
+    console.warn('Using the "limit" argument on "truncateWrapper" is deprecated. Please us the "opts" property on the main object instead. Read more at https://github.com/mrsteele/openai-tokens/wiki/%5BDeprecated%5D-No-longer-supporting-the-%22limit%22-argument-on-%22truncateWrapper%22')
   }
+  const { opts, ...body } = originalBody
+  const fn = body.input ? truncateEmbedding : truncateCompletion
+  return fn(body, limit || opts?.limit)
 }
 
 module.exports = {

src/truncate.test.js

@@ -25,4 +25,16 @@ describe('truncateWrapper', () => {
 
     expect(response.input).toMatchObject([str, str, 'small embedding'])
   })
+
+  test('should support supplied limits', () => {
+    const response = truncateWrapper({
+      model: 'text-embedding-ada-002',
+      opts: {
+        limit: 2
+      },
+      input: [bigStr, bigStr, 'small embedding']
+    })
+
+    expect(response.input).toMatchObject(['so not', 'so not', 'small embed'])
+  })
 })

src/validate.js

@@ -10,7 +10,7 @@ const validateWrapper = (body = {}) => {
   return {
     tokenLimit,
     tokenTotal,
-    valid: tokenTotal < tokenLimit,
+    valid: tokenTotal <= tokenLimit,
     cost: model.price * tokenTotal
   }
 }