fix error on building index using anthropic sonnet (#576)

pingcap · Jan 8, 2025 · 69e6098 · 69e6098
1 parent ae2b633
commit 69e6098
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/backend/app/rag/knowledge_graph/extractor.py b/backend/app/rag/knowledge_graph/extractor.py
@@ -112,7 +112,9 @@ def get_llm_output_config(self):
         elif "bedrock" in self.dspy_lm.provider.lower():
             # Fix: add bedrock branch to fix 'Malformed input request' error
             # subject must not be valid against schema {"required":["messages"]}: extraneous key [response_mime_type] is not permitted
-            return {}
+            return {
+                 "max_tokens": 8192
+            }
         else:
             return {
                 "response_mime_type": "application/json",

diff --git a/backend/app/utils/dspy.py b/backend/app/utils/dspy.py
@@ -52,7 +52,7 @@ def get_dspy_lm_by_llama_llm(llama_llm: BaseLLM) -> dspy.LM:
         bedrock = dspy.Bedrock(region_name=llama_llm.region_name)
         if llama_llm.model.startswith("anthropic"):
             return dspy.AWSAnthropic(
-                bedrock, model=llama_llm.model, max_new_tokens=llama_llm.max_tokens
+                bedrock, model=llama_llm.model, max_new_tokens=llama_llm.max_tokens or 8192
             )
         elif llama_llm.model.startswith("meta"):
             return dspy.AWSMeta(