Merge pull request #215 from restackio/fix-pydantic-serialization

Fix pydantic serialization error

Author: leawn

agent_rag/src/agents/chat_rag.py

@@ -1,10 +1,19 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import NonRetryableError, agent, import_functions, log
+from restack_ai.agent import (
+    NonRetryableError,
+    agent,
+    import_functions,
+    log,
+)
 
 with import_functions():
-    from src.functions.llm_chat import LlmChatInput, Message, llm_chat
+    from src.functions.llm_chat import (
+        LlmChatInput,
+        Message,
+        llm_chat,
+    )
     from src.functions.lookup_sales import lookup_sales
 
 

agent_rag/src/functions/llm_chat.py

@@ -53,4 +53,4 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
         raise NonRetryableError(error_message) from e
     else:
         log.info("llm_chat function completed", response=response)
-        return response
+        return response.model_dump()

agent_todo/src/agents/agent_todo.py

@@ -1,15 +1,30 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import NonRetryableError, agent, import_functions, log
-
-from src.workflows.todo_execute import TodoExecute, TodoExecuteParams
+from restack_ai.agent import (
+    NonRetryableError,
+    agent,
+    import_functions,
+    log,
+)
+
+from src.workflows.todo_execute import (
+    TodoExecute,
+    TodoExecuteParams,
+)
 
 with import_functions():
     from openai import pydantic_function_tool
 
-    from src.functions.llm_chat import LlmChatInput, Message, llm_chat
-    from src.functions.todo_create import TodoCreateParams, todo_create
+    from src.functions.llm_chat import (
+        LlmChatInput,
+        Message,
+        llm_chat,
+    )
+    from src.functions.todo_create import (
+        TodoCreateParams,
+        todo_create,
+    )
 
 
 class MessagesEvent(BaseModel):

agent_todo/src/functions/llm_chat.py

@@ -7,7 +7,9 @@
 from openai.types.chat.chat_completion_message_tool_call import (
     ChatCompletionMessageToolCall,
 )
-from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
 from pydantic import BaseModel
 from restack_ai.function import NonRetryableError, function, log
 
@@ -62,4 +64,4 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
         raise NonRetryableError(error_message) from e
     else:
         log.info("llm_chat function completed", response=response)
-        return response
+        return response.model_dump()

agent_todo/src/workflows/todo_execute.py

@@ -1,7 +1,12 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.workflow import NonRetryableError, import_functions, log, workflow
+from restack_ai.workflow import (
+    NonRetryableError,
+    import_functions,
+    log,
+    workflow,
+)
 
 with import_functions():
     from src.functions.get_random import RandomParams, get_random

agent_tool/src/agents/chat_tool_functions.py

@@ -1,13 +1,25 @@
 from datetime import timedelta
 
 from pydantic import BaseModel
-from restack_ai.agent import NonRetryableError, agent, import_functions, log
+from restack_ai.agent import (
+    NonRetryableError,
+    agent,
+    import_functions,
+    log,
+)
 
 with import_functions():
     from openai import pydantic_function_tool
 
-    from src.functions.llm_chat import LlmChatInput, Message, llm_chat
-    from src.functions.lookup_sales import LookupSalesInput, lookup_sales
+    from src.functions.llm_chat import (
+        LlmChatInput,
+        Message,
+        llm_chat,
+    )
+    from src.functions.lookup_sales import (
+        LookupSalesInput,
+        lookup_sales,
+    )
 
 
 class MessagesEvent(BaseModel):

agent_tool/src/functions/llm_chat.py

@@ -7,7 +7,9 @@
 from openai.types.chat.chat_completion_message_tool_call import (
     ChatCompletionMessageToolCall,
 )
-from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
 from pydantic import BaseModel
 from restack_ai.function import NonRetryableError, function, log
 
@@ -52,11 +54,15 @@ async def llm_chat(function_input: LlmChatInput) -> ChatCompletion:
                 Message(role="system", content=function_input.system_content or "")
             )
 
-        return client.chat.completions.create(
+        result = client.chat.completions.create(
             model=function_input.model or "gpt-4o-mini",
             messages=function_input.messages,
             tools=function_input.tools,
         )
+
+        log.info("llm_chat function completed", result=result)
+
+        return result.model_dump()
     except Exception as e:
         error_message = f"LLM chat failed: {e}"
         raise NonRetryableError(error_message) from e