From 5e02468a8b898380aa3e7ed6aacd0d1f5fc3ca97 Mon Sep 17 00:00:00 2001
From: bucolucas <bucolucas@gmail.com>
Date: Sun, 18 Aug 2024 13:20:33 -0500
Subject: [PATCH] Update AI provider usage and imports

---
 telegram_inference_bot.py | 35 ++++++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py
index 9f725fa..dfb2692 100644
--- a/telegram_inference_bot.py
+++ b/telegram_inference_bot.py
@@ -7,7 +7,7 @@ from telegram import Update
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
 from dotenv import load_dotenv
 from tools.base_tool import BaseTool
-from ai_providers import create_ai_provider
+from ai_providers import create_ai_provider, AnthropicProvider, OpenAIProvider
 
 # Load environment variables
 load_dotenv()
@@ -86,7 +86,7 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
             
             tool_response = call_tool(tool_call)
 
-            formatted_result = ai_provider.format_tool_result(tool_call, tool_response)
+            formatted_result = {"role": "function", "name": function_name, "content": json.dumps(tool_response)}
             messages.append(formatted_result)
 
             response = ai_provider.get_chat_response([{"role": "system", "content": system_prompt}] + messages)
@@ -95,13 +95,14 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
             toolUseCount += 1
 
         if toolUseCount == 0:
-            assistant_reply = ai_provider.format_assistant_reply(response)
+            assistant_reply = response.content[0].text if isinstance(ai_provider, AnthropicProvider) else response.choices[0].message.content
             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:]
 
-        await update.message.reply_text(ai_provider.get_reply_text(response))
+        reply_text = response.content[0].text if isinstance(ai_provider, AnthropicProvider) else response.choices[0].message.content
+        await update.message.reply_text(reply_text)
 
     except Exception as e:
         logging.error(f"An error occurred: {str(e)}")
@@ -116,19 +117,31 @@ def call_tool(function_call):
 
 async def switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     global ai_provider
-    model = ai_provider.switch_model()
-    logging.info(f"Switched to model: {model}")
-    await update.message.reply_text(f"Switched to model: {model}")
+    if isinstance(ai_provider, OpenAIProvider):
+        ai_provider.use_smart_model = not ai_provider.use_smart_model
+        ai_provider.model = ai_provider.get_model()
+        logging.info(f"Switched to model: {ai_provider.model}")
+        await update.message.reply_text(f"Switched to model: {ai_provider.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 ai_provider
-    ai_provider = create_ai_provider("openai" if ai_provider.name == "anthropic" else "anthropic")
-    logging.info(f"Switched to {ai_provider.name} provider")
-    await update.message.reply_text(f"Switched to {ai_provider.name} 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:
-    await update.message.reply_text(f"Currently using {ai_provider.name}: {ai_provider.get_model()}")
+    if isinstance(ai_provider, AnthropicProvider):
+        await update.message.reply_text(f"Currently using Anthropic: {ai_provider.model}")
+    else:
+        await update.message.reply_text(f"Currently using OpenAI: {ai_provider.model}")
 
 def main() -> None:
     # Create the Application and pass it your bot's token