refactor(api): type LLM generator results with TypedDict (#34621)

This commit is contained in:
Statxc 2026-04-07 03:06:08 +02:00 committed by GitHub
parent 89e23456f0
commit 0bce6b35b4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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,