Filling out missing user data

from stringdale import Define,V,E, Condition
from stringdale.chat import Chat

from pydantic import BaseModel,Field
from typing import Optional

class User(BaseModel):
    name: Optional[str] = Field(None, description='The name of the user')
    age: Optional[int] = Field(None, description='The age of the user')
    email: Optional[str] = Field(None, description='The email of the user')

ask_missing_data = Chat(
    model='gpt-4o-mini',
    messages=[
        {'role':'system','content':'''
        You are a helpful assistant asks the user for missing information.
        Do not ask for keys if they are not a part of the missing keys
        Ask the user for the following missing keys:
        {{missing_keys}}
        '''}])

format_new_info = Chat(
    model='gpt-4o-mini',
    messages=[
        {'role':'system','content':'''
        You are a helpful assistant that gets user data and makes sure it is complete.
        If you are not sure that you were given the relevant information, put None in the relevant field.
        Fill part of the missing keys based on the user data.
        If the user didnt provide info an a given missing key, leave it empty
        
        {% if current_info %}
        The information we have so far is:
        {{current_info}}
        {% endif %}
        The missing keys are:
        {{missing_keys}}
        '''},
        {'role':'user','content':'{{input}}'},
        
    ],
    output_schema=User)

class MissingKeys():
    def __init__(self,keys):
        self.keys = keys

    def get_missing_keys(self,obj):
        return [key for key in self.keys if getattr(obj,key,None) is None]

    def has_missing_keys(self,obj):
        return len(self.get_missing_keys(obj)) > 0

missing = MissingKeys(['name','age'])

with Define('collect user data',type='decision') as UserQA:
    V('Start',
        outputs=[
            ('get_missing_keys',missing.has_missing_keys),
            ('End',),
            ('state/current_info',)
        ])

    V('get_missing_keys',missing.get_missing_keys,
        outputs=[
            'ask_missing_data(missing_keys)',
            'state/missing_keys'

        ])

    V('ask_missing_data',ask_missing_data,
        inputs=['state/current_info(current_info)'],
        outputs=['break']    )

    V('break',is_break=True,outputs=['format_new_info(input)'])

    
    V('format_new_info',format_new_info,
        inputs=[
            'state/current_info(current_info)',
            'state/missing_keys(missing_keys=.)'
        ],
        outputs=[
            ('get_missing_keys(0=content)',Condition(missing.has_missing_keys,'(0=content)')),
            'End',
            'state/current_info'
        ])

UserQA.draw(direction='TB')

d = UserQA()

for trace in d.run(None):
    trace.pprint()
d.output

---
name: Start
input:
  0: null
output: null

---
name: get_missing_keys
input:
  0: null
output: [name, age]

---
name: ask_missing_data
input:
  missing_keys: [name, age]
  current_info: null
output:
  role: assistant
  content: Could you please provide your name and age?
  meta:
    input_tokens: 50
    output_tokens: 9

{'role': 'assistant',
 'content': 'Could you please provide your name and age?',
 'meta': {'input_tokens': 50, 'output_tokens': 9}}

for trace in d.run('my name is brian'):
    trace.pprint()
d.output

---
name: break
input:
  0: my name is brian
output: my name is brian

---
name: format_new_info
input:
  input: my name is brian
  current_info: null
  missing_keys: [name, age]
output:
  role: assistant
  content: name='brian' age=None email=None
  meta:
    input_tokens: 342
    output_tokens: 22

---
name: get_missing_keys
input:
  0: name='brian' age=None email=None
output: [age]

---
name: ask_missing_data
input:
  missing_keys: [age]
  current_info:
    role: assistant
    content: name='brian' age=None email=None
    meta:
      input_tokens: 342
      output_tokens: 22
output:
  role: assistant
  content: Could you please provide your age?
  meta:
    input_tokens: 47
    output_tokens: 7

{'role': 'assistant',
 'content': 'Could you please provide your age?',
 'meta': {'input_tokens': 47, 'output_tokens': 7}}

for trace in d.run('25 yo'):
    trace.pprint()
d.output

---
name: break
input:
  0: 25 yo
output: 25 yo

---
name: format_new_info
input:
  input: 25 yo
  current_info:
    role: assistant
    content: name='brian' age=None email=None
    meta:
      input_tokens: 342
      output_tokens: 22
  missing_keys: [age]
output:
  role: assistant
  content: name='brian' age=25 email=None
  meta:
    input_tokens: 383
    output_tokens: 23

---
name: End
input:
  0:
    role: assistant
    content: name='brian' age=25 email=None
    meta:
      input_tokens: 383
      output_tokens: 23
output:
  role: assistant
  content: name='brian' age=25 email=None
  meta:
    input_tokens: 383
    output_tokens: 23

{'role': 'assistant',
 'content': User(name='brian', age=25, email=None),
 'meta': {'input_tokens': 383, 'output_tokens': 23}}