From 0bce6b35b45b18495cca67fce7186f94341c7f14 Mon Sep 17 00:00:00 2001
From: Statxc <tyleradams93226@gmail.com>
Date: Tue, 7 Apr 2026 03:06:08 +0200
Subject: [PATCH] refactor(api): type LLM generator results with TypedDict
 (#34621)

---
 api/core/llm_generator/llm_generator.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/api/core/llm_generator/llm_generator.py b/api/core/llm_generator/llm_generator.py
index d39630ad95..8071770c0f 100644
--- a/api/core/llm_generator/llm_generator.py
+++ b/api/core/llm_generator/llm_generator.py
@@ -2,7 +2,7 @@ import json
 import logging
 import re
 from collections.abc import Sequence
-from typing import Protocol, cast
+from typing import Protocol, TypedDict, cast
 
 import json_repair
 from graphon.enums import WorkflowNodeExecutionMetadataKey
@@ -49,6 +49,17 @@ class WorkflowServiceInterface(Protocol):
         pass
 
 
+class CodeGenerateResultDict(TypedDict):
+    code: str
+    language: str
+    error: str
+
+
+class StructuredOutputResultDict(TypedDict):
+    output: str
+    error: str
+
+
 class LLMGenerator:
     @classmethod
     def generate_conversation_name(
@@ -293,7 +304,7 @@ class LLMGenerator:
         cls,
         tenant_id: str,
         args: RuleCodeGeneratePayload,
-    ):
+    ) -> CodeGenerateResultDict:
         if args.code_language == "python":
             prompt_template = PromptTemplateParser(PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE)
         else:
@@ -362,7 +373,9 @@ class LLMGenerator:
         return answer.strip()
 
     @classmethod
-    def generate_structured_output(cls, tenant_id: str, args: RuleStructuredOutputPayload):
+    def generate_structured_output(
+        cls, tenant_id: str, args: RuleStructuredOutputPayload
+    ) -> StructuredOutputResultDict:
         model_manager = ModelManager.for_tenant(tenant_id=tenant_id)
         model_instance = model_manager.get_model_instance(
             tenant_id=tenant_id,