mirror of
https://github.com/langgenius/dify.git
synced 2026-05-13 08:57:28 +08:00
refactor(api): type LLM generator results with TypedDict (#34621)
This commit is contained in:
parent
89e23456f0
commit
0bce6b35b4
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user