mirror of
https://github.com/ollama/ollama-python.git
synced 2026-01-14 06:07:17 +08:00
client/types: update web search and fetch API (#584)
--------- Co-authored-by: ParthSareen <parth.sareen@ollama.com>
This commit is contained in:
parent
d0f71bc8b8
commit
16f344f635
@ -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})
|
||||||
|
|||||||
@ -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
|
||||||
|
|||||||
@ -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),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@ -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):
|
||||||
|
|||||||
@ -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')
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user