client/types: update web search and fetch API (#584)
Some checks are pending
test / test (push) Waiting to run
test / lint (push) Waiting to run

---------

Co-authored-by: ParthSareen <parth.sareen@ollama.com>
This commit is contained in:
nicole pardal 2025-09-23 13:27:36 -07:00 committed by GitHub
parent d0f71bc8b8
commit 16f344f635
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 95 additions and 107 deletions

View File

@ -9,56 +9,47 @@ from typing import Union
from rich import print from rich import print
from ollama import WebCrawlResponse, WebSearchResponse, chat, web_crawl, web_search from ollama import WebFetchResponse, WebSearchResponse, chat, web_fetch, web_search
def format_tool_results(results: Union[WebSearchResponse, WebCrawlResponse]): def format_tool_results(
results: Union[WebSearchResponse, WebFetchResponse],
user_search: str,
):
output = []
if isinstance(results, WebSearchResponse): if isinstance(results, WebSearchResponse):
if not results.success: output.append(f'Search results for "{user_search}":')
error_msg = ', '.join(results.errors) if results.errors else 'Unknown error' for result in results.results:
return f'Web search failed: {error_msg}' output.append(f'{result.title}' if result.title else f'{result.content}')
output.append(f' URL: {result.url}')
output = [] output.append(f' Content: {result.content}')
for query, search_results in results.results.items(): output.append('')
output.append(f'Search results for "{query}":')
for i, result in enumerate(search_results, 1):
output.append(f'{i}. {result.title}')
output.append(f' URL: {result.url}')
output.append(f' Content: {result.content}')
output.append('')
return '\n'.join(output).rstrip() return '\n'.join(output).rstrip()
elif isinstance(results, WebCrawlResponse): elif isinstance(results, WebFetchResponse):
if not results.success: output.append(f'Fetch results for "{user_search}":')
error_msg = ', '.join(results.errors) if results.errors else 'Unknown error' output.extend(
return f'Web crawl failed: {error_msg}' [
f'Title: {results.title}',
output = [] f'URL: {user_search}' if user_search else '',
for url, crawl_results in results.results.items(): f'Content: {results.content}',
output.append(f'Crawl results for "{url}":') ]
for i, result in enumerate(crawl_results, 1): )
output.append(f'{i}. {result.title}') if results.links:
output.append(f' URL: {result.url}') output.append(f'Links: {", ".join(results.links)}')
output.append(f' Content: {result.content}') output.append('')
if result.links:
output.append(f' Links: {", ".join(result.links)}')
output.append('')
return '\n'.join(output).rstrip() return '\n'.join(output).rstrip()
# Set OLLAMA_API_KEY in the environment variable or use the headers parameter to set the authorization header # client = Client(headers={'Authorization': f"Bearer {os.getenv('OLLAMA_API_KEY')}"} if api_key else None)
# client = Client(headers={'Authorization': 'Bearer <OLLAMA_API_KEY>'}) available_tools = {'web_search': web_search, 'web_fetch': web_fetch}
available_tools = {'web_search': web_search, 'web_crawl': web_crawl} query = "what is ollama's new engine"
query = "ollama's new engine"
print('Query: ', query) print('Query: ', query)
messages = [{'role': 'user', 'content': query}] messages = [{'role': 'user', 'content': query}]
while True: while True:
response = chat(model='qwen3', messages=messages, tools=[web_search, web_crawl], think=True) response = chat(model='qwen3', messages=messages, tools=[web_search, web_fetch], think=True)
if response.message.thinking: if response.message.thinking:
print('Thinking: ') print('Thinking: ')
print(response.message.thinking + '\n\n') print(response.message.thinking + '\n\n')
@ -72,12 +63,20 @@ while True:
for tool_call in response.message.tool_calls: for tool_call in response.message.tool_calls:
function_to_call = available_tools.get(tool_call.function.name) function_to_call = available_tools.get(tool_call.function.name)
if function_to_call: if function_to_call:
result: WebSearchResponse | WebCrawlResponse = function_to_call(**tool_call.function.arguments) args = tool_call.function.arguments
print('Result from tool call name: ', tool_call.function.name, 'with arguments: ', tool_call.function.arguments) result: Union[WebSearchResponse, WebFetchResponse] = function_to_call(**args)
print('Result: ', format_tool_results(result)[:200]) 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 # 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_tool_results[: 2000 * 4], 'tool_name': tool_call.function.name})
else: else:
print(f'Tool {tool_call.function.name} not found') 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}) messages.append({'role': 'tool', 'content': f'Tool {tool_call.function.name} not found', 'tool_name': tool_call.function.name})

View File

@ -15,7 +15,7 @@ from ollama._types import (
ShowResponse, ShowResponse,
StatusResponse, StatusResponse,
Tool, Tool,
WebCrawlResponse, WebFetchResponse,
WebSearchResponse, WebSearchResponse,
) )
@ -37,7 +37,7 @@ __all__ = [
'ShowResponse', 'ShowResponse',
'StatusResponse', 'StatusResponse',
'Tool', 'Tool',
'WebCrawlResponse', 'WebFetchResponse',
'WebSearchResponse', 'WebSearchResponse',
] ]
@ -56,4 +56,4 @@ copy = _client.copy
show = _client.show show = _client.show
ps = _client.ps ps = _client.ps
web_search = _client.web_search web_search = _client.web_search
web_crawl = _client.web_crawl web_fetch = _client.web_fetch

View File

@ -66,8 +66,8 @@ from ollama._types import (
ShowResponse, ShowResponse,
StatusResponse, StatusResponse,
Tool, Tool,
WebCrawlRequest, WebFetchRequest,
WebCrawlResponse, WebFetchResponse,
WebSearchRequest, WebSearchRequest,
WebSearchResponse, WebSearchResponse,
) )
@ -633,13 +633,13 @@ class Client(BaseClient):
'/api/ps', '/api/ps',
) )
def web_search(self, queries: Sequence[str], max_results: int = 3) -> WebSearchResponse: def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse:
""" """
Performs a web search Performs a web search
Args: Args:
queries: The queries to search for query: The query to search for
max_results: The maximum number of results to return. max_results: The maximum number of results to return (default: 3)
Returns: Returns:
WebSearchResponse with the search results WebSearchResponse with the search results
@ -654,32 +654,30 @@ class Client(BaseClient):
'POST', 'POST',
'https://ollama.com/api/web_search', 'https://ollama.com/api/web_search',
json=WebSearchRequest( json=WebSearchRequest(
queries=queries, query=query,
max_results=max_results, max_results=max_results,
).model_dump(exclude_none=True), ).model_dump(exclude_none=True),
) )
def web_crawl(self, urls: Sequence[str]) -> WebCrawlResponse: def web_fetch(self, url: str) -> WebFetchResponse:
""" """
Gets the content of web pages for the provided URLs. Fetches the content of a web page for the provided URL.
Args: Args:
urls: The URLs to crawl url: The URL to fetch
Returns: Returns:
WebCrawlResponse with the crawl results WebFetchResponse with the fetched result
Raises:
ValueError: If OLLAMA_API_KEY environment variable is not set
""" """
if not self._client.headers.get('authorization', '').startswith('Bearer '): if not self._client.headers.get('authorization', '').startswith('Bearer '):
raise ValueError('Authorization header with Bearer token is required for web fetch') raise ValueError('Authorization header with Bearer token is required for web fetch')
return self._request( return self._request(
WebCrawlResponse, WebFetchResponse,
'POST', 'POST',
'https://ollama.com/api/web_crawl', 'https://ollama.com/api/web_fetch',
json=WebCrawlRequest( json=WebFetchRequest(
urls=urls, url=url,
).model_dump(exclude_none=True), ).model_dump(exclude_none=True),
) )
@ -752,13 +750,13 @@ class AsyncClient(BaseClient):
return cls(**(await self._request_raw(*args, **kwargs)).json()) return cls(**(await self._request_raw(*args, **kwargs)).json())
async def websearch(self, queries: Sequence[str], max_results: int = 3) -> WebSearchResponse: async def web_search(self, query: str, max_results: int = 3) -> WebSearchResponse:
""" """
Performs a web search Performs a web search
Args: Args:
queries: The queries to search for query: The query to search for
max_results: The maximum number of results to return. max_results: The maximum number of results to return (default: 3)
Returns: Returns:
WebSearchResponse with the search results WebSearchResponse with the search results
@ -768,27 +766,27 @@ class AsyncClient(BaseClient):
'POST', 'POST',
'https://ollama.com/api/web_search', 'https://ollama.com/api/web_search',
json=WebSearchRequest( json=WebSearchRequest(
queries=queries, query=query,
max_results=max_results, max_results=max_results,
).model_dump(exclude_none=True), ).model_dump(exclude_none=True),
) )
async def webcrawl(self, urls: Sequence[str]) -> WebCrawlResponse: async def web_fetch(self, url: str) -> WebFetchResponse:
""" """
Gets the content of web pages for the provided URLs. Fetches the content of a web page for the provided URL.
Args: Args:
urls: The URLs to crawl url: The URL to fetch
Returns: Returns:
WebCrawlResponse with the crawl results WebFetchResponse with the fetched result
""" """
return await self._request( return await self._request(
WebCrawlResponse, WebFetchResponse,
'POST', 'POST',
'https://ollama.com/api/web_crawl', 'https://ollama.com/api/web_fetch',
json=WebCrawlRequest( json=WebFetchRequest(
urls=urls, url=url,
).model_dump(exclude_none=True), ).model_dump(exclude_none=True),
) )

View File

@ -542,37 +542,28 @@ class ProcessResponse(SubscriptableBaseModel):
class WebSearchRequest(SubscriptableBaseModel): class WebSearchRequest(SubscriptableBaseModel):
queries: Sequence[str] query: str
max_results: Optional[int] = None max_results: Optional[int] = None
class WebSearchResult(SubscriptableBaseModel): class WebSearchResult(SubscriptableBaseModel):
title: str content: Optional[str] = None
url: str title: Optional[str] = None
content: str url: Optional[str] = None
class WebCrawlResult(SubscriptableBaseModel): class WebFetchRequest(SubscriptableBaseModel):
title: str
url: str url: str
content: str
links: Optional[Sequence[str]] = None
class WebSearchResponse(SubscriptableBaseModel): class WebSearchResponse(SubscriptableBaseModel):
results: Mapping[str, Sequence[WebSearchResult]] results: Sequence[WebSearchResult]
success: bool
errors: Optional[Sequence[str]] = None
class WebCrawlRequest(SubscriptableBaseModel): class WebFetchResponse(SubscriptableBaseModel):
urls: Sequence[str] title: Optional[str] = None
content: Optional[str] = None
links: Optional[Sequence[str]] = None
class WebCrawlResponse(SubscriptableBaseModel):
results: Mapping[str, Sequence[WebCrawlResult]]
success: bool
errors: Optional[Sequence[str]] = None
class RequestError(Exception): class RequestError(Exception):

View File

@ -1203,29 +1203,29 @@ def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.Monke
client = Client() client = Client()
with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'): with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'):
client.web_search(['test query']) client.web_search('test query')
def test_client_web_crawl_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch): def test_client_web_fetch_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv('OLLAMA_API_KEY', raising=False) monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
client = Client() client = Client()
with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'): with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'):
client.web_crawl(['https://example.com']) client.web_fetch('https://example.com')
def _mock_request_web_search(self, cls, method, url, json=None, **kwargs): def _mock_request_web_search(self, cls, method, url, json=None, **kwargs):
assert method == 'POST' assert method == 'POST'
assert url == 'https://ollama.com/api/web_search' assert url == 'https://ollama.com/api/web_search'
assert json is not None and 'queries' in json and 'max_results' in json assert json is not None and 'query' in json and 'max_results' in json
return httpxResponse(status_code=200, content='{"results": {}, "success": true}') return httpxResponse(status_code=200, content='{"results": {}, "success": true}')
def _mock_request_web_crawl(self, cls, method, url, json=None, **kwargs): def _mock_request_web_fetch(self, cls, method, url, json=None, **kwargs):
assert method == 'POST' assert method == 'POST'
assert url == 'https://ollama.com/api/web_crawl' assert url == 'https://ollama.com/api/web_fetch'
assert json is not None and 'urls' in json assert json is not None and 'url' in json
return httpxResponse(status_code=200, content='{"results": {}, "success": true}') return httpxResponse(status_code=200, content='{"results": {}, "success": true}')
@ -1234,15 +1234,15 @@ def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(Client, '_request', _mock_request_web_search) monkeypatch.setattr(Client, '_request', _mock_request_web_search)
client = Client() client = Client()
client.web_search(['what is ollama?'], max_results=2) client.web_search('what is ollama?', max_results=2)
def test_client_web_crawl_with_env_api_key(monkeypatch: pytest.MonkeyPatch): def test_client_web_fetch_with_env_api_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv('OLLAMA_API_KEY', 'test-key') monkeypatch.setenv('OLLAMA_API_KEY', 'test-key')
monkeypatch.setattr(Client, '_request', _mock_request_web_crawl) monkeypatch.setattr(Client, '_request', _mock_request_web_fetch)
client = Client() client = Client()
client.web_crawl(['https://example.com']) client.web_fetch('https://example.com')
def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch):
@ -1250,15 +1250,15 @@ def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.Monke
monkeypatch.setattr(Client, '_request', _mock_request_web_search) monkeypatch.setattr(Client, '_request', _mock_request_web_search)
client = Client(headers={'Authorization': 'Bearer custom-token'}) client = Client(headers={'Authorization': 'Bearer custom-token'})
client.web_search(['what is ollama?'], max_results=1) client.web_search('what is ollama?', max_results=1)
def test_client_web_crawl_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch): def test_client_web_fetch_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv('OLLAMA_API_KEY', raising=False) monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
monkeypatch.setattr(Client, '_request', _mock_request_web_crawl) monkeypatch.setattr(Client, '_request', _mock_request_web_fetch)
client = Client(headers={'Authorization': 'Bearer custom-token'}) client = Client(headers={'Authorization': 'Bearer custom-token'})
client.web_crawl(['https://example.com']) client.web_fetch('https://example.com')
def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch): def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch):
@ -1274,4 +1274,4 @@ def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyP
client = Client(headers={'Authorization': 'Bearer explicit-token'}) client = Client(headers={'Authorization': 'Bearer explicit-token'})
assert client._client.headers['authorization'] == 'Bearer explicit-token' assert client._client.headers['authorization'] == 'Bearer explicit-token'
client.web_search(['override check']) client.web_search('override check')