Merge pull request #6 from matiasvlevi/dev

Merge Dev

package.json

@@ -1,6 +1,6 @@
 {
   "name": "gptdoc",
-  "version": "1.0.8",
+  "version": "1.0.9",
   "license": "MIT",
   "author": "Matias Vazquez-Levi",
   "description": "Generate documentation comments with the help of OpenAI's models",

src/config/makeConfig.ts

@@ -1,7 +1,7 @@
 import { Config } from './types'
 import { CLIArgs } from "../cli";
 import { parseEnv } from './parseEnv';
-import { Models } from '../gpt';
+import { Models } from '../models';
 
 /** 
  * Merge a configuration with the defaults

src/core/GPTDocument.ts

@@ -4,6 +4,10 @@ import {
     JSDOCComment
 } from './regex';
 
+import {
+    Models
+} from '../models'
+
 import {
     GPT_DEBUG_COMMENT,
     GPT_PROMPT,
@@ -134,12 +138,13 @@ export class GPTDocument {
             // Generate a prompt
             prompt = GPT_PROMPT(project.config, this.meta.kind, this.source);
             prompt_tokens = GPTDocument.estimateTokenCount(prompt);
-            
+
             // Length guard, 
             // do not send a request to OpenAI if prompt is too lengthy
-            if (prompt_tokens >= 4000) {
+            const { maxTokens } = Models[project.config.openai.model];
+            if (prompt_tokens >= maxTokens) {
                 Logger.error(
-                    `Your prompt has ${prompt_tokens} tokens, maximum is 4000`, true
+                    `Your prompt has ${prompt_tokens} tokens, maximum is ${maxTokens}`, true
                 );
                 return source;
             }

src/core/Project.ts

@@ -5,7 +5,11 @@ import { File } from './File'
 import { Config } from "../config/index";
 
 import * as Logger from "../utils/Logger";
-import { Models, PriceRange } from "../gpt";
+
+import {
+    Models,
+    PriceRange
+} from '../models'
 
 /** @gpt */
 export class Project {
@@ -132,19 +136,15 @@ export class Project {
         this.prompt_tokens += value;
     }
 
-
     addResponseTokens(value: number) {
         this.response_tokens += value;
     }
 
     getTokenCost(): number | PriceRange {
         let price: number | PriceRange = 0.02;
 
-        for (let model in Models) {
-            if (this.config.openai.model.includes(model)) {
-                price = Models[model].price;
-                break;
-            }
+        if (this.config.openai.model in Models) {
+            price = Models[this.config.openai.model].price;
         }
 
         return price; 

src/gpt.ts

@@ -71,8 +71,6 @@ export async function OpenAICompletion(_config: Config, prompt: string) {
         GPT_CHAT_COMPLETION_CONFIG(config, prompt):
         GPT_COMPLETION_CONFIG(config, prompt);
 
-
-
     const res = await fetch(
         `https://api.openai.com/v1/${_config.chat ? 'chat/completions' : 'completions'}`, {
         headers: {
@@ -86,70 +84,3 @@ export async function OpenAICompletion(_config: Config, prompt: string) {
     return await res.json();
 }
 
-export interface PriceRange {
-    prompt: number;
-    response: number
-}
-
-interface ModelMeta {
-    price: number | PriceRange;
-    isChatModel: boolean;
-}
-
-/**
- * Lookup table for pricing
- */
-export const Models: { [key:string]: ModelMeta } = {
-    'text-ada-001': {
-        price: 0.0004,
-        isChatModel: false
-    },
-    'text-babbage-001': {
-        price: 0.0005,
-        isChatModel: false
-    },
-    'text-curie-001': {
-        price: 0.002,
-        isChatModel: false
-    },
-    'text-davinci-002': {
-        price: 0.02,
-        isChatModel: false
-    },
-    'text-davinci-003': {
-        price: 0.02,
-        isChatModel: false
-    },
-    'gpt-3.5-turbo': {
-        price: 0.002,
-        isChatModel: true
-    },
-    'gpt-4': {
-        price: {
-            prompt: 0.03,
-            response: 0.06
-        },
-        isChatModel: true
-    },
-    'gpt-4-0314': {
-        price: {
-            prompt: 0.03,
-            response: 0.06
-        },
-        isChatModel: true
-    },
-    'gpt-4-32k': {
-        price: {
-            prompt: 0.06,
-            response: 0.12
-        },
-        isChatModel: true
-    },
-    'gpt-4-32k-0314': {
-        price: {
-            prompt: 0.06,
-            response: 0.12
-        },
-        isChatModel: true
-    }
-}

src/models.ts

@@ -0,0 +1,78 @@
+export interface PriceRange {
+    prompt: number;
+    response: number
+}
+
+interface ModelMeta {
+    price: number | PriceRange;
+    maxTokens: number;
+    isChatModel: boolean;
+}
+
+/**
+ * Lookup table for pricing
+ */
+export const Models: { [key:string]: ModelMeta } = {
+    'text-ada-001': {
+        price: 0.0004,
+        maxTokens: 2000,
+        isChatModel: false
+    },
+    'text-babbage-001': {
+        price: 0.0005,
+        maxTokens: 2000,
+        isChatModel: false
+    },
+    'text-curie-001': {
+        price: 0.002,
+        maxTokens: 2000,
+        isChatModel: false
+    },
+    'text-davinci-002': {
+        price: 0.02,
+        maxTokens: 4000,
+        isChatModel: false
+    },
+    'text-davinci-003': {
+        price: 0.02,
+        maxTokens: 4000,
+        isChatModel: false
+    },
+    'gpt-3.5-turbo': {
+        price: 0.002,
+        maxTokens: 4000,
+        isChatModel: true
+    },
+    'gpt-4': {
+        price: {
+            prompt: 0.03,
+            response: 0.06
+        },
+        maxTokens: 8000,
+        isChatModel: true
+    },
+    'gpt-4-0314': {
+        price: {
+            prompt: 0.03,
+            response: 0.06
+        },
+        maxTokens: 8000,
+        isChatModel: true
+    },
+    'gpt-4-32k': {
+        price: {
+            prompt: 0.06,
+            response: 0.12
+        },
+        maxTokens: 32000,
+        isChatModel: true
+    },
+    'gpt-4-32k-0314': {
+        price: {
+            prompt: 0.06,
+            response: 0.12
+        },
+        maxTokens: 32000,
+        isChatModel: true
+    }
+}

src/utils/Logger.ts

@@ -1,4 +1,4 @@
-import { PriceRange } from "../gpt";
+import { PriceRange } from "../models";
 
 /** @gpt */
 export function error(text: string, exit: boolean = false) {