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

Node Start:
{'input': {0: None}, 'output': None}
================================================================================
Node get_missing_keys:
{'input': {0: None}, 'output': ['name', 'age']}
================================================================================
Node ask_missing_data:
{ 'input': {'current_info': None, 'missing_keys': ['name', 'age']},
  'output': { 'content': 'Please provide your name and age.',
              'meta': {'input_tokens': 151, 'output_tokens': 11},
              'role': 'assistant'}}
================================================================================

{'role': 'assistant',
 'content': 'Please provide your name and age.',
 'meta': {'input_tokens': 151, 'output_tokens': 11}}

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

Node break:
{'input': {0: 'my name is brian'}, 'output': 'my name is brian'}
================================================================================
Node format_new_info:
{ 'input': { 'current_info': None,
             'input': 'my name is brian',
             'missing_keys': ['name', 'age']},
  'output': { 'content': User(name='brian', age=None, email=None),
              'meta': {'input_tokens': 342, 'output_tokens': 22},
              'role': 'assistant'}}
================================================================================
Node get_missing_keys:
{'input': {0: User(name='brian', age=None, email=None)}, 'output': ['age']}
================================================================================
Node ask_missing_data:
{ 'input': { 'current_info': { 'content': User(name='brian', age=None, email=None),
                               'meta': { 'input_tokens': 342,
                                         'output_tokens': 22},
                               'role': 'assistant'},
             'missing_keys': ['age']},
  'output': { 'content': 'Please provide your age.',
              'meta': {'input_tokens': 148, 'output_tokens': 9},
              'role': 'assistant'}}
================================================================================

{'role': 'assistant',
 'content': 'Please provide your age.',
 'meta': {'input_tokens': 148, 'output_tokens': 9}}

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

Node break:
{'input': {0: '25 yo'}, 'output': '25 yo'}
================================================================================
Node format_new_info:
{ 'input': { 'current_info': { 'content': { 'age': None,
                                            'email': None,
                                            'name': 'brian'},
                               'meta': { 'input_tokens': 342,
                                         'output_tokens': 22},
                               'role': 'assistant'},
             'input': '25 yo',
             'missing_keys': ['age']},
  'output': { 'content': User(name='brian', age=25, email=None),
              'meta': {'input_tokens': 388, 'output_tokens': 23},
              'role': 'assistant'}}
================================================================================
Node End:
{ 'input': { 0: { 'content': User(name='brian', age=25, email=None),
                  'meta': {'input_tokens': 388, 'output_tokens': 23},
                  'role': 'assistant'}},
  'output': { 'content': User(name='brian', age=25, email=None),
              'meta': {'input_tokens': 388, 'output_tokens': 23},
              'role': 'assistant'}}
================================================================================

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