File: web-search.py

package info (click to toggle)
ollama-python 0.6.1-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 428 kB
  • sloc: python: 3,951; makefile: 5
file content (85 lines) | stat: -rw-r--r-- 2,872 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "rich",
#     "ollama",
# ]
# ///
from typing import Union

from rich import print

from ollama import WebFetchResponse, WebSearchResponse, chat, web_fetch, web_search


def format_tool_results(
  results: Union[WebSearchResponse, WebFetchResponse],
  user_search: str,
):
  output = []
  if isinstance(results, WebSearchResponse):
    output.append(f'Search results for "{user_search}":')
    for result in results.results:
      output.append(f'{result.title}' if result.title else f'{result.content}')
      output.append(f'   URL: {result.url}')
      output.append(f'   Content: {result.content}')
      output.append('')
    return '\n'.join(output).rstrip()

  elif isinstance(results, WebFetchResponse):
    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': f"Bearer {os.getenv('OLLAMA_API_KEY')}"} if api_key else None)
available_tools = {'web_search': web_search, 'web_fetch': web_fetch}

query = "what is ollama's new engine"
print('Query: ', query)

messages = [{'role': 'user', 'content': query}]
while True:
  response = chat(model='qwen3', messages=messages, tools=[web_search, web_fetch], think=True)
  if response.message.thinking:
    print('Thinking: ')
    print(response.message.thinking + '\n\n')
  if response.message.content:
    print('Content: ')
    print(response.message.content + '\n')

  messages.append(response.message)

  if response.message.tool_calls:
    for tool_call in response.message.tool_calls:
      function_to_call = available_tools.get(tool_call.function.name)
      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:')
        print(args)
        print()

        user_search = args.get('query', '') or args.get('url', '')
        formatted_tool_results = format_tool_results(result, user_search=user_search)

        print(formatted_tool_results[:300])
        print()

        # caps the result at ~2000 tokens
        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})
  else:
    # no more tool calls, we can stop the loop
    break