hdsp-jupyter-extension 2.0.21__py3-none-any.whl → 2.0.23__py3-none-any.whl

agent_server/langchain/models/gpt_oss_chat.py

@@ -98,7 +98,9 @@ class ChatGPTOSS(BaseChatModel):
                     "function": {
                         "name": tool.name,
                         "description": tool.description or "",
-                        "parameters": tool.args_schema.schema() if tool.args_schema else {"type": "object", "properties": {}},
+                        "parameters": tool.args_schema.schema()
+                        if tool.args_schema
+                        else {"type": "object", "properties": {}},
                     },
                 }
                 formatted_tools.append(tool_schema)
@@ -116,6 +118,7 @@ class ChatGPTOSS(BaseChatModel):
             temperature=self.temperature,
             max_tokens=self.max_tokens,
             streaming=self.streaming,
+            callbacks=self.callbacks,  # Preserve callbacks (e.g., LLMTraceLogger)
         )
         new_instance._tools = formatted_tools
         new_instance._tool_choice = tool_choice
@@ -131,15 +134,19 @@ class ChatGPTOSS(BaseChatModel):
         for msg in messages:
             if isinstance(msg, SystemMessage):
                 # Convert system to developer for gpt-oss instruction hierarchy
-                result.append({
-                    "role": "developer",
-                    "content": msg.content,
-                })
+                result.append(
+                    {
+                        "role": "developer",
+                        "content": msg.content,
+                    }
+                )
             elif isinstance(msg, HumanMessage):
-                result.append({
-                    "role": "user",
-                    "content": msg.content,
-                })
+                result.append(
+                    {
+                        "role": "user",
+                        "content": msg.content,
+                    }
+                )
             elif isinstance(msg, AIMessage):
                 ai_msg: Dict[str, Any] = {
                     "role": "assistant",
@@ -154,25 +161,31 @@ class ChatGPTOSS(BaseChatModel):
                             "type": "function",
                             "function": {
                                 "name": tc["name"],
-                                "arguments": json.dumps(tc["args"]) if isinstance(tc["args"], dict) else tc["args"],
+                                "arguments": json.dumps(tc["args"])
+                                if isinstance(tc["args"], dict)
+                                else tc["args"],
                             },
                         }
                         for tc in tool_calls
                     ]
                 result.append(ai_msg)
             elif isinstance(msg, ToolMessage):
-                result.append({
-                    "role": "tool",
-                    "tool_call_id": msg.tool_call_id,
-                    "content": msg.content,
-                })
+                result.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": msg.tool_call_id,
+                        "content": msg.content,
+                    }
+                )
             else:
                 # Fallback for other message types
                 role = getattr(msg, "role", "user")
-                result.append({
-                    "role": role,
-                    "content": msg.content,
-                })
+                result.append(
+                    {
+                        "role": role,
+                        "content": msg.content,
+                    }
+                )
 
         return result
 
@@ -204,12 +217,14 @@ class ChatGPTOSS(BaseChatModel):
                     args = json.loads(tc.function.arguments)
                 except json.JSONDecodeError:
                     args = {"raw": tc.function.arguments}
-                tool_calls_list.append({
-                    "name": tc.function.name,
-                    "args": args,
-                    "id": tc.id,
-                    "type": "tool_call",
-                })
+                tool_calls_list.append(
+                    {
+                        "name": tc.function.name,
+                        "args": args,
+                        "id": tc.id,
+                        "type": "tool_call",
+                    }
+                )
 
         ai_message = AIMessage(
             content=content,
@@ -260,7 +275,9 @@ class ChatGPTOSS(BaseChatModel):
                 request_kwargs["tool_choice"] = self._tool_choice
 
         # Make API call
-        logger.debug(f"ChatGPTOSS request: model={self.model}, messages_count={len(openai_messages)}")
+        logger.debug(
+            f"ChatGPTOSS request: model={self.model}, messages_count={len(openai_messages)}"
+        )
         response = self.client.chat.completions.create(**request_kwargs)
 
         return self._create_chat_result(response)
@@ -328,12 +345,14 @@ class ChatGPTOSS(BaseChatModel):
                             tool_calls_accum[idx]["arguments"] += tc.function.arguments
 
                     # Build tool call chunk for LangChain
-                    tool_call_chunks.append({
-                        "index": idx,
-                        "id": tool_calls_accum[idx]["id"],
-                        "name": tool_calls_accum[idx]["name"],
-                        "args": tool_calls_accum[idx]["arguments"],
-                    })
+                    tool_call_chunks.append(
+                        {
+                            "index": idx,
+                            "id": tool_calls_accum[idx]["id"],
+                            "name": tool_calls_accum[idx]["name"],
+                            "args": tool_calls_accum[idx]["arguments"],
+                        }
+                    )
 
             # Create chunk message
             chunk_message = AIMessageChunk(