fix: pre-commit run checking

Cinnamon · May 23, 2024 · 1b5c94b · 1b5c94b
1 parent 293b39b
commit 1b5c94b
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/libs/kotaemon/kotaemon/indices/qa/citation.py b/libs/kotaemon/kotaemon/indices/qa/citation.py
@@ -105,9 +105,9 @@ def invoke(self, context: str, question: str):
             print("CitationPipeline: finish invoking LLM")
             if not llm_output.messages:
                 return None
-            function_output = llm_output.messages[0].additional_kwargs["tool_calls"][0]["function"][
-                "arguments"
-            ]
+            function_output = llm_output.messages[0].additional_kwargs["tool_calls"][0][
+                "function"
+            ]["arguments"]
             output = QuestionAnswer.parse_raw(function_output)
         except Exception as e:
             print(e)
@@ -122,9 +122,9 @@ async def ainvoke(self, context: str, question: str):
             print("CitationPipeline: async invoking LLM")
             llm_output = await self.get_from_path("llm").ainvoke(messages, **llm_kwargs)
             print("CitationPipeline: finish async invoking LLM")
-            function_output = llm_output.messages[0].additional_kwargs["tool_calls"][0]["function"][
-                "arguments"
-            ]
+            function_output = llm_output.messages[0].additional_kwargs["tool_calls"][0][
+                "function"
+            ]["arguments"]
             output = QuestionAnswer.parse_raw(function_output)
         except Exception as e:
             print(e)