From c3bc0eb53bda7fdd712b1d541d986b6305ae3603 Mon Sep 17 00:00:00 2001 From: Balaji Munusamy Date: Fri, 2 May 2025 18:25:48 +0200 Subject: [PATCH] made bedrock complete generic --- lightrag/llm/bedrock.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lightrag/llm/bedrock.py b/lightrag/llm/bedrock.py index 615d3453..5f712603 100644 --- a/lightrag/llm/bedrock.py +++ b/lightrag/llm/bedrock.py @@ -94,12 +94,14 @@ async def bedrock_complete_if_cache( return response["output"]["message"]["content"][0]["text"] +# Generic Bedrock completion function async def bedrock_complete( prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs ) -> str: keyword_extraction = kwargs.pop("keyword_extraction", None) + model_name = kwargs["hashing_kv"].global_config["llm_model_name"] result = await bedrock_complete_if_cache( - "anthropic.claude-3-haiku-20240307-v1:0", + model_name, prompt, system_prompt=system_prompt, history_messages=history_messages,