Added ability to switch between smart and dumb models, and back and forth between openai and anthropic

telegram_inference_bot.py

@@ -5,6 +5,7 @@ import inspect
 import tempfile
 import base64
 import logging
+import anthropic
 from telegram import Update
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
 from openai import OpenAI
@@ -14,7 +15,11 @@ from tools.base_tool import BaseTool
 # Load environment variables
 load_dotenv()
 
-client = OpenAI()
+openai_client = OpenAI()
+
+anthropic_client = anthropic.Anthropic(
+    api_key=os.environ.get("ANTHROPIC_API_KEY")
+)
 
 GPT_4O = "gpt-4o"
 GPT_4O_MINI = "gpt-4o-mini"
@@ -25,6 +30,7 @@ model_max_tokens = {
 }
 
 use_smart_model = True
+use_anthropic = False
 
 # Set up logging to console and file
 logging.basicConfig(level=logging.WARNING, handlers=[
@@ -109,33 +115,7 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
         # Prepare messages for OpenAI API
         messages = [{"role": "system", "content": system_prompt}] + conversation_history[user_id]
 
-        # Check if there's an image to process
-        if user_id in user_images:
-            with open(user_images[user_id], "rb") as image_file:
-                response = client.chat.completions.create(
-                    model=GPT_4O_MINI,
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": [
-                                {"type": "text", "text": user_message},
-                                {
-                                    "type": "image_url",
-                                    "image_url": {
-                                        "url": f"data:image/jpeg;base64,{base64.b64encode(image_file.read()).decode('utf-8')}"
-                                    }
-                                },
-                            ],
-                        }
-                    ],
-                    max_tokens=16384
-                )
-            # Remove the temporary image file
-            os.remove(user_images[user_id])
-            del user_images[user_id]
-        else:
-            # Call OpenAI API for inference (text-only)
-            response = get_chat_response(client, messages, GPT_4O if use_smart_model else GPT_4O_MINI)
+        response = get_chat_response(messages)
 
         # Extract the assistant's reply
         assistant_message = response.choices[0].message
@@ -152,7 +132,7 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
                 })
 
                 # Call API again to get the final response
-                assistant_message = get_chat_response(client, messages, GPT_4O if use_smart_model else GPT_4O_MINI).choices[0].message
+                assistant_message = get_chat_response(messages).choices[0].message
                 if not hasattr(assistant_message, 'function_call') or not assistant_message.function_call:
                     assistant_reply = assistant_message.content
                     conversation_history[user_id].append({"role": "assistant", "content": assistant_reply})
@@ -182,8 +162,16 @@ def call_tool(function_call):
         if function_name in [f["name"] for f in tool.get_functions()]:
             return tool.execute(function_name, **eval(function_args))
 
-def get_chat_response(client, messages, model):
-    response = client.chat.completions.create(
+def get_chat_response(messages):
+    if use_anthropic:
+        response = get_openai_response(messages)
+    else:
+        response = get_claude_response(messages)
+    return response
+
+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=messages,
         functions=functions,
@@ -192,6 +180,14 @@ def get_chat_response(client, messages, model):
     )
     return response
 
+def get_claude_response(messages):
+    response = anthropic_client.messages.create(
+        messages=messages,
+        max_tokens=4096,
+        model="claude-3-5-sonnet-20240620"
+    )
+    return response
+
 def switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     global use_smart_model
     use_smart_model = not use_smart_model
@@ -199,6 +195,12 @@ def switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     logging.info(f"Switched to model: {model}")
     update.message.reply_text(f"Switched to model: {model}")
 
+def switch_anthropic(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    global use_anthropic
+    use_anthropic = not use_anthropic
+    logging.info("Using Anthropic" if use_anthropic else "Using OpenAI")
+    update.message.reply_text("Using Anthropic" if use_anthropic else "Using OpenAI")
+
 async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     model = GPT_4O if use_smart_model else GPT_4O_MINI
     await update.message.reply_text(f"Currently using model model: {model}")
@@ -212,6 +214,7 @@ def main() -> None:
     application.add_handler(CommandHandler("start", start))
     application.add_handler(CommandHandler("clear", clear))
     application.add_handler(CommandHandler("switch", switch))
+    application.add_handler(CommandHandler("toggle", switch_anthropic))
     application.add_handler(CommandHandler("status", status))
     application.add_handler(MessageHandler(filters.PHOTO, handle_image))
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))