Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Tool Registry Feature #234

Closed
wants to merge 3 commits into from
Closed

Add Tool Registry Feature #234

wants to merge 3 commits into from

Conversation

synacktraa
Copy link

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
Copy link

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
Copy link
Author

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
Copy link

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
Copy link

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

@synacktraa
Copy link
Author

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
Copy link
Author

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

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

@synacktraa
Copy link
Author

@TheEpic-dev I removed pydantic dependency :). You can take a look at this commit: 637f9a9

@synacktraa
Copy link
Author

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

@anthonywu
Copy link

@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
Copy link
Author

@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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants