add image transcriber function

agent.py

@@ -1,6 +1,7 @@
 from smolagents import ToolCallingAgent, DuckDuckGoSearchTool, VisitWebpageTool, AmazonBedrockServerModel
 
 from tools import AudioTranscriber, txt_reader, pdf_reader, excel_reader, math_calculator
+from utils import BEDROCK_MODEL_ID
 
 MAIN_PROMPT = """
 You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template:
@@ -20,7 +21,7 @@ class MainAgent:
     """Main agent that orchestrates file handling and web search tasks."""
 
     def __init__(self):
-        model = AmazonBedrockServerModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0")
+        model = AmazonBedrockServerModel(model_id=BEDROCK_MODEL_ID)
         file_handler_agent = ToolCallingAgent(
             name="FileHandlerAgent",
             description="Handles file-related tasks such as reading Excel, text, PDF files, etc.",

tools.py

@@ -2,6 +2,7 @@
 import ast
 import json
 import os
+import base64
 from time import sleep
 from uuid import uuid4
 
@@ -11,18 +12,18 @@ from pandas import read_excel
 from smolagents import tool, Tool
 
 from definitions import TranscriptionJob
-from utils import get_file, s3_upload_file, s3_download_file
+from utils import get_file, s3_upload_file, s3_download_file, bedrock_runtime, BEDROCK_MODEL_ID
 
 
 @tool
-def math_calculator(query: str) -> str:
+def math_calculator(expression: str) -> str:
     """A simple calculator tool that evaluates mathematical expressions.
 
     Args:
-        query (str): A mathematical expression as a string, e.g., "2 + 2 * 3".
+        expression (str): A mathematical expression as a string, e.g., "2 + 2 * 3".
     """
     try:
-        result = ast.literal_eval(query)
+        result = ast.literal_eval(expression)
         return str(result)
     except Exception as e:
         return f"Error evaluating expression: {e}"
@@ -137,3 +138,45 @@ class AudioTranscriber(Tool):  # pylint: disable=C0115
             return transcription
         except Exception as e:
             return f"Error starting transcription job for {file_name}: {e}"
+
+
+@tool
+def image_transcriber(text_prompt: str, task_id: str, file_name: str) -> str:
+    """Transcribes text from an image file
+
+    Args:
+        text_prompt (str): The text prompt to guide the transcription.
+        task_id (str): The ID of the task associated with the image file.
+        file_name (str): The name of the image file to transcribe.
+    """
+    try:
+        file_content = get_file(task_id)
+        base64_image = base64.b64encode(file_content).decode("utf-8")
+        response = bedrock_runtime.invoke_model(
+            modelId=BEDROCK_MODEL_ID,
+            body=json.dumps(
+                {
+                    "anthropic_version": "bedrock-2023-05-31",
+                    "max_tokens": 4096,
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "image",
+                                    "source": {
+                                        "type": "base64",
+                                        "media_type": f"image/{file_name.split('.')[-1]}",
+                                        "data": base64_image,
+                                    },
+                                },
+                                {"type": "text", "text": text_prompt},
+                            ],
+                        }
+                    ],
+                }
+            ),
+        ).read()
+        return json.loads(response)["message"]["content"][0]["text"]
+    except Exception as e:
+        return f"Error processing image file {file_name}: {e}"

utils.py

@@ -6,16 +6,11 @@ import requests
 from dotenv import load_dotenv
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-
+BEDROCK_MODEL_ID = "anthropic.claude-3-5-sonnet-20241022-v2:0"
 
 load_dotenv()
 
-bedrock_client = boto3.client(
-    "bedrock-runtime",
-    region_name=os.getenv("AWS_REGION"),
-    aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
-    aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
-)
+bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.getenv("AWS_REGION"))
 
 
 def get_file(task_id: str) -> BytesIO: