hueso/ollama_slixmpp_omemo_bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

satoshi

Browse Source
This commit is contained in:
hueso 2025-03-06 20:19:14 -03:00
parent 1576952d8f
commit 193aef1d0b
2 changed files with 38 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
requirements.txt
View File

@ -8,6 +8,7 @@ pytest
pytest-cov pytest-cov
pytest-mypy-plugins pytest-mypy-plugins
result result
slixmpp slixmpp==1.8.6
slixmpp_omemo slixmpp_omemo==0.9.1
twine twine
ollama

58
src/ollama_slixmpp_omemo_bot/ollama_bot.py
View File

@ -2,7 +2,7 @@ from enum import Enum
from typing import Dict, Optional from typing import Dict, Optional
import re import re
import ollama from ollama import Client
from slixmpp import ClientXMPP, JID from slixmpp import ClientXMPP, JID
from slixmpp.exceptions import IqTimeout, IqError from slixmpp.exceptions import IqTimeout, IqError
from slixmpp.stanza import Message from slixmpp.stanza import Message
@ -26,8 +26,8 @@ class LEVELS(Enum):
class LLMS(Enum): class LLMS(Enum):
LLAMA3 = "llama3" LLAMA3 = "based"
MISTRAL = "mistral" MISTRAL = "shreyshah/satoshi-7b-q4_k_m"
class OllamaBot(ClientXMPP): class OllamaBot(ClientXMPP):
@ -37,7 +37,7 @@ class OllamaBot(ClientXMPP):
def __init__(self, jid: JidStr, password: str): def __init__(self, jid: JidStr, password: str):
ClientXMPP.__init__(self, jid, password) ClientXMPP.__init__(self, jid, password)
self.model: LLMS = LLMS.LLAMA3 self.model: LLMS = LLMS.MISTRAL
self.prefix_re: re.Pattern = re.compile(r"^%s" % self.cmd_prefix) self.prefix_re: re.Pattern = re.compile(r"^%s" % self.cmd_prefix)
self.cmd_re: re.Pattern = re.compile( self.cmd_re: re.Pattern = re.compile(
r"^%s(?P<command>\w+)(?:\s+(?P<args>.*))?" % self.cmd_prefix r"^%s(?P<command>\w+)(?:\s+(?P<args>.*))?" % self.cmd_prefix
@ -82,35 +82,35 @@ class OllamaBot(ClientXMPP):
async def cmd_help(self, mto: JID, mtype: Optional[MessageTypes]) -> None: async def cmd_help(self, mto: JID, mtype: Optional[MessageTypes]) -> None:
body = ( body = (
"Hello, I am the ollama_slixmpp_omemo_bot!\n\n" "Hello, I am Satoshi Nakamoto!\n\n"
"The following commands are available:\n\n" # "The following commands are available:\n\n"
f"{self.cmd_prefix}verbose - Send message or reply with log messages.\n\n" # f"{self.cmd_prefix}verbose - Send message or reply with log messages.\n\n"
f"{self.cmd_prefix}error -Send message or reply only on error.\n\n" # f"{self.cmd_prefix}error -Send message or reply only on error.\n\n"
f"{self.cmd_prefix}llama3 - Enable the llama3 model.\n\n" # f"{self.cmd_prefix}llama3 - Enable the llama3 model.\n\n"
f"{self.cmd_prefix}mistral - Enable the mistral model.\n\n" # f"{self.cmd_prefix}mistral - Enable the mistral model.\n\n"
f"Typing anything else will be sent to {self.model.value}!\n\n" # f"Typing anything else will be sent to {self.model.value}!\n\n"
) )
return await self.encrypted_reply(mto, mtype, body) return await self.plain_reply(mto, mtype, body)
async def cmd_set_llama3(self, mto: JID, mtype: Optional[MessageTypes]) -> None: async def cmd_set_llama3(self, mto: JID, mtype: Optional[MessageTypes]) -> None:
self.model = LLMS.LLAMA3 self.model = LLMS.LLAMA3
body: str = f"""Model set to {LLMS.LLAMA3.value}""" body: str = f"""Model set to {LLMS.LLAMA3.value}"""
return await self.encrypted_reply(mto, mtype, body) return await self.plain_reply(mto, mtype, body)
async def cmd_set_mistral(self, mto: JID, mtype: Optional[MessageTypes]) -> None: async def cmd_set_mistral(self, mto: JID, mtype: Optional[MessageTypes]) -> None:
self.model = LLMS.MISTRAL self.model = LLMS.MISTRAL
body: str = f"""Model set to {LLMS.MISTRAL.value}""" body: str = f"""Model set to {LLMS.MISTRAL.value}"""
return await self.encrypted_reply(mto, mtype, body) return await self.plain_reply(mto, mtype, body)
async def cmd_verbose(self, mto: JID, mtype: Optional[MessageTypes]) -> None: async def cmd_verbose(self, mto: JID, mtype: Optional[MessageTypes]) -> None:
self.debug_level = LEVELS.DEBUG self.debug_level = LEVELS.DEBUG
body: str = """Debug level set to 'verbose'.""" body: str = """Debug level set to 'verbose'."""
return await self.encrypted_reply(mto, mtype, body) return await self.plain_reply(mto, mtype, body)
async def cmd_error(self, mto: JID, mtype: Optional[MessageTypes]) -> None: async def cmd_error(self, mto: JID, mtype: Optional[MessageTypes]) -> None:
self.debug_level = LEVELS.ERROR self.debug_level = LEVELS.ERROR
body: str = """Debug level set to 'error'.""" body: str = """Debug level set to 'error'."""
return await self.encrypted_reply(mto, mtype, body) return await self.plain_reply(mto, mtype, body)
async def message_handler( async def message_handler(
self, msg: Message, allow_untrusted: bool = False self, msg: Message, allow_untrusted: bool = False
@ -120,11 +120,16 @@ class OllamaBot(ClientXMPP):
mtype: Optional[MessageTypes] = msg["type"] mtype: Optional[MessageTypes] = msg["type"]
if mtype not in ("chat", "normal"): if mtype not in ("chat", "normal"):
return None return None
if not msg["body"]:
return None
log = open("log.txt", "a", 1)
if not self["xep_0384"].is_encrypted(msg): if not self["xep_0384"].is_encrypted(msg):
if self.debug_level == LEVELS.DEBUG: log.write(f"{mfrom}: {msg['body']}\n")
await self.plain_reply( ollama_server_response: Optional[str] = self.message_to_ollama_server(
mto, mtype, f"Echo unencrypted message: {msg['body']}" msg["body"]
) )
await self.plain_reply(mto, mtype, f"{ollama_server_response or ''}")
return None return None
try: try:
encrypted = msg["omemo_encrypted"] encrypted = msg["omemo_encrypted"]
@ -133,13 +138,14 @@ class OllamaBot(ClientXMPP):
) )
if body is not None: if body is not None:
decoded: str = body.decode("utf8") decoded: str = body.decode("utf8")
log.write(f"{mfrom}: {decoded}\n")
if self.is_command(decoded): if self.is_command(decoded):
await self.handle_command(mto, mtype, decoded) await self.handle_command(mto, mtype, decoded)
elif self.debug_level == LEVELS.DEBUG: elif self.debug_level == LEVELS.DEBUG:
ollama_server_response: Optional[str] = ( ollama_server_response: Optional[str] = (
self.message_to_ollama_server(decoded) self.message_to_ollama_server(decoded)
) )
await self.encrypted_reply( await self.plain_reply(
mto, mtype, f"{ollama_server_response or ''}" mto, mtype, f"{ollama_server_response or ''}"
) )
except MissingOwnKey: except MissingOwnKey:
@ -149,7 +155,7 @@ class OllamaBot(ClientXMPP):
"Error: Message not encrypted for me.", "Error: Message not encrypted for me.",
) )
except NoAvailableSession: except NoAvailableSession:
await self.encrypted_reply( await self.plain_reply(
mto, mto,
mtype, mtype,
"Error: Message uses an encrypted session I don't know about.", "Error: Message uses an encrypted session I don't know about.",
@ -178,6 +184,8 @@ class OllamaBot(ClientXMPP):
return None return None
async def plain_reply(self, mto: JID, mtype: Optional[MessageTypes], body): async def plain_reply(self, mto: JID, mtype: Optional[MessageTypes], body):
log = open("log.txt", "a", 1)
log.write(f"SATOSHI: {body}\n")
msg = self.make_message(mto=mto, mtype=mtype) msg = self.make_message(mto=mto, mtype=mtype)
msg["body"] = body msg["body"] = body
return msg.send() return msg.send()
@ -227,7 +235,11 @@ class OllamaBot(ClientXMPP):
def message_to_ollama_server(self, msg: Optional[str]) -> Optional[str]: def message_to_ollama_server(self, msg: Optional[str]) -> Optional[str]:
if msg is not None: if msg is not None:
response = ollama.chat( client = Client(
host="https://solarpunk.land/ollama/",
headers={"Authorization": "Bearer xxx"},
)
response = client.chat(
model=self.model.value, model=self.model.value,
messages=[{"role": "user", "content": f"{msg}"}], messages=[{"role": "user", "content": f"{msg}"}],
) )