Update llm.py (#8513)

* Update llm.py

* Update llm.py

* Update llm.py

* Update llm.py

---------

Co-authored-by: Aarushi <50577581+aarushik93@users.noreply.github.com>
pull/8514/head agpt-platform-beta-v0.2.0
Nicholas Tindle 2024-10-31 00:40:38 -05:00 committed by GitHub
parent 3ea50f3b70
commit 74d8aa6d90
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 4 deletions

View File

@ -519,18 +519,24 @@ class AITextGeneratorBlock(Block):
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):