From 8e3df8681c40368021a77a77ad802b7cc2c834ac Mon Sep 17 00:00:00 2001
From: cpacker <packercharles@gmail.com>
Date: Fri, 22 Nov 2024 11:07:50 -0800
Subject: [PATCH 1/2] fix: add more types to the type_map from python -> json
 schema

---
 letta/functions/schema_generator.py | 22 ++++++++++++++++++----
 letta/local_llm/utils.py            |  2 +-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/letta/functions/schema_generator.py b/letta/functions/schema_generator.py
index 92894d4c45..01d87157fb 100644
--- a/letta/functions/schema_generator.py
+++ b/letta/functions/schema_generator.py
@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Dict, Optional, Type, Union, get_args, get_origin
+from typing import Any, Dict, List, Optional, Type, Union, get_args, get_origin
 
 from docstring_parser import parse
 from pydantic import BaseModel
@@ -38,15 +38,29 @@ def type_to_json_schema_type(py_type):
 
     # Mapping of Python types to JSON schema types
     type_map = {
+        # Basic types
         int: "integer",
         str: "string",
         bool: "boolean",
         float: "number",
-        list[str]: "array",
-        # Add more mappings as needed
+        # Collections
+        List[str]: "array",
+        List[int]: "array",
+        list: "array",
+        tuple: "array",
+        set: "array",
+        # Dictionaries
+        dict: "object",
+        Dict[str, Any]: "object",
+        # Special types
+        None: "null",
+        type(None): "null",
+        # Optional types
+        Optional[str]: "string",
+        Union[str, None]: "string",
     }
     if py_type not in type_map:
-        raise ValueError(f"Python type {py_type} has no corresponding JSON schema type")
+        raise ValueError(f"Python type {py_type} has no corresponding JSON schema type - full map: {type_map}")
 
     return type_map.get(py_type, "string")  # Default to "string" if type not in map
 
diff --git a/letta/local_llm/utils.py b/letta/local_llm/utils.py
index 8b91f4b397..9d625e75b2 100644
--- a/letta/local_llm/utils.py
+++ b/letta/local_llm/utils.py
@@ -88,7 +88,7 @@ def num_tokens_from_functions(functions: List[dict], model: str = "gpt-4"):
     try:
         encoding = tiktoken.encoding_for_model(model)
     except KeyError:
-        print("Warning: model not found. Using cl100k_base encoding.")
+        warnings.warn("Warning: model not found. Using cl100k_base encoding.")
         encoding = tiktoken.get_encoding("cl100k_base")
 
     num_tokens = 0

From 36598bbde22d86aa10adc06008a99d2a51a82b90 Mon Sep 17 00:00:00 2001
From: cpacker <packercharles@gmail.com>
Date: Fri, 22 Nov 2024 12:12:29 -0800
Subject: [PATCH 2/2] fix: remove optional

---
 letta/functions/schema_generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/letta/functions/schema_generator.py b/letta/functions/schema_generator.py
index 01d87157fb..2969f6f8a9 100644
--- a/letta/functions/schema_generator.py
+++ b/letta/functions/schema_generator.py
@@ -56,7 +56,7 @@ def type_to_json_schema_type(py_type):
         None: "null",
         type(None): "null",
         # Optional types
-        Optional[str]: "string",
+        # Optional[str]: "string",  # NOTE: caught above ^
         Union[str, None]: "string",
     }
     if py_type not in type_map: