clean up

Author: rgbkrk

chatlab/__init__.py

@@ -32,6 +32,7 @@
 )
 from .registry import FunctionRegistry
 from spork import Markdown
+from instructor import Partial
 
 __version__ = __version__
 
@@ -51,4 +52,5 @@
     "FunctionRegistry",
     "ChatlabMetadata",
     "expose_exception_to_llm",
+    "Partial"
 ]

chatlab/chat.py

@@ -172,9 +172,9 @@ async def __process_stream(
                             tool_argument = ToolArguments(
                                 id=tool_call.id, name=tool_call.function.name, arguments=tool_call.function.arguments
                             )
-                            # Now we get the function and see if it has the ChatLabMetadata for a render func
-                            func = self.function_registry.get_chatlab_metadata(tool_call.function.name)
 
+                            # If the user provided a custom renderer, set it on the tool argument object for displaying
+                            func = self.function_registry.get_chatlab_metadata(tool_call.function.name)
                             if func is not None and func.render is not None:
                                 tool_argument.custom_render = func.render
 

chatlab/registry.py

@@ -65,7 +65,7 @@ class WhatTime(BaseModel):
 from openai.types.chat import ChatCompletionToolParam
 
 from .decorators import ChatlabMetadata
-
+from .errors import ChatLabError
 
 class APIManifest(TypedDict, total=False):
     """The schema for the API."""
@@ -80,13 +80,13 @@ class APIManifest(TypedDict, total=False):
     """
 
 
-class FunctionArgumentError(Exception):
+class FunctionArgumentError(ChatLabError):
     """Exception raised when a function is called with invalid arguments."""
 
     pass
 
 
-class UnknownFunctionError(Exception):
+class UnknownFunctionError(ChatLabError):
     """Exception raised when a function is called that is not registered."""
 
     pass

chatlab/views/tools.py

@@ -2,6 +2,8 @@
 from pydantic import ValidationError
 from spork import AutoUpdate
 
+import warnings
+
 from ..components.function_details import ChatFunctionComponent
 
 from ..registry import FunctionRegistry, FunctionArgumentError, UnknownFunctionError, extract_arguments, extract_model_from_function
@@ -101,23 +103,23 @@ def render(self):
                 possible_args = parser.parse(self.arguments)
 
                 Model = extract_model_from_function(self.name, self.custom_render)
-                model = Partial[Model](**possible_args)
-
-                kwargs = {}
+                model = Model.model_validate(possible_args)
 
                 # Pluck the kwargs out from the crafted model, as we can't pass the pydantic model as the arguments
                 # However any "inner" models should retain their pydantic Model nature
-                for k in model.__dict__.keys():
-                    kwargs[k] = getattr(model, k)
+                kwargs = {k: getattr(model, k) for k in model.__dict__.keys()}
 
-                return self.custom_render(**kwargs)
             except FunctionArgumentError:
                 return None
             except ValidationError:
                 return None
+
+            try:
+                return self.custom_render(**kwargs)
             except Exception as e:
-                #print(f"Exception in custom render for {self.name}.", e)
-                #print(self.arguments)
+                # Exception in userland code
+                # Would be preferable to bubble up, however 
+                # it might be due to us passing a not-quite model
                 return None
 
         return ChatFunctionComponent(name=self.name, verbage=self.verbage, input=self.arguments)