diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py index 85fc61b..855c16b 100644 --- a/telegram_inference_bot.py +++ b/telegram_inference_bot.py @@ -3,34 +3,15 @@ import os import importlib import inspect import logging -import anthropic from telegram import Update from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes -from openai import OpenAI from dotenv import load_dotenv from tools.base_tool import BaseTool +from ai_providers import create_ai_provider # Load environment variables load_dotenv() -openai_client = OpenAI() - -anthropic_client = anthropic.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 = True -use_anthropic = True - # Set up logging to console and file logging.basicConfig(level=logging.WARNING, handlers=[ logging.StreamHandler(), @@ -63,6 +44,9 @@ functions = [] for tool in tools: functions.extend(tool.get_functions()) +# Initialize AI provider +ai_provider = create_ai_provider("anthropic") + 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.") @@ -91,138 +75,72 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> messages = conversation_history[user_id] - response = get_chat_response(messages) - tool_calls = [] - if use_anthropic: - for message in response.content: - if message.type == "tool_use": - tool_calls.append(message) - else: - messages.append({"role": "assistant", "content": response.content}) - 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) + response = ai_provider.get_chat_response([{"role": "system", "content": system_prompt}] + messages) + tool_calls = ai_provider.format_tool_calls(response) toolUseCount = 0 while len(tool_calls) > 0 and toolUseCount < 50: - tool_call = tool_calls.pop(0) function_name = tool_call.name tool_response = call_tool(tool_call) - formatted_result = {} - - if use_anthropic: - formatted_result = {"role": "user", "content":[{"type": "tool_result", "tool_use_id": tool_call.id, "content": json.dumps(tool_response)}]} - else: - formatted_result = {"role": "function", "name": function_name, "content": json.dumps(tool_response)} - + formatted_result = ai_provider.format_tool_result(tool_call, tool_response) messages.append(formatted_result) - response = get_chat_response(messages) - assistant_message = "" - if use_anthropic: - for message in response.content: - if message.type == "tool_use": - tool_calls.append(message) - else: - messages.append({"role": "assistant", "content": response.content}) - 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) - assistant_reply = assistant_message + response = ai_provider.get_chat_response([{"role": "system", "content": system_prompt}] + messages) + tool_calls = ai_provider.format_tool_calls(response) toolUseCount += 1 - - if (toolUseCount == 0): - if use_anthropic: - assistant_reply = response.content - else: - assistant_reply = assistant_message + if toolUseCount == 0: + assistant_reply = ai_provider.format_assistant_reply(response) conversation_history[user_id].append({"role": "assistant", "content": assistant_reply}) if len(conversation_history[user_id]) > 20: conversation_history[user_id] = conversation_history[user_id][-20:] - if use_anthropic: - await update.message.reply_text(messages[-1]["content"][0].text) - else: - await update.message.reply_text(assistant_reply.content) + await update.message.reply_text(ai_provider.get_reply_text(response)) 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.") 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.loads(function_call.arguments if hasattr(function_call, 'arguments') else 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 - -def get_claude_response(messages): - anthropic_tools = [ - { - "name": function['name'], - "description": function['description'], - "input_schema": function['parameters'] if function['parameters'] not in [None, {}] else {"type": "object", "properties": {"param1": {"type": "string", "description": "Unnecessary"}}, "required": []} - } - for function in functions - ] - try: - response = anthropic_client.messages.create( - model="claude-3-5-sonnet-20240620", - system=system_prompt, - messages=messages, - max_tokens=8192, - tools=anthropic_tools - ) - except Exception as e: - logging.error(f"An error occurred: {str(e)}") - return None - - return response + return tool.execute(function_name, **function_args) 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}") + global ai_provider + if isinstance(ai_provider, OpenAIProvider): + ai_provider.use_smart_model = not ai_provider.use_smart_model + model = ai_provider.get_model() + logging.info(f"Switched to model: {model}") + await update.message.reply_text(f"Switched to model: {model}") + else: + await update.message.reply_text("Switching models is only available for OpenAI provider.") 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") + global ai_provider + if isinstance(ai_provider, AnthropicProvider): + ai_provider = create_ai_provider("openai") + logging.info("Switched to OpenAI provider") + await update.message.reply_text("Switched to OpenAI provider") + else: + ai_provider = create_ai_provider("anthropic") + logging.info("Switched to Anthropic provider") + await update.message.reply_text("Switched to Anthropic provider") 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") + if isinstance(ai_provider, AnthropicProvider): + await update.message.reply_text(f"Currently using Anthropic: {ai_provider.model}") 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(f"Currently using OpenAI: {ai_provider.get_model()}") def main() -> None: # Create the Application and pass it your bot's token