removed redundancy, added formatting

This commit is contained in:
nicole pardal 2025-09-23 12:00:03 -07:00
parent 5a3b22ad55
commit 81c71ba44e

View File

@ -15,55 +15,35 @@ from ollama import Client, WebFetchResponse, WebSearchResponse
def format_tool_results(
results: Union[WebSearchResponse, WebFetchResponse],
*,
query: Optional[str] = None,
url: Optional[str] = None,
user_search: str,
):
if isinstance(results, WebSearchResponse):
output = []
if isinstance(results.results, dict):
for q, search_results in results.results.items():
output.append(f'Search results for "{q}":')
for i, result in enumerate(search_results, 1):
title = getattr(result, 'title', None)
url_value = getattr(result, 'url', None)
output.append(f'{i}. {title}' if title else f'{i}. {getattr(result, "content", "")}')
if url_value:
output.append(f' URL: {url_value}')
output.append(f' Content: {getattr(result, "content", "")}')
output.append('')
else:
if query:
output.append(f'Search results for "{query}":')
for i, result in enumerate(results.results, 1):
title = getattr(result, 'title', None)
url_value = getattr(result, 'url', None)
output.append(f'{i}. {title}' if title else f'{i}. {getattr(result, "content", "")}')
if url_value:
output.append(f' URL: {url_value}')
output.append(f' Content: {getattr(result, "content", "")}')
output.append('')
output.append(f'Search results for "{user_search}":')
print()
for i, result in enumerate(results.results, 1):
title = getattr(result, 'title', None)
url_value = getattr(result, 'url', None)
output.append(f'{i}. {title}' if title else f'{i}. {getattr(result, "content", "")}')
if url_value:
output.append(f' URL: {url_value}')
output.append(f' Content: {getattr(result, "content", "")}')
output.append('')
return '\n'.join(output).rstrip()
elif isinstance(results, WebFetchResponse):
output = []
if url:
output.append(f'Fetch results for "{url}":')
output.extend(
[
f'Title: {results.title}',
f'URL: {url}' if url else '',
f'Content: {results.content}',
]
)
output.append(f'Fetch results for "{user_search}":')
output.extend([
f'Title: {results.title}',
f'URL: {user_search}' if user_search else '',
f'Content: {results.content}',
])
if results.links:
output.append(f'Links: {", ".join(results.links)}')
output.append('')
return '\n'.join(output).rstrip()
client = Client(headers={'Authorization': (os.getenv('OLLAMA_API_KEY'))})
available_tools = {'web_search': client.web_search, 'web_fetch': client.web_fetch}
@ -88,19 +68,24 @@ while True:
if function_to_call:
args = tool_call.function.arguments
result: Union[WebSearchResponse, WebFetchResponse] = function_to_call(**args)
print('Result from tool call name: ', tool_call.function.name, 'with arguments: ', args)
print('Result from tool call name:', tool_call.function.name, 'with arguments:')
print(args)
print()
user_search = args.get('query', '') or args.get('url', '')
if tool_call.function.name == 'web_search':
formatted = format_tool_results(result, query=args.get('query'))
formatted_tool_results = format_tool_results(result, user_search=user_search)
elif tool_call.function.name == 'web_fetch':
formatted = format_tool_results(result, url=args.get('url'))
formatted_tool_results = format_tool_results(result, user_search=user_search)
else:
formatted = format_tool_results(result)
formatted_tool_results = format_tool_results(result)
print('Result: ', formatted[:200])
print('Result:')
print(formatted_tool_results[:200])
print()
# caps the result at ~2000 tokens
messages.append({'role': 'tool', 'content': formatted[: 2000 * 4], 'tool_name': tool_call.function.name})
messages.append({'role': 'tool', 'content': formatted_tool_results[: 2000 * 4], 'tool_name': tool_call.function.name})
else:
print(f'Tool {tool_call.function.name} not found')
messages.append({'role': 'tool', 'content': f'Tool {tool_call.function.name} not found', 'tool_name': tool_call.function.name})