@tanstack/ai
Version:
Type-safe TypeScript AI SDK for streaming chat, tool calling, agents, structured outputs, and multimodal generation.
213 lines (212 loc) • 7.67 kB
JavaScript
import { EventType } from "@ag-ui/core";
import { toRunErrorPayload } from "../error-payload.js";
import { MAX_TOKENS_KEYS } from "../../utilities/sampling-keys.js";
import { BaseSummarizeAdapter } from "./adapter.js";
const MAX_TOKENS_KEY_BY_ADAPTER = {
openai: "max_output_tokens",
anthropic: "max_tokens",
grok: "max_tokens",
groq: "max_completion_tokens",
gemini: "maxOutputTokens",
openrouter: "maxCompletionTokens"
};
const KNOWN_MAX_TOKENS_KEYS = MAX_TOKENS_KEYS;
function isKnownMaxTokensAdapter(adapterName) {
return adapterName === "ollama" || MAX_TOKENS_KEY_BY_ADAPTER[adapterName] !== void 0;
}
function applyDefaultTemperature(adapterName, temperature, modelOptions) {
const merged = { ...modelOptions };
if (adapterName === "ollama") {
const existing = merged.options && typeof merged.options === "object" ? merged.options : void 0;
if (existing && "temperature" in existing) return merged;
merged.options = { temperature, ...existing };
return merged;
}
if ("temperature" in merged) return merged;
merged.temperature = temperature;
return merged;
}
function applyMaxLength(adapterName, maxLength, modelOptions) {
const merged = { ...modelOptions };
if (adapterName === "ollama") {
const callerSetFlatLimit = KNOWN_MAX_TOKENS_KEYS.some(
(k) => typeof merged[k] === "number"
);
const existing = merged.options && typeof merged.options === "object" ? merged.options : void 0;
if (callerSetFlatLimit || existing && typeof existing.num_predict === "number") {
return merged;
}
merged.options = { num_predict: maxLength, ...existing };
return merged;
}
const key = MAX_TOKENS_KEY_BY_ADAPTER[adapterName];
if (key === void 0) return merged;
const callerSetLimit = KNOWN_MAX_TOKENS_KEYS.some(
(k) => typeof merged[k] === "number"
);
if (callerSetLimit) return merged;
merged[key] = maxLength;
return merged;
}
class ChatStreamSummarizeAdapter extends BaseSummarizeAdapter {
name;
textAdapter;
constructor(textAdapter, model, name = "chat-stream-summarize") {
super({}, model);
this.name = name;
this.textAdapter = textAdapter;
}
async summarize(options) {
const systemPrompt = this.buildSummarizationPrompt(options);
let summary = "";
const id = this.generateId();
let model = options.model;
let usage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
options.logger.request(
`activity=summarize provider=${this.name} model=${options.model} text-length=${options.text.length} maxLength=${options.maxLength ?? "unset"}`,
{ provider: this.name, model: options.model }
);
try {
for await (const chunk of this.textAdapter.chatStream(
this.buildTextOptions(options, systemPrompt)
)) {
if (chunk.type === "TEXT_MESSAGE_CONTENT") {
if (chunk.content) {
summary = chunk.content;
} else if (chunk.delta) {
summary += chunk.delta;
}
model = chunk.model || model;
}
if (chunk.type === "RUN_FINISHED") {
if (chunk.usage) {
usage = chunk.usage;
}
}
if (chunk.type === "RUN_ERROR") {
const message = (chunk.error && typeof chunk.error.message === "string" ? chunk.error.message : null) ?? "Summarization failed";
const code = chunk.error && typeof chunk.error.code === "string" ? chunk.error.code : void 0;
const err = new Error(message);
if (code) {
;
err.code = code;
}
throw err;
}
}
} catch (error) {
options.logger.errors(`${this.name}.summarize fatal`, {
error: toRunErrorPayload(error, `${this.name}.summarize failed`),
source: `${this.name}.summarize`
});
throw error;
}
return { id, model, summary, usage };
}
async *summarizeStream(options) {
const systemPrompt = this.buildSummarizationPrompt(options);
options.logger.request(
`activity=summarizeStream provider=${this.name} model=${options.model} text-length=${options.text.length} maxLength=${options.maxLength ?? "unset"}`,
{ provider: this.name, model: options.model }
);
const id = this.generateId();
let summary = "";
let model = options.model;
let usage = {
promptTokens: 0,
completionTokens: 0,
totalTokens: 0
};
try {
for await (const chunk of this.textAdapter.chatStream(
this.buildTextOptions(options, systemPrompt)
)) {
if (chunk.type === "TEXT_MESSAGE_CONTENT") {
if (chunk.content) {
summary = chunk.content;
} else if (chunk.delta) {
summary += chunk.delta;
}
if (chunk.model) model = chunk.model;
}
if (chunk.type === "RUN_FINISHED") {
if (chunk.usage) usage = chunk.usage;
if (chunk.model) model = chunk.model;
yield {
type: EventType.CUSTOM,
name: "generation:result",
value: { id, model, summary, usage },
model,
timestamp: Date.now()
};
}
yield chunk;
}
} catch (error) {
options.logger.errors(`${this.name}.summarizeStream fatal`, {
error: toRunErrorPayload(error, `${this.name}.summarizeStream failed`),
source: `${this.name}.summarizeStream`
});
throw error;
}
}
/**
* Build the TextOptions passed to the underlying chatStream. Provider
* `modelOptions` from the summarize call are forwarded as-is so knobs like
* Anthropic cache headers, Gemini safety settings, or Ollama tuning params
* still reach the wire layer.
*/
buildTextOptions(options, systemPrompt) {
let working = {
...options.modelOptions
};
working = applyDefaultTemperature(this.name, 0.3, working);
if (options.maxLength !== void 0) {
if (!isKnownMaxTokensAdapter(this.name)) {
options.logger.warn(
`summarize: maxLength=${options.maxLength} could not be mapped to a provider token key for adapter name "${this.name}" — it was dropped from modelOptions (the prompt still asks the model to stay under it). Construct ChatStreamSummarizeAdapter with a recognised provider name to forward the cap.`,
{ provider: this.name }
);
}
working = applyMaxLength(this.name, options.maxLength, working);
}
const modelOptions = working;
return {
model: options.model,
messages: [{ role: "user", content: options.text }],
systemPrompts: [systemPrompt],
modelOptions,
logger: options.logger
};
}
buildSummarizationPrompt(options) {
let prompt = "You are a professional summarizer. ";
switch (options.style) {
case "bullet-points":
prompt += "Provide a summary in bullet point format. ";
break;
case "paragraph":
prompt += "Provide a summary in paragraph format. ";
break;
case "concise":
prompt += "Provide a very concise summary in 1-2 sentences. ";
break;
case void 0:
prompt += "Provide a clear and concise summary. ";
break;
default:
prompt += "Provide a clear and concise summary. ";
}
if (options.focus && options.focus.length > 0) {
prompt += `Focus on the following aspects: ${options.focus.join(", ")}. `;
}
if (options.maxLength) {
prompt += `Keep the summary under ${options.maxLength} tokens. `;
}
return prompt;
}
}
export {
ChatStreamSummarizeAdapter
};
//# sourceMappingURL=chat-stream-summarize.js.map