"""
Matrix integration tools for OllamaTools.
This module provides functionality to connect to Matrix homeservers,
join rooms, send messages, and listen for incoming messages.
It integrates with Ollama to process messages and generate responses.
Dependencies:
- matrix-nio: For Matrix protocol support
- matrix-nio[e2e]: For end-to-end encryption support (optional)
- pyyaml: For configuration handling
"""
import os
import sys
import yaml
import logging
import asyncio
import re
from typing import Dict, List, Optional, Union, Callable, Any, Set
from collections import defaultdict
# Configure logging
logger = logging.getLogger(__name__)
try:
from nio import (
AsyncClient,
MatrixRoom,
RoomMessageText,
LoginResponse,
JoinError
)
except ImportError:
raise ImportError(
"Matrix support requires the 'matrix-nio' package. "
"Please install it with 'pip install matrix-nio' or "
"'pip install matrix-nio[e2e]' for end-to-end encryption support."
)
# Make ollama import optional
try:
from ollama import AsyncClient as OllamaAsync
OLLAMA_AVAILABLE = True
except ImportError:
OLLAMA_AVAILABLE = False
logger.warning(
"The 'ollama' package is not installed. "
"Message processing with Ollama will not be available. "
"Install it with 'pip install ollama'."
)
class MatrixClient:
"""
A client for connecting to Matrix homeservers and interacting with rooms.
This client can connect to a Matrix homeserver, join rooms, send messages,
and listen for incoming messages. It integrates with Ollama to process
messages and generate responses.
"""
# Command patterns
MODEL_COMMAND_PATTERN = re.compile(r'^/model\s+([a-zA-Z0-9:_\-\.]+)$')
HELP_COMMAND_PATTERN = re.compile(r'^/help$')
LIST_MODELS_COMMAND_PATTERN = re.compile(r'^/models$')
def __init__(
self,
homeserver: str,
user_id: str,
password: str,
device_id: str = None,
ollama_host: str = "localhost:11434",
ollama_model: str = "llama3:latest",
ollama_prompt: str = None,
available_models: List[str] = None,
admin_users: List[str] = None,
ssl: bool = True
):
"""
Initialize a new Matrix client.
Args:
homeserver: The URL of the Matrix homeserver
user_id: The Matrix user ID (e.g., @user:example.com)
password: The password for the Matrix account
device_id: A unique identifier for this device (for E2E encryption)
ollama_host: The host address for the Ollama API
ollama_model: The Ollama model to use for generating responses
ollama_prompt: An optional system prompt for the Ollama model
available_models: List of available Ollama models
admin_users: List of Matrix user IDs that can change models for all users
ssl: Whether to use SSL for the Matrix connection
"""
self.homeserver = homeserver
self.user_id = user_id
self.password = password
self.device_id = device_id
self.ollama_host = ollama_host
self.ollama_model = ollama_model
self.ollama_prompt = ollama_prompt
self.ssl = ssl
# Initialize the Matrix client
self.client = AsyncClient(
homeserver=self.homeserver,
device_id=self.device_id,
ssl=self.ssl
)
# Track when the client connected to avoid processing old messages
self.connection_time = None
# Track joined rooms
self.rooms = []
# Message callback
self.message_callback = None
# User context tracking
self.user_contexts = defaultdict(list)
self.max_context_length = 10 # Maximum number of messages to keep in context
# User model preferences
self.user_models = {}
# Available models
self.available_models = available_models or [ollama_model]
# Admin users who can change models
self.admin_users = set(admin_users or [])
async def connect(self) -> bool:
"""
Connect to the Matrix homeserver and log in.
Returns:
bool: True if the connection was successful, False otherwise
"""
try:
# Extract username from user_id (e.g., "@user:example.com" -> "user")
username = self.user_id.split(':')[0].lstrip('@')
# Set the user ID for the client
self.client.user = username
response = await self.client.login(
password=self.password
)
if isinstance(response, LoginResponse):
logger.info(f"Successfully connected to {self.homeserver} as {self.user_id}")
self.connection_time = asyncio.get_event_loop().time()
return True
else:
logger.error(f"Failed to connect to {self.homeserver}: {response}")
return False
except Exception as e:
logger.error(f"Error connecting to Matrix server: {e}")
return False
async def join_room(self, room_id: str) -> bool:
"""
Join a Matrix room.
Args:
room_id: The ID of the room to join (e.g., #room:example.com)
Returns:
bool: True if joining was successful, False otherwise
"""
try:
response = await self.client.join(room_id)
if isinstance(response, JoinError):
logger.error(f"Failed to join room {room_id}: {response}")
return False
logger.info(f"Successfully joined room {room_id}")
self.rooms.append(room_id)
return True
except Exception as e:
logger.error(f"Error joining room {room_id}: {e}")
return False
async def send_message(self, room_id: str, message: str) -> bool:
"""
Send a message to a Matrix room.
Args:
room_id: The ID of the room to send the message to
message: The message to send
Returns:
bool: True if sending was successful, False otherwise
"""
try:
response = await self.client.room_send(
room_id=room_id,
message_type="m.room.message",
content={
"msgtype": "m.text",
"body": message
}
)
logger.debug(f"Sent message to room {room_id}")
return True
except Exception as e:
logger.error(f"Error sending message to room {room_id}: {e}")
return False
async def process_with_ollama(self, message: str, sender: str = None, room_id: str = None) -> str:
"""
Process a message with Ollama and return the response.
Args:
message: The message to process
sender: The Matrix user ID of the sender
room_id: The ID of the room where the message was sent
Returns:
str: The response from Ollama
"""
if not OLLAMA_AVAILABLE:
return "Ollama is not available. Please install the 'ollama' package with 'pip install ollama'."
try:
ollama_client = OllamaAsync(host=self.ollama_host)
# Determine which model to use
model = self.ollama_model
if sender and sender in self.user_models:
model = self.user_models[sender]
messages = []
# Add system prompt if provided
if self.ollama_prompt:
messages.append({
"role": "system",
"content": self.ollama_prompt
})
# Add conversation context if available
if sender and room_id:
context_key = f"{sender}:{room_id}"
for ctx_msg in self.user_contexts[context_key]:
messages.append(ctx_msg)
# Add user message
user_message = {
"role": "user",
"content": message
}
messages.append(user_message)
# Store user message in context
if sender and room_id:
context_key = f"{sender}:{room_id}"
self.user_contexts[context_key].append(user_message)
# Trim context if it gets too long
if len(self.user_contexts[context_key]) > self.max_context_length:
self.user_contexts[context_key] = self.user_contexts[context_key][-self.max_context_length:]
response = await ollama_client.chat(
model=model,
messages=messages
)
response_message = response["message"]
message_content = response_message["content"]
# Store assistant response in context
if sender and room_id:
context_key = f"{sender}:{room_id}"
self.user_contexts[context_key].append({
"role": "assistant",
"content": message_content
})
# Trim context if it gets too long
if len(self.user_contexts[context_key]) > self.max_context_length:
self.user_contexts[context_key] = self.user_contexts[context_key][-self.max_context_length:]
logger.debug(f"Received response from Ollama (model: {model}): {message_content}")
return message_content
except Exception as e:
logger.error(f"Error processing message with Ollama: {e}")
return f"Error processing message: {e}"
async def _message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None:
"""
Internal callback for handling incoming messages.
Args:
room: The room the message was sent in
event: The message event
"""
if not isinstance(event, RoomMessageText):
return
# Get message timestamp
message_timestamp = event.server_timestamp / 1000
# Skip messages from before the client connected
if self.connection_time and message_timestamp < self.connection_time:
logger.debug("Skipping old message")
return
# Skip messages from the bot itself
if event.sender == self.user_id:
logger.debug("Skipping message from self")
return
message = event.body
room_id = room.room_id
sender = event.sender
logger.info(f"Received message in room {room_id} from {sender}: {message}")
# If a custom callback is set, use it
if self.message_callback:
await self.message_callback(room, event, self)
return
# Check if this is a command
if message.startswith('/'):
await self._handle_command(message, room_id, sender)
return
# Extract username from user_id (e.g., "@user:example.com" -> "user")
username = self.user_id.split(':')[0].lstrip('@')
# Check if the message contains the bot's username (case insensitive)
if username.lower() in message.lower():
logger.info(f"Username mention detected in message: {message}")
# Process with Ollama and send response
response = await self.process_with_ollama(message, sender, room_id)
await self.send_message(room_id, response)
return
# For other messages without username mention, don't automatically respond
async def _handle_command(self, message: str, room_id: str, sender: str) -> None:
"""
Handle commands sent to the bot.
Args:
message: The command message
room_id: The ID of the room where the command was sent
sender: The Matrix user ID of the sender
"""
# Check if this is a model change command
model_match = self.MODEL_COMMAND_PATTERN.match(message)
if model_match:
model_name = model_match.group(1)
await self._change_model(model_name, room_id, sender)
return
# Check if this is a help command
if self.HELP_COMMAND_PATTERN.match(message):
await self._send_help(room_id)
return
# Check if this is a list models command
if self.LIST_MODELS_COMMAND_PATTERN.match(message):
await self._list_models(room_id)
return
# Unknown command
await self.send_message(room_id, f"Unknown command: {message}\nType /help for available commands.")
async def _change_model(self, model_name: str, room_id: str, sender: str) -> None:
"""
Change the Ollama model for a user.
Args:
model_name: The name of the model to change to
room_id: The ID of the room where the command was sent
sender: The Matrix user ID of the sender
"""
# Check if the model is available
if model_name not in self.available_models:
await self.send_message(
room_id,
f"Model '{model_name}' is not available. Available models: {', '.join(self.available_models)}"
)
return
# Change the model for the user
self.user_models[sender] = model_name
# Clear the user's context
for key in list(self.user_contexts.keys()):
if key.startswith(f"{sender}:"):
self.user_contexts[key] = []
await self.send_message(room_id, f"Changed model to '{model_name}' for {sender}")
logger.info(f"Changed model to '{model_name}' for user {sender}")
async def _send_help(self, room_id: str) -> None:
"""
Send help information to a room.
Args:
room_id: The ID of the room to send the help to
"""
help_text = """
Available commands:
- /help - Show this help message
- /models - List available models
- /model <model_name> - Change to a different model
You can also mention me by my username in your message and I'll respond.
For example: "Hey {username}, what's the weather like today?"
""".format(username=self.user_id.split(':')[0].lstrip('@'))
await self.send_message(room_id, help_text.strip())
async def _list_models(self, room_id: str) -> None:
"""
List available models to a room.
Args:
room_id: The ID of the room to send the list to
"""
models_text = f"Available models: {', '.join(self.available_models)}"
await self.send_message(room_id, models_text)
async def get_available_models(self) -> List[str]:
"""
Get a list of available Ollama models.
Returns:
List[str]: A list of available model names
"""
if not OLLAMA_AVAILABLE:
return [self.ollama_model]
try:
ollama_client = OllamaAsync(host=self.ollama_host)
response = await ollama_client.list()
models = [model["name"] for model in response["models"]]
return models
except Exception as e:
logger.error(f"Error getting available models: {e}")
return [self.ollama_model]
async def update_available_models(self) -> None:
"""
Update the list of available models from the Ollama server.
"""
models = await self.get_available_models()
self.available_models = models
logger.info(f"Updated available models: {', '.join(models)}")
def add_admin_user(self, user_id: str) -> None:
"""
Add a user to the list of admin users.
Args:
user_id: The Matrix user ID to add as an admin
"""
self.admin_users.add(user_id)
logger.info(f"Added admin user: {user_id}")
def remove_admin_user(self, user_id: str) -> None:
"""
Remove a user from the list of admin users.
Args:
user_id: The Matrix user ID to remove as an admin
"""
if user_id in self.admin_users:
self.admin_users.remove(user_id)
logger.info(f"Removed admin user: {user_id}")
def set_message_callback(self, callback: Callable) -> None:
"""
Set a custom callback for handling incoming messages.
The callback should be an async function that takes three arguments:
- room: The MatrixRoom object
- event: The RoomMessageText event
- client: This MatrixClient instance
Args:
callback: The callback function
"""
self.message_callback = callback
async def listen(self) -> None:
"""
Start listening for messages in joined rooms.
This method will block until the client is stopped.
"""
# Set up the message callback
self.client.add_event_callback(self._message_callback, RoomMessageText)
# Start syncing with the server
await self.client.sync_forever(timeout=30000, full_state=True)
async def disconnect(self) -> None:
"""
Disconnect from the Matrix homeserver.
"""
await self.client.logout()
await self.client.close()
logger.info(f"Disconnected from {self.homeserver}")
class MatrixConfig:
"""
Configuration handler for Matrix integration.
"""
@staticmethod
def load_config(config_path: str) -> Dict:
"""
Load configuration from a YAML file.
Args:
config_path: Path to the configuration file
Returns:
Dict: The configuration as a dictionary
"""
try:
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
logger.info(f"Loaded configuration from {config_path}")
return config
except Exception as e:
logger.error(f"Error loading configuration from {config_path}: {e}")
raise
@staticmethod
def save_config(config: Dict, config_path: str) -> None:
"""
Save configuration to a YAML file.
Args:
config: The configuration as a dictionary
config_path: Path to save the configuration to
"""
try:
# Create directory if it doesn't exist
os.makedirs(os.path.dirname(config_path), exist_ok=True)
with open(config_path, 'w') as f:
yaml.dump(config, f, default_flow_style=False)
logger.info(f"Saved configuration to {config_path}")
except Exception as e:
logger.error(f"Error saving configuration to {config_path}: {e}")
raise
@staticmethod
def create_default_config() -> Dict:
"""
Create a default configuration.
Returns:
Dict: The default configuration
"""
return {
"MATRIX": {
"HOMESERVER": "https://matrix.example.org",
"USER": "@user:example.org",
"PASSWORD": "your-password",
"DEVICE_ID": "ollamatools-matrix",
"ROOMS": [
"#room:example.org"
],
"ADMIN_USERS": [
"@admin:example.org"
]
},
"OLLAMA": {
"HOST": "localhost:11434",
"MODEL": "llama3:latest",
"PROMPT": "You are a helpful assistant.",
"AVAILABLE_MODELS": [
"llama3:latest",
"mistral:latest",
"gemma:latest"
]
}
}
# Async helper functions for the module API
async def _connect_to_matrix(
homeserver: str,
user_id: str,
password: str,
device_id: str = None,
ollama_host: str = "localhost:11434",
ollama_model: str = "llama3:latest",
ollama_prompt: str = None,
available_models: List[str] = None,
admin_users: List[str] = None,
ssl: bool = True
) -> Optional[MatrixClient]:
"""
Connect to a Matrix homeserver.
Args:
homeserver: The URL of the Matrix homeserver
user_id: The Matrix user ID (e.g., @user:example.com)
password: The password for the Matrix account
device_id: A unique identifier for this device (for E2E encryption)
ollama_host: The host address for the Ollama API
ollama_model: The Ollama model to use for generating responses
ollama_prompt: An optional system prompt for the Ollama model
available_models: List of available Ollama models
admin_users: List of Matrix user IDs that can change models for all users
ssl: Whether to use SSL for the Matrix connection
Returns:
Optional[MatrixClient]: A connected Matrix client, or None if connection failed
"""
client = MatrixClient(
homeserver=homeserver,
user_id=user_id,
password=password,
device_id=device_id,
ollama_host=ollama_host,
ollama_model=ollama_model,
ollama_prompt=ollama_prompt,
available_models=available_models,
admin_users=admin_users,
ssl=ssl
)
success = await client.connect()
if success:
return client
else:
return None
async def _join_matrix_room(client: MatrixClient, room_id: str) -> bool:
"""
Join a Matrix room.
Args:
client: A connected Matrix client
room_id: The ID of the room to join (e.g., #room:example.com)
Returns:
bool: True if joining was successful, False otherwise
"""
return await client.join_room(room_id)
async def _send_matrix_message(client: MatrixClient, room_id: str, message: str) -> bool:
"""
Send a message to a Matrix room.
Args:
client: A connected Matrix client
room_id: The ID of the room to send the message to
message: The message to send
Returns:
bool: True if sending was successful, False otherwise
"""
return await client.send_message(room_id, message)
async def _listen_for_matrix_messages(client: MatrixClient) -> None:
"""
Start listening for messages in joined rooms.
Args:
client: A connected Matrix client
"""
await client.listen()
# Synchronous wrapper functions for the module API
def connect_to_matrix(
homeserver: str,
user_id: str,
password: str,
device_id: str = None,
ollama_host: str = "localhost:11434",
ollama_model: str = "llama3:latest",
ollama_prompt: str = None,
available_models: List[str] = None,
admin_users: List[str] = None,
ssl: bool = True
) -> Optional[MatrixClient]:
"""
Connect to a Matrix homeserver.
Args:
homeserver: The URL of the Matrix homeserver
user_id: The Matrix user ID (e.g., @user:example.com)
password: The password for the Matrix account
device_id: A unique identifier for this device (for E2E encryption)
ollama_host: The host address for the Ollama API
ollama_model: The Ollama model to use for generating responses
ollama_prompt: An optional system prompt for the Ollama model
available_models: List of available Ollama models
admin_users: List of Matrix user IDs that can change models for all users
ssl: Whether to use SSL for the Matrix connection
Returns:
Optional[MatrixClient]: A connected Matrix client, or None if connection failed
"""
return asyncio.run(_connect_to_matrix(
homeserver=homeserver,
user_id=user_id,
password=password,
device_id=device_id,
ollama_host=ollama_host,
ollama_model=ollama_model,
ollama_prompt=ollama_prompt,
available_models=available_models,
admin_users=admin_users,
ssl=ssl
))
def join_matrix_room(client: MatrixClient, room_id: str) -> bool:
"""
Join a Matrix room.
Args:
client: A connected Matrix client
room_id: The ID of the room to join (e.g., #room:example.com)
Returns:
bool: True if joining was successful, False otherwise
"""
return asyncio.run(_join_matrix_room(client, room_id))
def send_matrix_message(client: MatrixClient, room_id: str, message: str) -> bool:
"""
Send a message to a Matrix room.
Args:
client: A connected Matrix client
room_id: The ID of the room to send the message to
message: The message to send
Returns:
bool: True if sending was successful, False otherwise
"""
return asyncio.run(_send_matrix_message(client, room_id, message))
def listen_for_matrix_messages(client: MatrixClient) -> None:
"""
Start listening for messages in joined rooms.
This function will block until the client is stopped.
Args:
client: A connected Matrix client
"""
asyncio.run(_listen_for_matrix_messages(client))
def load_matrix_config(config_path: str) -> Dict:
"""
Load Matrix configuration from a YAML file.
Args:
config_path: Path to the configuration file
Returns:
Dict: The configuration as a dictionary
"""
return MatrixConfig.load_config(config_path)
def save_matrix_config(config: Dict, config_path: str) -> None:
"""
Save Matrix configuration to a YAML file.
Args:
config: The configuration as a dictionary
config_path: Path to save the configuration to
"""
MatrixConfig.save_config(config, config_path)
def create_default_matrix_config() -> Dict:
"""
Create a default Matrix configuration.
Returns:
Dict: The default configuration
"""
return MatrixConfig.create_default_config()