Update llm.py (#8513)

* Update llm.py * Update llm.py * Update llm.py * Update llm.py --------- Co-authored-by: Aarushi <[email protected]>
Significant-Gravitas · Oct 31, 2024 · 74d8aa6 · 74d8aa6
1 parent 3ea50f3
commit 74d8aa6
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/autogpt_platform/backend/backend/blocks/llm.py b/autogpt_platform/backend/backend/blocks/llm.py
@@ -519,18 +519,24 @@ def __init__(self):
             test_mock={"llm_call": lambda *args, **kwargs: "Response text"},
         )
 
-    def llm_call(self, input_data: AIStructuredResponseGeneratorBlock.Input) -> str:
+    def llm_call(
+        self,
+        input_data: AIStructuredResponseGeneratorBlock.Input,
+        credentials: APIKeyCredentials,
+    ) -> str:
         block = AIStructuredResponseGeneratorBlock()
-        response = block.run_once(input_data, "response")
+        response = block.run_once(input_data, "response", credentials=credentials)
         self.merge_stats(block.execution_stats)
         return response["response"]
 
-    def run(self, input_data: Input, **kwargs) -> BlockOutput:
+    def run(
+        self, input_data: Input, *, credentials: APIKeyCredentials, **kwargs
+    ) -> BlockOutput:
         object_input_data = AIStructuredResponseGeneratorBlock.Input(
             **{attr: getattr(input_data, attr) for attr in input_data.model_fields},
             expected_format={},
         )
-        yield "response", self.llm_call(object_input_data)
+        yield "response", self.llm_call(object_input_data, credentials)
 
 
 class SummaryStyle(Enum):