Langgraph type fixes (#361)

Author: JonesAndrew

src/art/langgraph/llm_wrapper.py

@@ -139,31 +139,29 @@ def _log(self, completion_id, input, output):
             entry = {"input": input, "output": output, "tools": self.tools}
             self.logger.log(f"{completion_id}", entry)
 
-    def invoke(self, input, config=None):
+    def invoke(self, input, config=None, **kwargs):
         completion_id = str(uuid.uuid4())
 
-        async def execute():
+        def execute():
             result = self.llm.invoke(input, config=config)
             self._log(completion_id, input, result)
             return result
 
         result = execute()
 
-        if hasattr(result, "get") and result.get("parsed"):
-            return result.get("parsed")
-
-        if hasattr(result, "tool_calls"):
-            for tool_call in result.tool_calls:
+        tool_calls = getattr(result, "tool_calls", None)
+        if tool_calls:
+            for tool_call in tool_calls:
                 if isinstance(tool_call["args"], str):
                     tool_call["args"] = json.loads(tool_call["args"])
 
         if self.structured_output:
             return self.structured_output.model_validate(
-                result.tool_calls[0]["args"] if result.tool_calls else None
+                tool_calls[0]["args"] if tool_calls else None
             )
         return result
 
-    async def ainvoke(self, input, config=None):
+    async def ainvoke(self, input, config=None, **kwargs):
         completion_id = str(uuid.uuid4())
 
         async def execute():
@@ -178,17 +176,15 @@ async def execute():
 
         result = await execute()
 
-        if hasattr(result, "get") and result.get("parsed"):
-            return result.get("parsed")
-
-        if hasattr(result, "tool_calls"):
-            for tool_call in result.tool_calls:
+        tool_calls = getattr(result, "tool_calls", None)
+        if tool_calls:
+            for tool_call in tool_calls:
                 if isinstance(tool_call["args"], str):
                     tool_call["args"] = json.loads(tool_call["args"])
 
         if self.structured_output:
             return self.structured_output.model_validate(
-                result.tool_calls[0]["args"] if result.tool_calls else None
+                tool_calls[0]["args"] if tool_calls else None
             )
         return result
 
@@ -220,23 +216,24 @@ def with_config(
     ):
         art_config = CURRENT_CONFIG.get()
         self.logger = art_config["logger"]
-        max_tokens = config.get("max_tokens")
 
         if hasattr(self.llm, "bound"):
-            self.llm.bound = ChatOpenAI(
-                base_url=art_config["base_url"],
-                api_key=art_config["api_key"],
-                model=art_config["model"],
-                temperature=1.0,
-                max_tokens=max_tokens,
+            setattr(
+                self.llm,
+                "bound",
+                ChatOpenAI(
+                    base_url=art_config["base_url"],
+                    api_key=art_config["api_key"],
+                    model=art_config["model"],
+                    temperature=1.0,
+                ),
             )
         else:
             self.llm = ChatOpenAI(
                 base_url=art_config["base_url"],
                 api_key=art_config["api_key"],
                 model=art_config["model"],
                 temperature=1.0,
-                max_tokens=max_tokens,
             )
 
         return self

src/art/langgraph/message_utils.py

@@ -19,6 +19,7 @@
 from openai.types.chat.chat_completion_function_message_param import (
     ChatCompletionFunctionMessageParam,
 )
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
 from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 from openai.types.chat.chat_completion_system_message_param import (
     ChatCompletionSystemMessageParam,
@@ -50,7 +51,7 @@ def make_message_param(role: str, **kwargs) -> ChatCompletionMessageParam:
     return cls(**kwargs)
 
 
-def langchain_msg_to_openai(msg: BaseMessage) -> Message:
+def langchain_msg_to_openai(msg: BaseMessage):
     if isinstance(msg, HumanMessage):
         role = "user"
     elif isinstance(msg, AIMessage):
@@ -71,23 +72,30 @@ def langchain_msg_to_openai(msg: BaseMessage) -> Message:
     result = {"role": role, "content": content}
 
     # Handle tool calls or function call if present
-    if hasattr(msg, "tool_calls") and msg.tool_calls:
-        result["tool_calls"] = msg.tool_calls
-    if hasattr(msg, "tool_call_id"):
-        result["tool_call_id"] = msg.tool_call_id
-    if hasattr(msg, "function_call") and msg.function_call:
-        result["function_call"] = msg.function_call
+    tool_calls = getattr(msg, "tool_calls", None)
+    if tool_calls:
+        result["tool_calls"] = tool_calls
+
+    tool_call_id = getattr(msg, "tool_call_id", None)
+    if tool_call_id:
+        result["tool_call_id"] = tool_call_id
+
+    function_call = getattr(msg, "function_call", None)
+    if function_call:
+        result["function_call"] = function_call
 
     return result
 
 
 def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
-    converted: MessagesAndChoices = []
+    converted = []
 
     for msg in messages:
-        if hasattr(msg, "response_metadata") and "logprobs" in msg.response_metadata:
-            if msg.tool_calls:
-                for tool_call in msg.tool_calls:
+        response_metadata = getattr(msg, "response_metadata")
+        if response_metadata and "logprobs" in response_metadata:
+            tool_calls = getattr(msg, "tool_calls", None)
+            if tool_calls:
+                for tool_call in tool_calls:
                     tool_call["function"] = {
                         "arguments": json.dumps(tool_call["args"]),
                         "name": tool_call["name"],
@@ -96,11 +104,13 @@ def convert_langgraph_messages(messages: List[object]) -> MessagesAndChoices:
 
             converted.append(
                 Choice(
-                    message=ChatCompletionAssistantMessageParam(
-                        role="assistant", content=msg.content, tool_calls=msg.tool_calls
+                    message=ChatCompletionMessage(
+                        role="assistant",
+                        content=getattr(msg, "content"),
+                        tool_calls=tool_calls,
                     ),
                     index=0,
-                    **msg.response_metadata,
+                    **response_metadata,
                 )
             )
         elif isinstance(msg, BaseMessage):