feat: replace manager by utils

Chainlit · Nov 21, 2024 · b03d76d · b03d76d
1 parent f1902e7
commit b03d76d
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 44 deletions.
diff --git a/src/api.ts b/src/api.ts
@@ -5,8 +5,8 @@ import { ReadStream } from 'fs';
 import { v4 as uuidv4 } from 'uuid';
 
 import { LiteralClient } from '.';
-import { PromptCacheManager } from './cache/prompt-cache-manager';
 import { sharedCache } from './cache/sharedcache';
+import { getPromptCacheKey, putPrompt } from './cache/utils';
 import {
   Dataset,
   DatasetExperiment,
@@ -2156,7 +2156,7 @@ export class API {
   ) {
     const { id, name, version } = variables;
     const cachedPrompt = sharedCache.get(
-      PromptCacheManager.getPromptCacheKey({ id, name, version })
+      getPromptCacheKey({ id, name, version })
     );
     const timeout = cachedPrompt ? 1000 : undefined;
 
@@ -2176,7 +2176,7 @@ export class API {
       }
 
       const prompt = new Prompt(this, promptData);
-      PromptCacheManager.putPrompt(prompt);
+      putPrompt(prompt);
       return prompt;
     } catch (error) {
       return cachedPrompt;

diff --git a/src/cache/prompt-cache-manager.ts b/src/cache/prompt-cache-manager.ts
diff --git a/src/cache/utils.ts b/src/cache/utils.ts
@@ -0,0 +1,27 @@
+import { Prompt } from '../prompt-engineering/prompt';
+import { sharedCache } from './sharedcache';
+
+export function getPromptCacheKey({
+  id,
+  name,
+  version
+}: {
+  id?: string;
+  name?: string;
+  version?: number;
+}): string {
+  if (id) {
+    return id;
+  } else if (name && typeof version === 'number') {
+    return `${name}:${version}`;
+  } else if (name) {
+    return name;
+  }
+  throw new Error('Either id or name must be provided');
+}
+
+export function putPrompt(prompt: Prompt): void {
+  sharedCache.put(prompt.id, prompt);
+  sharedCache.put(prompt.name, prompt);
+  sharedCache.put(`${prompt.name}:${prompt.version}`, prompt);
+}
diff --git a/tests/api.test.ts b/tests/api.test.ts
@@ -3,8 +3,8 @@ import 'dotenv/config';
 import { v4 as uuidv4 } from 'uuid';
 
 import { ChatGeneration, IGenerationMessage, LiteralClient } from '../src';
-import { PromptCacheManager } from '../src/cache/prompt-cache-manager';
 import { sharedCache } from '../src/cache/sharedcache';
+import { putPrompt } from '../src/cache/utils';
 import { Dataset } from '../src/evaluation/dataset';
 import { Score } from '../src/evaluation/score';
 import { Prompt, PromptConstructor } from '../src/prompt-engineering/prompt';
@@ -687,7 +687,7 @@ is a templated list.`;
 
     it('should fallback to cache when getPromptById DB call fails', async () => {
       const prompt = new Prompt(client.api, mockPromptData);
-      PromptCacheManager.putPrompt(prompt);
+      putPrompt(prompt);
 
       jest
         .spyOn(client.api as any, 'makeGqlCall')
@@ -699,7 +699,7 @@ is a templated list.`;
 
     it('should fallback to cache when getPrompt DB call fails', async () => {
       const prompt = new Prompt(client.api, mockPromptData);
-      PromptCacheManager.putPrompt(prompt);
+      putPrompt(prompt);
       jest.spyOn(axios, 'post').mockRejectedValueOnce(new Error('DB Error'));
 
       const result = await client.api.getPrompt(prompt.id);

diff --git a/tests/cache.test.ts b/tests/cache.test.ts
@@ -1,6 +1,6 @@
 import { API } from '../src/api';
-import { PromptCacheManager } from '../src/cache/prompt-cache-manager';
 import { sharedCache } from '../src/cache/sharedcache';
+import { getPromptCacheKey, putPrompt } from '../src/cache/utils';
 import { Prompt, PromptConstructor } from '../src/prompt-engineering/prompt';
 
 describe('Cache', () => {
@@ -38,10 +38,10 @@ describe('Cache', () => {
     mockPrompt = new Prompt(api, mockPromptData);
   });
 
-  describe('PromptCacheManager', () => {
+  describe('Cache Utils', () => {
     describe('getPromptCacheKey', () => {
       it('should return id when provided', () => {
-        const key = PromptCacheManager.getPromptCacheKey({
+        const key = getPromptCacheKey({
           id: 'test-id',
           name: 'test-name',
           version: 1
@@ -50,28 +50,28 @@ describe('Cache', () => {
       });
 
       it('should return name:version when id not provided but name and version are', () => {
-        const key = PromptCacheManager.getPromptCacheKey({
+        const key = getPromptCacheKey({
           name: 'test-name',
           version: 1
         });
         expect(key).toBe('test-name:1');
       });
 
       it('should return name when only name provided', () => {
-        const key = PromptCacheManager.getPromptCacheKey({ name: 'test-name' });
+        const key = getPromptCacheKey({ name: 'test-name' });
         expect(key).toBe('test-name');
       });
 
       it('should throw error when neither id nor name provided', () => {
-        expect(() =>
-          PromptCacheManager.getPromptCacheKey({ version: 0 })
-        ).toThrow('Either id or name must be provided');
+        expect(() => getPromptCacheKey({ version: 0 })).toThrow(
+          'Either id or name must be provided'
+        );
       });
     });
 
     describe('putPrompt', () => {
       it('should store prompt with multiple keys', () => {
-        PromptCacheManager.putPrompt(mockPrompt);
+        putPrompt(mockPrompt);
 
         expect(sharedCache.get('test-id')).toEqual(mockPrompt);
         expect(sharedCache.get('test-name')).toEqual(mockPrompt);