How to Use AutoGen with Local Models (GPT4All, Ooba, Ollama)

# Install AutoGen
pip install pyautogen

# Depending on your backend:
pip install gpt4all              # GPT4All
pip install llama-cpp-python    # Oobabooga
# Ollama: install from https://ollama.ai

# Install Ollama
curl -fsSL https://ollama.ai/install.sh | sh

# Pull models
ollama pull llama3:8b         # Fast, good for simple tasks
ollama pull llama3:70b        # Best quality, needs strong hardware
ollama pull mistral:7b        # Great for coding tasks
ollama pull codellama:13b     # Specialized for code generation
ollama pull mixtral:8x7b      # Strong all-rounder

# Start Ollama server (runs automatically after install on most systems)
ollama serve

# autogen_ollama_config.py
import autogen

# Ollama config — note the base_url pointing to local server
ollama_llm_config = {
    "config_list": [
        {
            "model": "llama3:8b",
            "api_key": "ollama",          # Placeholder — Ollama doesn't require a real key
            "base_url": "http://localhost:11434/v1",  # OpenAI-compatible endpoint
            "api_type": "openai"
        }
    ],
    "temperature": 0,
    "timeout": 300  # Local models can be slow — set a generous timeout
}

# Create agents exactly as you would with OpenAI
assistant = autogen.AssistantAgent(
    name="LocalAssistant",
    llm_config=ollama_llm_config,
    system_message="You are a helpful coding assistant. Write clean Python code."
)

user = autogen.UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=5,
    code_execution_config={"work_dir": "local_output", "use_docker": False},
    is_termination_msg=lambda msg: "DONE" in msg.get("content", "").upper()
)

user.initiate_chat(
    assistant,
    message="Write a Python script that reads a CSV file and calculates summary statistics."
)

# Multi-model fallback config
ollama_config_with_fallback = {
    "config_list": [
        {
            "model": "llama3:70b",
            "api_key": "ollama",
            "base_url": "http://localhost:11434/v1",
            "api_type": "openai"
        },
        {
            "model": "llama3:8b",  # Fallback if 70B is too slow
            "api_key": "ollama",
            "base_url": "http://localhost:11434/v1",
            "api_type": "openai"
        }
    ],
    "temperature": 0,
    "timeout": 600
}

# Clone and setup Oobabooga
git clone https://github.com/oobabooga/text-generation-webui
cd text-generation-webui
pip install -r requirements.txt

# Start with OpenAI API extension enabled
python server.py --api --extensions openai --listen
# Server starts at localhost:5000/v1 by default

# autogen_ooba_config.py
import autogen

ooba_llm_config = {
    "config_list": [
        {
            "model": "mistral-7b-instruct-v0.2",  # Must match loaded model name in Ooba
            "api_key": "none",
            "base_url": "http://localhost:5000/v1",
            "api_type": "openai"
        }
    ],
    "temperature": 0.1,
    "timeout": 600,
    "max_tokens": 2048
}

assistant = autogen.AssistantAgent(
    name="OobaAssistant",
    llm_config=ooba_llm_config,
    system_message="""You are a data analysis expert.
    Analyze data, write Python code to process it, and explain your findings."""
)

user = autogen.UserProxyAgent(
    name="Analyst",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=8,
    code_execution_config={"work_dir": "analysis_output", "use_docker": False}
)

user.initiate_chat(
    assistant,
    message="Analyze a dataset: I have a CSV with columns [date, sales, returns, region]. What metrics should I calculate and how?"
)

# autogen_gpt4all_config.py
import autogen
from gpt4all import GPT4All

# GPT4All doesn't have a native OpenAI-compatible server,
# so we wrap it in a custom LLM class

class GPT4AllLLM:
    def __init__(self, model_name: str = "Meta-Llama-3-8B-Instruct.Q4_0.gguf"):
        self.model = GPT4All(model_name)
    
    def create(self, messages: list, **kwargs) -> dict:
        """OpenAI-compatible chat completion interface."""
        # Convert messages to a single prompt
        prompt = self._messages_to_prompt(messages)
        
        with self.model.chat_session():
            response = self.model.generate(
                prompt,
                max_tokens=kwargs.get("max_tokens", 512),
                temp=kwargs.get("temperature", 0.1)
            )
        
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": response
                    },
                    "finish_reason": "stop"
                }
            ],
            "model": "gpt4all",
            "usage": {"total_tokens": len(response.split())}
        }
    
    def _messages_to_prompt(self, messages: list) -> str:
        """Convert chat message format to a single string prompt."""
        prompt_parts = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                prompt_parts.append(f"System: {content}")
            elif role == "user":
                prompt_parts.append(f"Human: {content}")
            elif role == "assistant":
                prompt_parts.append(f"Assistant: {content}")
        prompt_parts.append("Assistant:")
        return "\n".join(prompt_parts)


# For a simpler approach, run GPT4All's built-in API server
# gpt4all --api --port 4891
# Then use it exactly like Ollama:

gpt4all_config = {
    "config_list": [
        {
            "model": "gpt4all-falcon-newbpe-q4_0",
            "api_key": "none",
            "base_url": "http://localhost:4891/v1",
            "api_type": "openai"
        }
    ],
    "temperature": 0,
    "timeout": 300
}

# local_llm_factory.py
import autogen
from typing import Literal

def get_llm_config(
    backend: Literal["ollama", "oobabooga", "gpt4all", "openai"] = "ollama",
    model: str = None,
    temperature: float = 0
) -> dict:
    """Factory function for AutoGen LLM configs across backends."""
    
    configs = {
        "ollama": {
            "model": model or "llama3:8b",
            "api_key": "ollama",
            "base_url": "http://localhost:11434/v1",
            "api_type": "openai"
        },
        "oobabooga": {
            "model": model or "mistral-7b-instruct",
            "api_key": "none",
            "base_url": "http://localhost:5000/v1",
            "api_type": "openai"
        },
        "gpt4all": {
            "model": model or "gpt4all-falcon",
            "api_key": "none",
            "base_url": "http://localhost:4891/v1",
            "api_type": "openai"
        },
        "openai": {
            "model": model or "gpt-4o",
            "api_key": "YOUR_OPENAI_API_KEY",
            "api_type": "openai"
        }
    }
    
    return {
        "config_list": [configs[backend]],
        "temperature": temperature,
        "timeout": 120 if backend == "openai" else 600
    }


# Usage — swap backends without touching agent code
import os

backend = os.getenv("LLM_BACKEND", "ollama")
llm_config = get_llm_config(backend=backend, model=os.getenv("LLM_MODEL"))

assistant = autogen.AssistantAgent(
    name="FlexibleAssistant",
    llm_config=llm_config
)

# hybrid_local_setup.py
import autogen
from local_llm_factory import get_llm_config

# Strong model for orchestration (70B if hardware allows)
orchestrator_config = get_llm_config("ollama", model="llama3:70b")

# Lighter model for specialized tasks
worker_config = get_llm_config("ollama", model="codellama:13b")

planner = autogen.AssistantAgent(
    name="Planner",
    llm_config=orchestrator_config,
    system_message="Break down complex tasks and assign them to specialized agents."
)

coder = autogen.AssistantAgent(
    name="Coder",
    llm_config=worker_config,
    system_message="Write Python code based on specifications provided."
)

reviewer = autogen.AssistantAgent(
    name="Reviewer",
    llm_config=worker_config,
    system_message="Review code for correctness and style. Flag issues clearly."
)

group_chat = autogen.GroupChat(
    agents=[planner, coder, reviewer],
    messages=[],
    max_round=10
)

manager = autogen.GroupChatManager(
    groupchat=group_chat,
    llm_config=orchestrator_config
)

user = autogen.UserProxyAgent(
    name="User",
    human_input_mode="NEVER",
    code_execution_config={"work_dir": "local_output", "use_docker": False}
)

user.initiate_chat(manager, message="Build a simple REST API with FastAPI that has CRUD endpoints for a todo list.")

# Prompting strategies for weaker local models

# 1. More explicit output format instructions
system_message = """You are a coding assistant.
ALWAYS format your responses as:
ANALYSIS: [your analysis here]
CODE:
# [python block start]
[code here]
# [block end]
EXPLANATION: [explanation here]
END"""

# 2. Shorter, more focused tasks
# Instead of: "Build a complete user authentication system"
# Use: "Write a password hashing function using bcrypt"

# 3. Chain smaller tasks instead of one large prompt
tasks = [
    "Define the data model for a user account",
    "Write the password validation function",
    "Write the session token generation function",
    "Write the login endpoint combining the above"
]

# 4. Increase temperature slightly for creative tasks, keep at 0 for code
creative_config = {**llm_config, "temperature": 0.3}
code_config = {**llm_config, "temperature": 0.0}

# data_classifier.py
import re

SENSITIVE_PATTERNS = [
    r'\b\d{3}-\d{2}-\d{4}\b',          # SSN
    r'\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b',  # Credit card
    r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',  # Email
    r'\b(?:\+1[-\s]?)?\(?\d{3}\)?[-\s]\d{3}[-\s]\d{4}\b'  # Phone
]

def contains_sensitive_data(text: str) -> bool:
    for pattern in SENSITIVE_PATTERNS:
        if re.search(pattern, text):
            return True
    return False

def safe_llm_call(text: str, use_local: bool = False):
    """Route to local model if sensitive data detected."""
    if contains_sensitive_data(text) and not use_local:
        print("Warning: Sensitive data detected. Routing to local model.")
        return get_llm_config("ollama")
    return get_llm_config("openai")