Upgrade to pydanticv2 (#63)

* Upgrade to Pydantic 2 #15

* update examples to use pydantic2

---------

Co-authored-by: Mike Harris <[email protected]>

examples/automatic_dataframe_extraction/auto_multi_dataframe.py

@@ -5,18 +5,10 @@
 
 
 class RowData(OpenAISchema):
- row: List[Any] = Field(..., description="The values for each row")
- citation: str = Field(
- ..., description="The citation for this row from the original source data"
- )
+ row: List[Any] = Field(..., description="Correct values for each row")
 
 
 class Dataframe(OpenAISchema):
- """
- Class representing a dataframe. This class is used to convert
- data into a frame that can be used by pandas.
- """
-
  name: str = Field(..., description="The name of the dataframe")
  data: List[RowData] = Field(
  ...,
@@ -30,15 +22,17 @@ class Dataframe(OpenAISchema):
  def to_pandas(self):
  import pandas as pd
 
- columns = self.columns + ["citation"]
- data = [row.row + [row.citation] for row in self.data]
+ columns = self.columns
+ data = [row.row for row in self.data]
 
  return pd.DataFrame(data=data, columns=columns)
 
 
 class Database(OpenAISchema):
  """
  A set of correct named and defined tables as dataframes
+ Each one should have the right number of columns and correct
+ values for each.
  """
 
  tables: List[Dataframe] = Field(
@@ -50,7 +44,7 @@ class Database(OpenAISchema):
 def dataframe(data: str) -> Database:
  completion = openai.ChatCompletion.create(
  model="gpt-4-0613",
- temperature=0.1,
+ temperature=0.0,
  functions=[Database.openai_schema],
  function_call={"name": Database.openai_schema["name"]},
  messages=[
@@ -89,14 +83,13 @@ def dataframe(data: str) -> Database:
  print(df.to_pandas())
  """
  People
- Name Age City Favorite Sport
- 0 John 25 New York Basketball
- 1 Mike 30 San Francisco Baseball
- 2 Sarah 20 Los Angeles Tennis
- 3 Mary 35 Chicago None
-
+ ID Name Age City Favorite Sport
+ 0 1 John 25 New York Basketball
+ 1 2 Mike 30 San Francisco Baseball
+ 2 3 Sarah 20 Los Angeles Tennis
+ 3 4 Mary 35 Chicago None
  Teams
- Team Name Captain Number of Players
- 0 Tigers John 12
- 1 Lions Mike 10
+ ID Team Name Captain Number of Players
+ 0 1  Tigers John 12
+ 1 2  Lions Mike 10
  """

examples/query_planner_execution/query_planner_execution.py

@@ -65,20 +65,20 @@ async def execute(self, dependency_func):
  query=self.question,
  )
  await asyncio.sleep(1)
- pprint(resp.dict())
+ pprint(resp.model_dump())
  return resp
 
  sub_queries = dependency_func(self.dependancies)
  computed_queries = await asyncio.gather(
  *[q.execute(dependency_func=dependency_func) for q in sub_queries]
  )
  sub_answers = MergedResponses(responses=computed_queries)
- merged_query = f"{self.question}\nContext: {sub_answers.json()}"
+ merged_query = f"{self.question}\nContext: {sub_answers.model_dump_json()}"
  resp = ComputeQuery(
  query=merged_query,
  )
  await asyncio.sleep(2)
- pprint(resp.dict())
+ pprint(resp.model_dump())
  return resp
 
 
@@ -105,8 +105,8 @@ def dependencies(self, idz: List[int]) -> List[Query]:
  return [q for q in self.query_graph if q.id in idz]
 
 
-Query.update_forward_refs()
-QueryPlan.update_forward_refs()
+Query.model_rebuild()
+QueryPlan.model_rebuild()
 
 
 def query_planner(question: str, plan=False) -> QueryPlan:

examples/recursive_filepaths/parse_recursive_paths.py

@@ -1,45 +1,3 @@
-"""
-This script parses a string representation of a filesystem structure into a tree-like directory structure.
-
-The 'Node' class represents a node in this tree, which can be either a file or a folder. Files cannot have 
-children, while folders can.
-
-The 'DirectoryTree' class contains a single root folder from which all other files/folders can be reached.
-The 'parse_tree_to_filesystem' function uses OpenAI's GPT-3 model to convert a string representation of a 
-directory tree into a 'DirectoryTree' object. This object can then be manipulated programmatically as needed,
-with methods such as 'print_paths' available for convenience.
-
-Please note: Recursive models currently work if they are wrapped by a non-recursive one. This is why we are
-passing a 'DirectoryTree' (which contains a single 'Node') as the function call, not a 'Node' directly. This
-is due to a limitation in how Pydantic generates schemas for recursive objects, which creates 
-'dict_keys(['$ref', 'definitions'])'. Instead of writing a resolver for such references, we can simply wrap the 
-recursive class in a non-recursive one so the function_call class never has a cyclic reference.
-
-Example usage:
->>> root = parse_tree_to_filesystem(
-... '''
-... root
-... ├── folder1
-... │ ├── file1.txt
-... │ └── file2.txt
-... └── folder2
-... ├── file3.txt
-... └── subfolder1
-... └── file4.txt
-... '''
-... )
->>> root.print_paths()
-# Expected output:
-# >>> root NodeType.FOLDER
-# >>> root/folder1 NodeType.FOLDER
-# >>> root/folder1/file1.txt NodeType.FILE
-# >>> root/folder1/file2.txt NodeType.FILE
-# >>> root/folder2 NodeType.FOLDER
-# >>> root/folder2/file3.txt NodeType.FILE
-# >>> root/folder2/subfolder1 NodeType.FOLDER
-# >>> root/folder2/subfolder1/file4.txt NodeType.FILE
-"""
-
 import enum
 from typing import List
 
@@ -115,8 +73,8 @@ def print_paths(self):
  self.root.print_paths()
 
 
-Node.update_forward_refs()
-DirectoryTree.update_forward_refs()
+Node.model_rebuild()
+DirectoryTree.model_rebuild()
 
 
 @retry(stop=stop_after_attempt(3))

openai_function_call/function_calls.py

@@ -22,8 +22,8 @@
 
 import json
 from functools import wraps
-from typing import Any, Callable, Optional, List, Type
-from pydantic import validate_arguments, BaseModel, create_model, Field
+from typing import Any, Callable
+from pydantic import BaseModel, validate_arguments
 
 
 def _remove_a_key(d, remove_key) -> None:
@@ -65,7 +65,7 @@ def sum(a: int, b: int) -> int:
  def __init__(self, func: Callable) -> None:
  self.func = func
  self.validate_func = validate_arguments(func)
- parameters = self.validate_func.model.schema()
+ parameters = self.validate_func.model.model_json_schema()
  parameters["properties"] = {
  k: v
  for k, v in parameters["properties"].items()
@@ -127,7 +127,7 @@ def openai_schema(cls):
  Returns:
  model_json_schema (dict): A dictionary in the format of OpenAI's schema as jsonschema
  """
- schema = cls.schema()
+ schema = cls.model_json_schema()
  parameters = {
  k: v for k, v in schema.items() if k not in ("title", "description")
  }