Use a more general name (#31)

Fixed tool call to work with vLLM and OCI DataScience Model Deployment API

Co-authored-by: Anup Ojah <[email protected]>

tests/unit_tests/test_extraction.py

@@ -428,8 +428,8 @@ async def test_e2e_existing_schema_policy_behavior(strict_mode):
     class MyRecognizedSchema(BaseModel):
         """A recognized schema that the pipeline can handle."""
 
-        user_id: str
-        notes: str
+        user_id: str  # type: ignore
+        notes: str  # type: ignore
 
     # Our existing data includes 2 top-level keys: recognized, unknown
     existing_schemas = {
@@ -543,8 +543,8 @@ async def test_e2e_existing_schema_policy_tuple_behavior(strict_mode):
     class MyRecognizedSchema(BaseModel):
         """A recognized schema that the pipeline can handle."""
 
-        user_id: str
-        notes: str
+        user_id: str  # type: ignore
+        notes: str  # type: ignore
 
     existing_schemas = [
         (

trustcall/_base.py

@@ -664,7 +664,11 @@ def __init__(
         existing_schema_policy: bool | Literal["ignore"] = True,
     ):
         new_tools: list = [PatchDoc]
-        tool_choice = PatchDoc.__name__ if not enable_deletes else "any"
+        tool_choice = (
+            {"type": "function", "function": {"name": "PatchDoc"}}
+            if not enable_deletes
+            else "any"
+        )
         if enable_inserts:  # Also let the LLM know that we can extract NEW schemas.
             tools_ = [
                 schema
@@ -1052,7 +1056,7 @@ def _tear_down(
 
     async def ainvoke(
         self, state: ExtendedExtractState, config: RunnableConfig
-    ) -> dict:
+    ) -> Command[Literal["sync", "__end__"]]:
         """Generate a JSONPatch to correct the validation error and heal the tool call.
 
         Assumptions:
@@ -1075,7 +1079,9 @@ async def ainvoke(
             goto=("sync",),
         )
 
-    def invoke(self, state: ExtendedExtractState, config: RunnableConfig) -> dict:
+    def invoke(
+        self, state: ExtendedExtractState, config: RunnableConfig
+    ) -> Command[Literal["sync", "__end__"]]:
         try:
             msg = self.bound.invoke(state.messages, config)
         except Exception: