From 0003f29a88b3a794a5c8bfa010d4804bbec0e316 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 24 May 2024 14:59:59 +0300 Subject: [PATCH] feat: Use custom model for summarization --- agents-api/agents_api/activities/summarization.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/agents-api/agents_api/activities/summarization.py b/agents-api/agents_api/activities/summarization.py index 1eaa2003d..f2ebd387b 100644 --- a/agents-api/agents_api/activities/summarization.py +++ b/agents-api/agents_api/activities/summarization.py @@ -15,7 +15,7 @@ ) from agents_api.common.protocol.entries import Entry from ..model_registry import JULEP_MODELS -from ..env import model_inference_url, model_api_key +from ..env import model_inference_url, model_api_key, summarization_model_name from agents_api.rec_sum.entities import get_entities from agents_api.rec_sum.summarize import summarize_messages from agents_api.rec_sum.trim import trim_messages @@ -209,8 +209,8 @@ async def summarization(session_id: str) -> None: assert len(entries) > 0, "no need to summarize on empty entries list" trimmed_messages, entities = await asyncio.gather( - trim_messages(entries), - get_entities(entries), + trim_messages(entries, model=summarization_model_name), + get_entities(entries, model=summarization_model_name), ) summarized = await summarize_messages(trimmed_messages)