diff --git a/zhenxun/services/llm/api.py b/zhenxun/services/llm/api.py
index 3bb45897..644c6de0 100644
--- a/zhenxun/services/llm/api.py
+++ b/zhenxun/services/llm/api.py
@@ -182,7 +182,7 @@ async def embed_documents(
 
 
 async def generate_structured(
-    message: str | LLMMessage | list[LLMContentPart],
+    message: str | UniMessage | LLMMessage | list[LLMContentPart],
     response_model: type[T],
     *,
     model: ModelName = None,
diff --git a/zhenxun/services/llm/session.py b/zhenxun/services/llm/session.py
index a07a7987..5f2e646b 100644
--- a/zhenxun/services/llm/session.py
+++ b/zhenxun/services/llm/session.py
@@ -449,7 +449,7 @@ class AI:
 
     async def generate_structured(
         self,
-        message: str | LLMMessage | list[LLMContentPart] | None,
+        message: str | UniMessage | LLMMessage | list[LLMContentPart] | None,
         response_model: type[T],
         *,
         model: ModelName = None,