Code style updates

This commit is contained in:
ParthSareen 2024-11-18 16:59:32 -08:00
parent 0d9eec0264
commit ed3ba8a734
4 changed files with 66 additions and 69 deletions

View File

@ -297,6 +297,14 @@ class Client(BaseClient):
"""
Create a chat response using the requested model.
Args:
tools (Sequence[Union[Mapping[str, Any], Tool, Callable]]):
A JSON schema as a dict, an Ollama Tool or a Python Function.
Python functions need to follow Google style docstrings to be converted to an Ollama Tool.
For more information, see: https://google.github.io/styleguide/pyguide.html#38-Docstrings
stream (bool): Whether to stream the response.
format (Optional[Literal['', 'json']]): The format of the response.
Raises `RequestError` if a model is not provided.
Raises `ResponseError` if the request could not be fulfilled.
@ -1084,10 +1092,7 @@ def _copy_tools(tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable
return []
for unprocessed_tool in tools:
if callable(unprocessed_tool):
yield convert_function_to_tool(unprocessed_tool)
else:
yield Tool.model_validate(unprocessed_tool)
yield convert_function_to_tool(unprocessed_tool) if callable(unprocessed_tool) else Tool.model_validate(unprocessed_tool)
def _as_path(s: Optional[Union[str, PathLike]]) -> Union[Path, None]:
@ -1162,6 +1167,8 @@ def _parse_host(host: Optional[str]) -> str:
'https://[0001:002:003:0004::1]:56789/path'
>>> _parse_host('[0001:002:003:0004::1]:56789/path/')
'http://[0001:002:003:0004::1]:56789/path'
>>> _parse_host('http://host.docker.internal:11434/path')
'http://host.docker.internal:11434/path'
"""
host, port = host or '', 11434

View File

@ -336,7 +336,7 @@ class ModelDetails(SubscriptableBaseModel):
class ListResponse(SubscriptableBaseModel):
class Model(SubscriptableBaseModel):
name: Optional[str] = None
model: Optional[str] = None
modified_at: Optional[datetime] = None
digest: Optional[str] = None
size: Optional[ByteSize] = None

View File

@ -1,4 +1,5 @@
from __future__ import annotations
from collections import defaultdict
import inspect
from typing import Callable, Union
@ -7,96 +8,73 @@ from ollama._types import Tool
def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]:
parsed_docstring = {'description': ''}
parsed_docstring = defaultdict(str)
if not doc_string:
return parsed_docstring
lowered_doc_string = doc_string.lower()
if 'args:' not in lowered_doc_string:
parsed_docstring['description'] = lowered_doc_string.strip()
return parsed_docstring
key = hash(doc_string)
parsed_docstring[key] = ''
for line in lowered_doc_string.splitlines():
if line.startswith('args:'):
key = 'args'
elif line.startswith('returns:') or line.startswith('yields:') or line.startswith('raises:'):
key = '_'
else:
parsed_docstring['description'] = lowered_doc_string.split('args:')[0].strip()
args_section = lowered_doc_string.split('args:')[1]
if 'returns:' in lowered_doc_string:
# Return section can be captured and used
args_section = args_section.split('returns:')[0]
if 'yields:' in lowered_doc_string:
args_section = args_section.split('yields:')[0]
cur_var = None
for line in args_section.split('\n'):
line = line.strip()
if not line:
continue
if ':' not in line:
# Continuation of the previous parameter's description
if cur_var:
parsed_docstring[cur_var] += f' {line}'
continue
# For the case with: `param_name (type)`: ...
if '(' in line:
param_name = line.split('(')[0]
param_desc = line.split('):')[1]
# For the case with: `param_name: ...`
else:
param_name, param_desc = line.split(':', 1)
# maybe change to a list and join later
parsed_docstring[key] += f'{line.strip()}\n'
parsed_docstring[param_name.strip()] = param_desc.strip()
cur_var = param_name.strip()
last_key = None
for line in parsed_docstring['args'].splitlines():
line = line.strip()
if ':' in line and not line.startswith('args'):
# Split on first occurrence of '(' or ':' to separate arg name from description
split_char = '(' if '(' in line else ':'
arg_name, rest = line.split(split_char, 1)
last_key = arg_name.strip()
# Get description after the colon
arg_description = rest.split(':', 1)[1].strip() if split_char == '(' else rest.strip()
parsed_docstring[last_key] = arg_description
elif last_key and line:
parsed_docstring[last_key] += ' ' + line
return parsed_docstring
def convert_function_to_tool(func: Callable) -> Tool:
doc_string_hash = hash(inspect.getdoc(func))
parsed_docstring = _parse_docstring(inspect.getdoc(func))
schema = type(
func.__name__,
(pydantic.BaseModel,),
{
'__annotations__': {k: v.annotation for k, v in inspect.signature(func).parameters.items()},
'__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in inspect.signature(func).parameters.items()},
'__signature__': inspect.signature(func),
'__doc__': inspect.getdoc(func),
'__doc__': parsed_docstring[doc_string_hash],
},
).model_json_schema()
properties = {}
required = []
parsed_docstring = _parse_docstring(schema.get('description'))
for k, v in schema.get('properties', {}).items():
prop = {
'description': parsed_docstring.get(k, ''),
'type': v.get('type'),
# If type is missing, the default is string
types = {t.get('type', 'string') for t in v.get('anyOf')} if 'anyOf' in v else {v.get('type', 'string')}
if 'null' in types:
schema['required'].remove(k)
types.discard('null')
schema['properties'][k] = {
'description': parsed_docstring[k],
'type': ', '.join(types),
}
if 'anyOf' in v:
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:
if prop['type'] != 'null':
required.append(k)
properties[k] = prop
schema['properties'] = properties
tool = Tool(
function=Tool.Function(
name=func.__name__,
description=parsed_docstring.get('description'),
parameters=Tool.Function.Parameters(
type='object',
properties=schema.get('properties', {}),
required=required,
),
description=schema.get('description', ''),
parameters=Tool.Function.Parameters(**schema),
)
)

View File

@ -188,7 +188,7 @@ def test_function_with_only_description():
tool = convert_function_to_tool(only_description).model_dump()
assert tool['function']['description'] == 'a function with only a description.'
assert tool['function']['parameters'] == {'type': 'object', 'properties': {}, 'required': []}
assert tool['function']['parameters'] == {'type': 'object', 'properties': {}, 'required': None}
def only_description_with_args(x: int, y: int):
"""
@ -226,3 +226,15 @@ def test_function_with_yields():
assert tool['function']['description'] == 'a function with yields section.'
assert tool['function']['parameters']['properties']['x']['description'] == 'the first number'
assert tool['function']['parameters']['properties']['y']['description'] == 'the second number'
def test_function_with_no_types():
def no_types(a, b):
"""
A function with no types.
"""
pass
tool = convert_function_to_tool(no_types).model_dump()
assert tool['function']['parameters']['properties']['a']['type'] == 'string'
assert tool['function']['parameters']['properties']['b']['type'] == 'string'