# 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.
*   [[programming/python/modules/argparse-module|argparse]]: To handle command-line arguments.

## Prerequisites

1.  **Install Ollama**: Download and install Ollama from ollama.com.
2.  **Pull a Model**: Open your terminal and run `ollama pull llama3` (or any other model like `mistral`, `gemma`).
3.  **Start Ollama**: Ensure the Ollama app is running in the background.

## Installation

```bash
pip install ollama
```

## The Code

Save this as `ollama_tool.py`.

```python
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

1.  **Interactive Chat**:
    ```bash
    python ollama_tool.py
    ```

2.  **Single Prompt**:
    ```bash
    python ollama_tool.py "Why is the sky blue?"
    ```

3.  **Use a Different Model**:
    ```bash
    python ollama_tool.py --model mistral
    ```

[[programming/python/python]]