feat(endpoints): Add conv ID to headers passed to TGI (#1511)

huggingface · Oct 10, 2024 · 6451bee · 6451bee
1 parent 81c4393
commit 6451bee
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 2 deletions.
diff --git a/src/lib/server/endpoints/endpoints.ts b/src/lib/server/endpoints/endpoints.ts
@@ -29,6 +29,7 @@ import endpointLangserve, {
 } from "./langserve/endpointLangserve";
 
 import type { Tool, ToolCall, ToolResult } from "$lib/types/Tool";
+import type { ObjectId } from "mongodb";
 
 export type EndpointMessage = Omit<Message, "id">;
 
@@ -41,6 +42,7 @@ export interface EndpointParameters {
 	tools?: Tool[];
 	toolResults?: ToolResult[];
 	isMultimodal?: boolean;
+	conversationId?: ObjectId;
 }
 
 interface CommonEndpoint {

diff --git a/src/lib/server/endpoints/openai/endpointOai.ts b/src/lib/server/endpoints/openai/endpointOai.ts
@@ -149,7 +149,7 @@ export async function endpointOai(
 				"Tools are not supported for 'completions' mode, switch to 'chat_completions' instead"
 			);
 		}
-		return async ({ messages, preprompt, continueMessage, generateSettings }) => {
+		return async ({ messages, preprompt, continueMessage, generateSettings, conversationId }) => {
 			const prompt = await buildPrompt({
 				messages,
 				continueMessage,
@@ -171,12 +171,22 @@ export async function endpointOai(
 
 			const openAICompletion = await openai.completions.create(body, {
 				body: { ...body, ...extraBody },
+				headers: {
+					"ChatUI-Conversation-ID": conversationId?.toString() ?? "",
+				},
 			});
 
 			return openAICompletionToTextGenerationStream(openAICompletion);
 		};
 	} else if (completion === "chat_completions") {
-		return async ({ messages, preprompt, generateSettings, tools, toolResults }) => {
+		return async ({
+			messages,
+			preprompt,
+			generateSettings,
+			tools,
+			toolResults,
+			conversationId,
+		}) => {
 			let messagesOpenAI: OpenAI.Chat.Completions.ChatCompletionMessageParam[] =
 				await prepareMessages(messages, imageProcessor, !model.tools && model.multimodal);
 
@@ -240,6 +250,9 @@ export async function endpointOai(
 
 			const openChatAICompletion = await openai.chat.completions.create(body, {
 				body: { ...body, ...extraBody },
+				headers: {
+					"ChatUI-Conversation-ID": conversationId?.toString() ?? "",
+				},
 			});
 
 			return openAIChatToTextGenerationStream(openChatAICompletion);

diff --git a/src/lib/server/endpoints/tgi/endpointTgi.ts b/src/lib/server/endpoints/tgi/endpointTgi.ts
@@ -43,6 +43,7 @@ export function endpointTgi(input: z.input<typeof endpointTgiParametersSchema>):
 		tools,
 		toolResults,
 		isMultimodal,
+		conversationId,
 	}) => {
 		const messagesWithResizedFiles = await Promise.all(
 			messages.map((message) => prepareMessage(Boolean(isMultimodal), message, imageProcessor))
@@ -72,6 +73,7 @@ export function endpointTgi(input: z.input<typeof endpointTgiParametersSchema>):
 						info.headers = {
 							...info.headers,
 							Authorization: authorization,
+							"ChatUI-Conversation-ID": conversationId?.toString() ?? "",
 						};
 					}
 					return fetch(endpointUrl, info);

diff --git a/src/lib/server/textGeneration/generate.ts b/src/lib/server/textGeneration/generate.ts
@@ -18,6 +18,7 @@ export async function* generate(
 		generateSettings: assistant?.generateSettings,
 		toolResults,
 		isMultimodal: model.multimodal,
+		conversationId: conv._id,
 	})) {
 		// text generation completed
 		if (output.generated_text) {

diff --git a/src/lib/server/textGeneration/tools.ts b/src/lib/server/textGeneration/tools.ts
@@ -196,6 +196,7 @@ export async function* runTools(
 				type: input.type === "file" ? "str" : input.type,
 			})),
 		})),
+		conversationId: conv._id,
 	})) {
 		// model natively supports tool calls
 		if (output.token.toolCalls) {