|
1 | 1 | import { describe, expect, it, vi, beforeEach } from "vitest"; |
2 | 2 | import { ModelProviderName, IAgentRuntime } from "../types"; |
3 | 3 | import { models } from "../models"; |
4 | | -import { generateText, generateTrueOrFalse, splitChunks, trimTokens } from "../generation"; |
| 4 | +import { |
| 5 | + generateText, |
| 6 | + generateTrueOrFalse, |
| 7 | + splitChunks, |
| 8 | + trimTokens, |
| 9 | +} from "../generation"; |
5 | 10 | import type { TiktokenModel } from "js-tiktoken"; |
6 | 11 |
|
7 | 12 | // Mock the elizaLogger |
@@ -42,6 +47,8 @@ describe("Generation", () => { |
42 | 47 | getSetting: vi.fn().mockImplementation((key: string) => { |
43 | 48 | if (key === "LLAMACLOUD_MODEL_LARGE") return false; |
44 | 49 | if (key === "LLAMACLOUD_MODEL_SMALL") return false; |
| 50 | + if (key === "TOGETHER_MODEL_LARGE") return false; |
| 51 | + if (key === "TOGETHER_MODEL_SMALL") return false; |
45 | 52 | return undefined; |
46 | 53 | }), |
47 | 54 | } as unknown as IAgentRuntime; |
@@ -122,53 +129,56 @@ describe("Generation", () => { |
122 | 129 | }); |
123 | 130 | }); |
124 | 131 |
|
125 | | - describe("trimTokens", () => { |
126 | | - const model = "gpt-4" as TiktokenModel; |
127 | | - |
128 | | - it("should return empty string for empty input", () => { |
129 | | - const result = trimTokens("", 100, model); |
130 | | - expect(result).toBe(""); |
131 | | - }); |
132 | | - |
133 | | - it("should throw error for negative maxTokens", () => { |
134 | | - expect(() => trimTokens("test", -1, model)).toThrow("maxTokens must be positive"); |
135 | | - }); |
136 | | - |
137 | | - it("should return unchanged text if within token limit", () => { |
138 | | - const shortText = "This is a short text"; |
139 | | - const result = trimTokens(shortText, 10, model); |
140 | | - expect(result).toBe(shortText); |
141 | | - }); |
142 | | - |
143 | | - it("should truncate text to specified token limit", () => { |
144 | | - // Using a longer text that we know will exceed the token limit |
145 | | - const longText = "This is a much longer text that will definitely exceed our very small token limit and need to be truncated to fit within the specified constraints." |
146 | | - const result = trimTokens(longText, 5, model); |
147 | | - |
148 | | - // The exact result will depend on the tokenizer, but we can verify: |
149 | | - // 1. Result is shorter than original |
150 | | - expect(result.length).toBeLessThan(longText.length); |
151 | | - // 2. Result is not empty |
152 | | - expect(result.length).toBeGreaterThan(0); |
153 | | - // 3. Result is a proper substring of the original text |
154 | | - expect(longText.includes(result)).toBe(true); |
155 | | - }); |
156 | | - |
157 | | - it("should handle non-ASCII characters", () => { |
158 | | - const unicodeText = "Hello 👋 World 🌍"; |
159 | | - const result = trimTokens(unicodeText, 5, model); |
160 | | - expect(result.length).toBeGreaterThan(0); |
161 | | - }); |
162 | | - |
163 | | - it("should handle multiline text", () => { |
164 | | - const multilineText = `Line 1 |
| 132 | + describe("trimTokens", () => { |
| 133 | + const model = "gpt-4" as TiktokenModel; |
| 134 | + |
| 135 | + it("should return empty string for empty input", () => { |
| 136 | + const result = trimTokens("", 100, model); |
| 137 | + expect(result).toBe(""); |
| 138 | + }); |
| 139 | + |
| 140 | + it("should throw error for negative maxTokens", () => { |
| 141 | + expect(() => trimTokens("test", -1, model)).toThrow( |
| 142 | + "maxTokens must be positive" |
| 143 | + ); |
| 144 | + }); |
| 145 | + |
| 146 | + it("should return unchanged text if within token limit", () => { |
| 147 | + const shortText = "This is a short text"; |
| 148 | + const result = trimTokens(shortText, 10, model); |
| 149 | + expect(result).toBe(shortText); |
| 150 | + }); |
| 151 | + |
| 152 | + it("should truncate text to specified token limit", () => { |
| 153 | + // Using a longer text that we know will exceed the token limit |
| 154 | + const longText = |
| 155 | + "This is a much longer text that will definitely exceed our very small token limit and need to be truncated to fit within the specified constraints."; |
| 156 | + const result = trimTokens(longText, 5, model); |
| 157 | + |
| 158 | + // The exact result will depend on the tokenizer, but we can verify: |
| 159 | + // 1. Result is shorter than original |
| 160 | + expect(result.length).toBeLessThan(longText.length); |
| 161 | + // 2. Result is not empty |
| 162 | + expect(result.length).toBeGreaterThan(0); |
| 163 | + // 3. Result is a proper substring of the original text |
| 164 | + expect(longText.includes(result)).toBe(true); |
| 165 | + }); |
| 166 | + |
| 167 | + it("should handle non-ASCII characters", () => { |
| 168 | + const unicodeText = "Hello 👋 World 🌍"; |
| 169 | + const result = trimTokens(unicodeText, 5, model); |
| 170 | + expect(result.length).toBeGreaterThan(0); |
| 171 | + }); |
| 172 | + |
| 173 | + it("should handle multiline text", () => { |
| 174 | + const multilineText = `Line 1 |
165 | 175 | Line 2 |
166 | 176 | Line 3 |
167 | 177 | Line 4 |
168 | 178 | Line 5`; |
169 | | - const result = trimTokens(multilineText, 5, model); |
170 | | - expect(result.length).toBeGreaterThan(0); |
171 | | - expect(result.length).toBeLessThan(multilineText.length); |
172 | | - }); |
173 | | - }); |
| 179 | + const result = trimTokens(multilineText, 5, model); |
| 180 | + expect(result.length).toBeGreaterThan(0); |
| 181 | + expect(result.length).toBeLessThan(multilineText.length); |
| 182 | + }); |
| 183 | + }); |
174 | 184 | }); |
0 commit comments