Comparing Pydantic

1
class Resume(BaseModel):
2
    name: str
3
    skills: List[str]
4
5
def create_prompt(input_text: str) -> str:
6
    PROMPT_TEMPLATE = f"""Parse the following resume and return a structured representation of the data in the schema below.
7
Resume:
8
---
9
{input_text}
10
---
11
12
Schema:
13
{Resume.model_json_schema()['properties']}
14
15
Output JSON:
16
"""
17
    return PROMPT_TEMPLATE
18
19
def extract_resume(input_text: str) -> Union[Resume, None]:
20
    prompt = create_prompt(input_text)
21
    chat_completion = client.chat.completions.create(
22
        model="gpt-4", messages=[{"role": "system", "content": prompt}]
23
    )
24
    try:
25
        output = chat_completion.choices[0].message.content
26
        if output:
27
            return Resume.model_validate_json(output)
28
        return None
29
    except Exception as e:
30
        raise e

1
...
2
+class Education(BaseModel):
3
+    school: str
4
+    degree: str
5
+    year: int
6
7
class Resume(BaseModel):
8
    name: str
9
    skills: List[str]
10
+   education: List[Education]
11
12
def create_prompt(input_text: str) -> str:
13
    additional_models = ""
14
+    if "$defs" in Resume.model_json_schema():
15
+        additional_models += f"\nUse these other schema definitions as +well:\n{Resume.model_json_schema()['$defs']}"
16
    PROMPT_TEMPLATE = f"""Parse the following resume and return a structured representation of the data in the schema below.
17
Resume:
18
---
19
{input_text}
20
---
21
22
Schema:
23
{Resume.model_json_schema()['properties']}
24
25
+ {additional_models}
26
27
Output JSON:
28
""".strip()
29
    return PROMPT_TEMPLATE
30
...

1
+ The output is:
2
{
3
  "name": "John Doe",
4
  ... // truncated for brevity
5
}

1
def extract_resume(input_text: str) -> Union[Resume, None]:
2
    prompt = create_prompt(input_text)
3
    print(prompt)
4
    chat_completion = client.chat.completions.create(
5
        model="gpt-4", messages=[{"role": "system", "content": prompt}]
6
    )
7
    try:
8
        output = chat_completion.choices[0].message.content
9
        print(output)
10
        if output:
11
+            # Extract JSON block using regex
12
+            json_match = re.search(r"\{.*?\}", output, re.DOTALL)
13
+            if json_match:
14
+                json_output = json_match.group(0)
15
                return Resume.model_validate_json(output)
16
        return None
17
    except Exception as e:
18
        raise e

1
+class ResumeArray(BaseModel):
2
+    resumes: List[Resume]

1
+def extract_resume(input_text: str) -> Union[List[Resume], None]:
2
+    prompt = create_prompt(input_text) # Also requires changes
3
    chat_completion = client.chat.completions.create(
4
        model="gpt-4", messages=[{"role": "system", "content": prompt}]
5
    )
6
    try:
7
        output = chat_completion.choices[0].message.content
8
        if output:
9
            # Extract JSON block using regex
10
            json_match = re.search(r"\{.*?\}", output, re.DOTALL)
11
            if json_match:
12
                json_output = json_match.group(0)
13
                try:
14
+                  parsed = json.loads(json_output)
15
+                  if isinstance(parsed, list):
16
+                      return list(map(Resume.model_validate_json, parsed))
17
+                  else:
18
+                      return [ResumeArray(**parsed)]
19
        return None
20
    except Exception as e:
21
        raise e

JSON Schema:
{'name': {'title': 'Name', 'type': 'string'}, 'skills': {'items': {'type': 'string'}, 'title': 'Skills', 'type': 'array'}, 'education': {'anyOf': [{'$ref': '#/$defs/Education'}, {'type': 'null'}]}}
Use these other JSON schema definitions as well:
{'Education': {'properties': {'degree': {'title': 'Degree', 'type': 'string'}, 'major': {'title': 'Major', 'type': 'string'}, 'school': {'title': 'School', 'type': 'string'}, 'year': {'title': 'Year', 'type': 'integer'}}, 'required': ['degree', 'major', 'school', 'year'], 'title': 'Education', 'type': 'object'}}

{
  "name": 
  {
    "title": "Name", 
    "type": "string", 
    "value": "John Doe"
  }, 
  "skills": 
  {
    "items": 
    {
      "type": "string", 
      "values": 
      [
        "Python", 
        "JavaScript", 
        "React"
      ]
    ... // truncated for brevity

1
Parse the following resume and return a structured representation of the data in the schema below.
2
Resume:
3
---
4
John Doe
5
Python, Rust
6
University of California, Berkeley, B.S. in Computer Science, 2020
7
---
8
9
+JSON Schema:
10
+{
11
+  "name": string,
12
+  "skills": string[]
13
+  "education": {
14
+    "school": string,
15
+    "degree": string,
16
+    "year": integer
17
+  }[]
18
+}
19
20
Output JSON:

Classify the company described in this text into the best
of the following categories:
Text:
---
{some_text}
---
Categories:
- Technology: Companies involved in the development and production of technology products or services
- Healthcare: Includes companies in pharmaceuticals, biotechnology, medical devices.
- Real estate: Includes real estate investment trusts (REITs) and companies involved in real estate development.
The best category is:

1
class FinancialCategory(Enum):
2
    technology = (
3
        "Technology",
4
        "Companies involved in the development and production of technology products or services.",
5
    )
6
    ...
7
    real_estate = (
8
        "Real Estate",
9
        "Includes real estate investment trusts (REITs) and companies involved in real estate development.",
10
    )
11
12
    def __init__(self, category, description):
13
        self._category = category
14
        self._description = description
15
16
    @property
17
    def category(self):
18
        return self._category
19
20
    @property
21
    def description(self):
22
        return self._description

1
    @classmethod
2
    def from_string(cls, category: str) -> "FinancialCategory":
3
        for c in cls:
4
            if c.category == category:
5
                return c
6
        raise ValueError(f"Invalid category: {category}")

1
def print_categories_and_descriptions():
2
    for category in FinancialCategory:
3
        print(f"{category.category}: {category.description}")
4
5
def create_prompt(text: str) -> str:
6
    additional_models = ""
7
    print_categories_and_descriptions()
8
    PROMPT_TEMPLATE = f"""Classify the company described in this text into the best
9
of the following categories:
10
11
Text:
12
---
13
{text}
14
---
15
16
Categories:
17
{print_categories_and_descriptions()}
18
19
The best category is:
20
"""
21
    return PROMPT_TEMPLATE

1
def classify_company(text: str) -> FinancialCategory:
2
    prompt = create_prompt(text)
3
    chat_completion = client.chat.completions.create(
4
        model="gpt-4", messages=[{"role": "system", "content": prompt}]
5
    )
6
    try:
7
        output = chat_completion.choices[0].message.content
8
        if output:
9
            # Use our helper function!
10
            return FinancialCategory.from_string(output)
11
        return None
12
    except Exception as e:
13
        raise e

1
from baml_client import baml as b
2
3
async def main():
4
  resume = await b.ExtractResume(resume_text="""John Doe
5
Python, Rust
6
University of California, Berkeley, B.S. in Computer Science, 2020""")
7
8
  assert resume.name == "John Doe"