This commit is contained in:
2024-08-18 18:01:05 -05:00
+29 -81
View File
@@ -6,7 +6,6 @@ import logging
import asyncio
from telegram import error as TelegramErrors, Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
from openai import OpenAI
from dotenv import load_dotenv
from tools.base_tool import BaseTool
from anthropic import Anthropic
@@ -14,24 +13,11 @@ from anthropic import Anthropic
# Load environment variables
load_dotenv()
openai_client = OpenAI()
anthropic_client = Anthropic(
api_key=os.environ.get("ANTHROPIC_API_KEY"),
default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
)
GPT_4O = "gpt-4o"
GPT_4O_MINI = "gpt-4o-mini"
model_max_tokens = {
GPT_4O: 4096,
GPT_4O_MINI: 16384
}
use_smart_model = False
use_anthropic = True
# Set up logging to console and file
logging.basicConfig(level=logging.WARNING, handlers=[
logging.StreamHandler(),
@@ -67,9 +53,18 @@ functions = []
for tool in tools:
functions.extend(tool.get_functions())
def get_keyboard():
keyboard = [
['/status', '/reset']
]
return ReplyKeyboardMarkup(keyboard, resize_keyboard=True)
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
logging.info("Bot started")
await update.message.reply_text("Hello! I'm your AI assistant. How can I help you today? You can send me images and then ask questions about them.")
await update.message.reply_text(
"Hello! I'm your AI assistant. How can I help you today? You can send me images and then ask questions about them.",
reply_markup=get_keyboard()
)
async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
user_id = update.effective_user.id
@@ -79,7 +74,7 @@ async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
tool.clear()
logging.info(f"Cleared conversation history and image for user {user_id}")
await update.message.reply_text("Conversation history and image cleared. Let's start fresh!")
await update.message.reply_text("Conversation history and image cleared. Let's start fresh!", reply_markup=get_keyboard())
async def update_status_message(context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
keyboard = [
@@ -113,17 +108,12 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
response = get_chat_response(messages)
tool_calls = []
if use_anthropic:
fullMessage = []
for message_part in response.content:
fullMessage.append(message_part)
if message_part.type == "tool_use":
tool_calls.append(message_part)
messages.append({"role": "assistant", "content": fullMessage})
else:
assistant_message = response.choices[0].message
if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None:
tool_calls.append(assistant_message.function_call)
fullMessage = []
for message_part in response.content:
fullMessage.append(message_part)
if message_part.type == "tool_use":
tool_calls.append(message_part)
messages.append({"role": "assistant", "content": fullMessage})
toolUseCount = 0
@@ -146,30 +136,17 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
messages.append(formatted_result)
response = get_chat_response(messages)
assistant_message = ""
if use_anthropic:
fullMessage = []
for message_part in response.content:
fullMessage.append(message_part)
if message_part.type == "tool_use":
tool_calls.append(message_part)
messages.append({"role": "assistant", "content": fullMessage})
else:
assistant_message = response.choices[0].message
conversation_history[user_id].append({"role": "assistant", "content": assistant_message})
if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None:
tool_calls.append(assistant_message.function_call)
else:
conversation_history[user_id].append({"role": "assistant", "content": assistant_message})
assistant_reply = assistant_message
fullMessage = []
for message_part in response.content:
fullMessage.append(message_part)
if message_part.type == "tool_use":
tool_calls.append(message_part)
messages.append({"role": "assistant", "content": fullMessage})
toolUseCount += 1
if (toolUseCount == 0):
if use_anthropic:
assistant_reply = response.content
else:
assistant_reply = assistant_message
assistant_reply = response.content
conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
if len(conversation_history[user_id]) > 20:
@@ -185,28 +162,17 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
except Exception as e:
logging.error(f"An error occurred: {str(e)}")
await update.message.reply_text("Sorry, an error occurred while processing your request.")
await update.message.reply_text("Sorry, an error occurred while processing your request.", reply_markup=get_keyboard())
def call_tool(function_call):
function_name = function_call.name if use_anthropic else function_call.name
function_args = json.dumps(function_call.input) if use_anthropic else function_call.arguments
function_name = function_call.name
function_args = json.dumps(function_call.input)
for tool in tools:
if function_name in [f["name"] for f in tool.get_functions()]:
return tool.execute(function_name, **json.loads(function_args))
def get_chat_response(messages):
return get_claude_response(messages) if use_anthropic else get_openai_response(messages)
def get_openai_response(messages):
model = GPT_4O if use_smart_model else GPT_4O_MINI
response = openai_client.chat.completions.create(
model=model,
messages = [{"role": "system", "content": system_prompt}] + messages,
functions=functions,
function_call="auto",
max_tokens=model_max_tokens[model]
)
return response
return get_claude_response(messages)
def get_claude_response(messages):
anthropic_tools = [
@@ -231,26 +197,8 @@ def get_claude_response(messages):
return response
async def switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
global use_smart_model
use_smart_model = not use_smart_model
model = GPT_4O if use_smart_model else GPT_4O_MINI
logging.info(f"Switched to model: {model}")
await update.message.reply_text(f"Switched to model: {model}")
async def switch_providers(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await clear(update, context)
global use_anthropic
use_anthropic = not use_anthropic
logging.info("Using Anthropic" if use_anthropic else "Using OpenAI")
await update.message.reply_text("Using Anthropic" if use_anthropic else "Using OpenAI")
async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
if use_anthropic:
await update.message.reply_text("Currently using claude-3-5-sonnet-20240620")
else:
model = GPT_4O if use_smart_model else GPT_4O_MINI
await update.message.reply_text(f"Currently using: {model}")
await update.message.reply_text("Currently using claude-3-5-sonnet-20240620", reply_markup=get_keyboard())
async def abort_processing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
query = update.callback_query