Add Tool Registry Feature

[synacktraa]

Description

This PR introduces a new ToolRegistry class to the Ollama library. The Tool Registry is designed to manage various tools, including functions (both synchronous and asynchronous), pydantic.BaseModel, typing.TypedDict, and typing.NamedTuple. It provides functionality to generate schemas for these tools, which are essential for LLM tool-calling, and allows for the invocation of tools using their metadata.

Features

• Create and manage a registry of tools
• Support for multiple tool types: functions, pydantic models, TypedDict, and NamedTuple
• Generate schemas for registered tools
• Invoke tools using their metadata (name & raw arguments) -- Handles almost all edge cases
• Option to override existing tools with the same name

Changes I've made:

• Added a new directory tool_calling which stores all functionality of tool registry.
• Added unit tests to cover all major functionality of the ToolRegistry class. (All tests passed)
• Updated _types.py
• Updated tools example in examples directory

How to use this functionality?

Creating a tool registry

import ollama

registry = ollama.ToolRegistry()

Adding tools to the registry instance

import json
import asyncio
from typing import Literal, TypedDict

registry = ollama.ToolRegistry()

@registry.register
def get_flight_times(departure: str, arrival: str) -> str:
  """
  Get flight times.
  :param departure: Departure location code
  :param arrival: Arrival location code
  """
  flights = {
    'NYC-LAX': {'departure': '08:00 AM', 'arrival': '11:30 AM', 'duration': '5h 30m'},
    'LAX-NYC': {'departure': '02:00 PM', 'arrival': '10:30 PM', 'duration': '5h 30m'},
    'LHR-JFK': {'departure': '10:00 AM', 'arrival': '01:00 PM', 'duration': '8h 00m'},
    'JFK-LHR': {'departure': '09:00 PM', 'arrival': '09:00 AM', 'duration': '7h 00m'},
    'CDG-DXB': {'departure': '11:00 AM', 'arrival': '08:00 PM', 'duration': '6h 00m'},
    'DXB-CDG': {'departure': '03:00 AM', 'arrival': '07:30 AM', 'duration': '7h 30m'},
  }

  key = f'{departure}-{arrival}'.upper()
  return json.dumps(flights.get(key, {'error': 'Flight not found'}))

@registry.register
class User(TypedDict): # It can also be `pydantic.BaseModel`, or `typing.NamedTuple`.
    """
    User Information
    :param name: Name of the user
    :param role: Role assigned to the user 
    """
    name: str
    role: Literal['admin', 'developer', 'tester']

"""
# Tools can also be registered using:
registry.register(get_flight_times)
registry.register(User)

# OR

registry.register_multiple(get_flight_times, User)
"""

Get tool schema list

tools = registry.tools
print(json.dumps(tools, indent=3))

Invoking the tool

res = ollama.chat(
    model='llama3-groq-tool-use:latest',
    tools=tools,
    messages=[{
        'role': 'user',
        'content': "What is the flight time from New York (NYC) to Los Angeles (LAX)?"
    }]
)
tool_call = res['message']['tool_calls'][0]['function']
print(f"{tool_call=}")
tool_output = registry.invoke(**tool_call)
print(f"{tool_output=}")

tool_call={'name': 'get_flight_times', 'arguments': {'arrival': 'LAX', 'departure': 'NYC'}}
tool_output='{"departure": "08:00 AM", "arrival": "11:30 AM", "duration": "5h 30m"}'

[TheEpic-dev]

I would recommend packaging this into an external third-party library, rather than including this in the ollama-python bindings.

There are already alternative implementations such as LangChain's tool calling feature, and as not everyone requires tools, adding extra dependencies such as pydantic feels like unnecessary bloat to me.

[synacktraa]

Thank you for your feedback on my PR.

Would it be helpful if I refactored the PR to make this feature optional, perhaps as a separate module within ollama-python that users can opt into? This way, we could avoid adding dependencies to the core package while still providing the functionality for those who need it.

[TheEpic-dev]

Thank you for your feedback on my PR.

Would it be helpful if I refactored the PR to make this feature optional, perhaps as a separate module within ollama-python that users can opt into? This way, we could avoid adding dependencies to the core package while still providing the functionality for those who need it.

I can't speak for the team, but personally, I think the best approach is a third-party library.

You can add it to the Libraries section of the ollama README so anyone can find and use it, it gives them more choice, and decreases the maintenance burden for Ollama maintainers.

But as I said, this is my opinion. Maybe they have a different vision.

[Josephrp]

i think it's fantastic :-) thanks for this , overdue even :-)

[synacktraa]

Thank you for your feedback on my PR. Would it be helpful if I refactored the PR to make this feature optional, perhaps as a separate module within ollama-python that users can opt into? This way, we could avoid adding dependencies to the core package while still providing the functionality for those who need it.

I can't speak for the team, but personally, I think the best approach is a third-party library.

You can add it to the Libraries section of the ollama README so anyone can find and use it, it gives them more choice, and decreases the maintenance burden for Ollama maintainers.

But as I said, this is my opinion. Maybe they have a different vision.

No problem, I understand what you're trying to say. If you guys decide to close this PR, I'll go and release it as a third party library. Thanks for your time!

[synacktraa]

i think it's fantastic :-) thanks for this , overdue even :-)

I am glad you found this feature useful. You're most welcome!

[synacktraa]

@TheEpic-dev I removed pydantic dependency :). You can take a look at this commit: https://github.com/ollama/ollama-python/pull/234/commits/637f9a9b702f279e4ea3a62363ec269f2714aa6c

[synacktraa]

Hi, I have released this feature as a separate library: https://github.com/synacktraa/tool-parse. I will now close this PR👍

[anthonywu]

@synacktraa I believe you would be interested in reviewing #238 or collaborating on your new tool-parse library. I think we're attacking the same problem from slightly different places but can converge.

[synacktraa]

@synacktraa I believe you would be interested in reviewing #238 or collaborating on your new tool-parse library. I think we're attacking the same problem from slightly different places but can converge.

Nice implementation without introducing any dependencies, although I ended up making pydantic optional there is still docstring_parser. I was already thinking about using annotated type for adding parameter description like autogen and your implementation does. With the help of annotated even docstring_parser can be made optional.