diff --git a/examples/web-search-crawl.py b/examples/web-search-crawl.py index 8a490b4..a139bdc 100644 --- a/examples/web-search-crawl.py +++ b/examples/web-search-crawl.py @@ -2,28 +2,59 @@ # requires-python = ">=3.11" # dependencies = [ # "rich", -# "ollama", +# "ollama @ file:///Users/nicole/Desktop/ollama-python", # ] # /// -from typing import Union +import os +from typing import Optional, Union + +from rich import print from ollama import Client, WebFetchResponse, WebSearchResponse -def format_tool_results(results: Union[WebSearchResponse, WebFetchResponse]): +def format_tool_results( + results: Union[WebSearchResponse, WebFetchResponse], + *, + query: Optional[str] = None, + url: Optional[str] = None, +): if isinstance(results, WebSearchResponse): output = [] - for i, result in enumerate(results.results, 1): - output.append(f'{i}. {result.content}') - output.append('') + 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('') return '\n'.join(output).rstrip() elif isinstance(results, WebFetchResponse): - output = [ + 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}', - ] + ]) if results.links: output.append(f'Links: {", ".join(results.links)}') output.append('') @@ -53,12 +84,21 @@ while True: for tool_call in response.message.tool_calls: function_to_call = available_tools.get(tool_call.function.name) if function_to_call: - result: Union[WebSearchResponse, WebFetchResponse] = function_to_call(**tool_call.function.arguments) - print('Result from tool call name: ', tool_call.function.name, 'with arguments: ', tool_call.function.arguments) - print('Result: ', format_tool_results(result)[:200]) + 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) + + if tool_call.function.name == 'web_search': + formatted = format_tool_results(result, query=args.get('query')) + elif tool_call.function.name == 'web_fetch': + formatted = format_tool_results(result, url=args.get('url')) + else: + formatted = format_tool_results(result) + + print('Result: ', formatted[:200]) # caps the result at ~2000 tokens - messages.append({'role': 'tool', 'content': format_tool_results(result)[: 2000 * 4], 'tool_name': tool_call.function.name}) + messages.append({'role': 'tool', 'content': formatted[: 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})