mirror of
https://github.com/kharonsec/ollama-python
synced 2026-05-09 16:42:13 +02:00
88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
import json
|
|
import ollama
|
|
import asyncio
|
|
|
|
|
|
# Simulates an API call to get flight times
|
|
# In a real application, this would fetch data from a live database or API
|
|
def get_flight_times(departure: str, arrival: str) -> str:
|
|
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'}))
|
|
|
|
|
|
async def run(model: str):
|
|
client = ollama.AsyncClient()
|
|
# Initialize conversation with a user query
|
|
messages = [{'role': 'user', 'content': 'What is the flight time from New York (NYC) to Los Angeles (LAX)?'}]
|
|
|
|
# First API call: Send the query and function description to the model
|
|
response = await client.chat(
|
|
model=model,
|
|
messages=messages,
|
|
tools=[
|
|
{
|
|
'type': 'function',
|
|
'function': {
|
|
'name': 'get_flight_times',
|
|
'description': 'Get the flight times between two cities',
|
|
'parameters': {
|
|
'type': 'object',
|
|
'properties': {
|
|
'departure': {
|
|
'type': 'string',
|
|
'description': 'The departure city (airport code)',
|
|
},
|
|
'arrival': {
|
|
'type': 'string',
|
|
'description': 'The arrival city (airport code)',
|
|
},
|
|
},
|
|
'required': ['departure', 'arrival'],
|
|
},
|
|
},
|
|
},
|
|
],
|
|
)
|
|
|
|
# Add the model's response to the conversation history
|
|
messages.append(response['message'])
|
|
|
|
# Check if the model decided to use the provided function
|
|
if not response['message'].get('tool_calls'):
|
|
print("The model didn't use the function. Its response was:")
|
|
print(response['message']['content'])
|
|
return
|
|
|
|
# Process function calls made by the model
|
|
if response['message'].get('tool_calls'):
|
|
available_functions = {
|
|
'get_flight_times': get_flight_times,
|
|
}
|
|
for tool in response['message']['tool_calls']:
|
|
function_to_call = available_functions[tool['function']['name']]
|
|
function_response = function_to_call(tool['function']['arguments']['departure'], tool['function']['arguments']['arrival'])
|
|
# Add function response to the conversation
|
|
messages.append(
|
|
{
|
|
'role': 'tool',
|
|
'content': function_response,
|
|
}
|
|
)
|
|
|
|
# Second API call: Get final response from the model
|
|
final_response = await client.chat(model=model, messages=messages)
|
|
print(final_response['message']['content'])
|
|
|
|
|
|
# Run the async function
|
|
asyncio.run(run('mistral'))
|