ollama / ollama-python

Ollama Python library
https://ollama.com
MIT License
2.69k stars 221 forks source link

Add type overloads for chat() method in _client.py #132

Closed simonottenhauskenbun closed 5 days ago

simonottenhauskenbun commented 2 months ago

This PR adds two type overloads in _client.py to reduce type warnings. The PR should only affect typing and no logic.

Before PR:

import ollama
response = ollama.chat(model='llama2', messages=[
  {
    'role': 'user',
    'content': 'Why is the sky blue?',
  },
])
# before PR: response is now Union[Mapping[str, Any], Iterator[Mapping[str, Any]]]
print(response['message']['content'])

New behaviour:

response = ollama.chat(model='llama2', stream=False, ...)
# response is of type Mapping[str, Any]

response = ollama.chat(model='llama2', stream=True, ...)
# response is of type Iterator[Mapping[str, Any]]

response = ollama.chat(model='llama2', ...)
# response is also of type Mapping[str, Any], since stream=False is the default.
mxyng commented 3 weeks ago

This is a good change. I was hesitant to add it originally since there are so many methods that should be overloaded.

@royjhan is going to carry this change if it's ok with you @simonottenhauskenbun and implement overloads for the other streaming methods

simonottenhauskenbun commented 5 days ago

This is a good change. I was hesitant to add it originally since there are so many methods that should be overloaded.

@royjhan is going to carry this change if it's ok with you @simonottenhauskenbun and implement overloads for the other streaming methods

Sorry for the late replay, I was on vaccation :) Thanks for implementing all the overloads. This PR is now longer needed.