Update agents/tools/ai_tools.py

agents/tools/ai_tools.py

@@ -7,21 +7,25 @@
 import base64, chess, os
 from agents.models.llms import (
     LLM_WEB_SEARCH,
-    LLM_IMAGE_TO_FEN,
+    LLM_WEB_BROWSER,
     LLM_IMAGE_ANALYSIS,
     LLM_AUDIO_ANALYSIS,
     LLM_VIDEO_ANALYSIS,
     LLM_YOUTUBE_ANALYSIS,
     LLM_DOCUMENT_ANALYSIS,
-    LLM_ARITHMETIC,
     LLM_CODE_GENERATION,
     LLM_CODE_EXECUTION,
-    LLM_WEB_BROWSER,
+    LLM_IMAGE_TO_FEN,
+    LLM_ALGEBRAIC_CHESS_NOTATION,
     LLM_FINAL_ANSWER,
     THINKING_LEVEL_TOOLS,
     THINKING_LEVEL_FINAL_ANSWER
 )
-from agents.models.prompts import PROMPT_IMG_TO_FEN, PROMPT_FINAL_ANSWER
+from agents.models.prompts import (
+    PROMPT_IMG_TO_FEN,
+    PROMPT_ALGEBRAIC_CHESS_NOTATION,
+    PROMPT_FINAL_ANSWER
+)
 from crewai.tools import tool
 from crewai_tools import StagehandTool
 from google import genai
@@ -106,123 +110,6 @@ class AITools():
         except Exception as e:
             raise RuntimeError(f"Processing failed: {str(e)}")
 
-    @tool("Image to FEN Tool")
-    def img_to_fen_tool(question: str, file_path: str) -> str:
-        """Given a chess question and image file, return the FEN.
-    
-        Args:
-            question (str): The chess question
-            file_path (str): The image file path
-            
-        Returns:
-            str: FEN of the chess position
-            
-        Raises:
-            RuntimeError: If processing fails
-        """
-        try:         
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-
-            with open(file_path, "rb") as f:
-                img_bytes = f.read()
-                img_b64 = base64.b64encode(img_bytes).decode("ascii")
-            
-            prompt = PROMPT_IMG_TO_FEN.format(question=question)
-            
-            content = types.Content(
-                parts=[
-                    types.Part(text=prompt),
-                    types.Part(
-                        inline_data=types.Blob(
-                            mime_type="image/png",
-                            data=base64.b64decode(img_b64),
-                        )
-                    )
-                ]
-            )
-            
-            response = client.models.generate_content(
-                model=LLM_IMAGE_TO_FEN,
-                contents=[content],
-                config=types.GenerateContentConfig(
-                    thinking_config=types.ThinkingConfig(
-                        thinking_level=THINKING_LEVEL_TOOLS
-                    )
-                )
-            )
-            
-            fen = None
-            
-            for part in response.parts:
-                if part.text is not None:
-                    fen = part.text.strip()
-                    break
-
-            board = chess.Board(fen) # FEN validation
-            
-            print(f"🤖 FEN: {fen}")
-
-            return fen;
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
-    @tool("Algebraic Chess Notation Tool")
-    def algebraic_chess_notation_tool(question: str, file_path: str, best_move: str) -> str:
-        """Given a chess question, image file, and best move with continuation in UCI notation, answer the question in algebraic notation.
-    
-        Args:
-            question (str): The chess question
-            file_path (str): The image file path
-            best_move (str): The best move with continuation in UCI notation
-            
-        Returns:
-            str: Answer to the question in algebraic notation
-            
-        Raises:
-            RuntimeError: If processing fails
-        """
-        try:         
-            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-
-            with open(file_path, "rb") as f:
-                img_bytes = f.read()
-                img_b64 = base64.b64encode(img_bytes).decode("ascii")
-            
-            prompt = PROMPT_UCI_TO_ALGEBRAIC.format(question=question, best_move=best_move)
-            
-            content = types.Content(
-                parts=[
-                    types.Part(text=prompt),
-                    types.Part(
-                        inline_data=types.Blob(
-                            mime_type="image/png",
-                            data=base64.b64decode(img_b64),
-                        )
-                    )
-                ]
-            )
-            
-            response = client.models.generate_content(
-                model=LLM_UCI_TO_ALGEBRAIC,
-                contents=[content],
-                config=types.GenerateContentConfig(
-                    thinking_config=types.ThinkingConfig(
-                        thinking_level=THINKING_LEVEL_TOOLS
-                    )
-                )
-            )
-            
-            for part in response.parts:
-                if part.text is not None:
-                    result = part.text.strip()
-                    break
-
-            print(f"🤖 Algebraic notation: {result}")
-
-            return result;
-        except Exception as e:
-            raise RuntimeError(f"Processing failed: {str(e)}")
-
     @tool("Image Analysis Tool")
     def image_analysis_tool(question: str, file_path: str) -> str:
         """Given a question and image file, analyze the image to answer the question.
@@ -468,36 +355,120 @@ class AITools():
         except Exception as e:
             raise RuntimeError(f"Processing failed: {str(e)}")
 
-    @tool("Final Answer Tool")
-    def final_answer_tool(question: str, answer: str) -> str:
-        """Given a question and initial answer, generate the final answer.
+    @tool("Image to FEN Tool")
+    def img_to_fen_tool(question: str, file_path: str) -> str:
+        """Given a chess question and image file, return the FEN.
     
         Args:
-            question (str): Question to answer
-            answer (str): The initial answer
+            question (str): The chess question
+            file_path (str): The image file path
             
         Returns:
-            str: Final answer to the question
+            str: FEN of the chess position
             
         Raises:
             RuntimeError: If processing fails
         """
-        try:
+        try:         
             client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-        
-            prompt = PROMPT_FINAL_ANSWER.format(question=question, answer=answer)
-        
+
+            with open(file_path, "rb") as f:
+                img_bytes = f.read()
+                img_b64 = base64.b64encode(img_bytes).decode("ascii")
+            
+            prompt = PROMPT_IMG_TO_FEN.format(question=question)
+            
+            content = types.Content(
+                parts=[
+                    types.Part(text=prompt),
+                    types.Part(
+                        inline_data=types.Blob(
+                            mime_type="image/png",
+                            data=base64.b64decode(img_b64),
+                        )
+                    )
+                ]
+            )
+            
+            response = client.models.generate_content(
+                model=LLM_IMAGE_TO_FEN,
+                contents=[content],
+                config=types.GenerateContentConfig(
+                    thinking_config=types.ThinkingConfig(
+                        thinking_level=THINKING_LEVEL_TOOLS
+                    )
+                )
+            )
+            
+            fen = None
+            
+            for part in response.parts:
+                if part.text is not None:
+                    fen = part.text.strip()
+                    break
+
+            board = chess.Board(fen) # FEN validation
+            
+            print(f"🤖 FEN: {fen}")
+
+            return fen;
+        except Exception as e:
+            raise RuntimeError(f"Processing failed: {str(e)}")
+
+    @tool("Algebraic Chess Notation Tool")
+    def algebraic_chess_notation_tool(question: str, file_path: str, best_move: str) -> str:
+        """Given a chess question, image file, and best move with continuation in UCI notation, answer the question in algebraic notation.
+    
+        Args:
+            question (str): The chess question
+            file_path (str): The image file path
+            best_move (str): The best move with continuation in UCI notation
+            
+        Returns:
+            str: Answer to the question in algebraic notation
+            
+        Raises:
+            RuntimeError: If processing fails
+        """
+        try:         
+            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+
+            with open(file_path, "rb") as f:
+                img_bytes = f.read()
+                img_b64 = base64.b64encode(img_bytes).decode("ascii")
+            
+            prompt = PROMPT_ALGEBRAIC_CHESS_NOTATION.format(question=question, best_move=best_move)
+            
+            content = types.Content(
+                parts=[
+                    types.Part(text=prompt),
+                    types.Part(
+                        inline_data=types.Blob(
+                            mime_type="image/png",
+                            data=base64.b64decode(img_b64),
+                        )
+                    )
+                ]
+            )
+            
             response = client.models.generate_content(
-                model=LLM_FINAL_ANSWER, 
-                contents=[prompt],
+                model=LLM_ALGEBRAIC_CHESS_NOTATION,
+                contents=[content],
                 config=types.GenerateContentConfig(
                     thinking_config=types.ThinkingConfig(
-                        thinking_level=THINKING_LEVEL_FINAL_ANSWER
+                        thinking_level=THINKING_LEVEL_TOOLS
                     )
                 )
             )
             
-            return response.text
+            for part in response.parts:
+                if part.text is not None:
+                    result = part.text.strip()
+                    break
+
+            print(f"🤖 Algebraic notation: {result}")
+
+            return result;
         except Exception as e:
             raise RuntimeError(f"Processing failed: {str(e)}")