Local LLM Wrapper with Ollama
This guide demonstrates how to interact with open-source Large Language Models (LLMs) running locally on your machine using Ollama. This allows you to use powerful models like Llama 3, Mistral, and Gemma without sending data to the cloud.
Modules Used:
ollama: The official Python library for the Ollama API.- argparse: To handle command-line arguments.
Prerequisites
- Install Ollama: Download and install Ollama from ollama.com.
- Pull a Model: Open your terminal and run
ollama pull llama3(or any other model likemistral,gemma). - Start Ollama: Ensure the Ollama app is running in the background.
Installation
pip install ollamaThe Code
Save this as ollama_tool.py.
import ollama
import argparse
import sys
def single_query(model, prompt):
"""Send a single prompt and print the response."""
print(f"Generating response from {model}...")
try:
response = ollama.generate(model=model, prompt=prompt)
print("\n" + response['response'])
except ollama.ResponseError as e:
print(f"Error: {e.error}")
except Exception as e:
print(f"An error occurred: {e}")
def chat_mode(model):
"""Interactive chat session with history."""
print(f"--- Chatting with {model} (Type 'quit' to exit) ---")
messages = []
while True:
user_input = input("\nYou: ")
if user_input.lower() in ['quit', 'exit']:
break
if not user_input.strip():
continue
# Add user message to history
messages.append({'role': 'user', 'content': user_input})
print("Ollama: ", end="", flush=True)
try:
# Stream the response
stream = ollama.chat(
model=model,
messages=messages,
stream=True,
)
collected_response = ""
for chunk in stream:
content = chunk['message']['content']
print(content, end="", flush=True)
collected_response += content
print() # Newline
# Add assistant response to history
messages.append({'role': 'assistant', 'content': collected_response})
except Exception as e:
print(f"\nError: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Ollama Local LLM Wrapper")
parser.add_argument("prompt", nargs="?", help="The prompt to send (leave empty for chat mode)")
parser.add_argument("--model", default="llama3", help="Model to use (default: llama3)")
args = parser.parse_args()
# Check if model exists (optional, but good practice)
try:
ollama.show(args.model)
except:
print(f"Warning: Model '{args.model}' might not be pulled yet. Run 'ollama pull {args.model}' if this fails.")
if args.prompt:
single_query(args.model, args.prompt)
else:
chat_mode(args.model)Usage
-
Interactive Chat:
python ollama_tool.py -
Single Prompt:
python ollama_tool.py "Why is the sky blue?" -
Use a Different Model:
python ollama_tool.py --model mistral