obscura/ollama-python
1
0
Fork 0
mirror of https://github.com/ollama/ollama-python.git synced 2026-11-05 23:32:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplified parsing and fixed required - added tests

Browse Source
This commit is contained in:
ParthSareen
2024-11-17 23:52:13 -08:00
parent e7bb55f5b4
commit 7396ab6cdd
2 changed files with 18 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ollama/_utils.py
+14 -10
View File
@@ -69,18 +69,22 @@ def convert_function_to_tool(func: Callable) -> Tool:
required = []
parsed_docstring = _parse_docstring(schema.get('description'))
for k, v in schema.get('properties', {}).items():
properties[k] = {}
properties[k]['description'] = parsed_docstring.get(k, '')
prop = {
'description': parsed_docstring.get(k, ''),
'type': v.get('type'),
}
if 'anyOf' in v:
anyof_list = v['anyOf']
properties[k]['type'] = [anyof_type.get('type', 'string') for anyof_type in anyof_list if anyof_type.get('type', 'string') != 'null']
if len(properties[k]['type']) == 1:
properties[k]['type'] = properties[k]['type'][0]
else:
properties[k]['type'] = str(properties[k]['type'])
is_optional = any(t.get('type') == 'null' for t in v['anyOf'])
types = [t.get('type', 'string') for t in v['anyOf'] if t.get('type') != 'null']
prop['type'] = types[0] if len(types) == 1 else str(types)
if not is_optional:
required.append(k)
else:
properties[k]['type'] = v.get('type', None)
required.append(k)
if prop['type'] != 'null':
required.append(k)
properties[k] = prop
schema['properties'] = properties
tests/test_utils.py
+4 -1
View File
@@ -103,6 +103,8 @@ def test_function_with_all_types():
assert tool['function']['parameters']['properties']['z']['type'] == 'array'
assert tool['function']['parameters']['properties']['w']['type'] == 'object'
assert set(x.strip().strip("'") for x in tool['function']['parameters']['properties']['v']['type'].removeprefix('[').removesuffix(']').split(',')) == {'string', 'integer'}
assert tool['function']['parameters']['properties']['v']['type'] != 'null'
assert tool['function']['parameters']['required'] == ['x', 'y', 'z', 'w']
else:
assert tool['function']['parameters']['properties']['z']['type'] == 'array'
assert tool['function']['parameters']['properties']['w']['type'] == 'object'
@@ -110,8 +112,9 @@ def test_function_with_all_types():
assert tool['function']['parameters']['properties']['s']['type'] == 'array'
assert tool['function']['parameters']['properties']['t']['type'] == 'array'
assert tool['function']['parameters']['properties']['l']['type'] == 'array'
print('hi')
assert tool['function']['parameters']['properties']['o']['type'] == 'integer'
assert tool['function']['parameters']['properties']['o']['type'] != 'null'
assert tool['function']['parameters']['required'] == ['x', 'y', 'z', 'w', 'd', 's', 't', 'l']
def test_function_docstring_parsing():