from lotus.models import LMWithTools
import lotus
from lotus.sem_ops.sem_map import sem_map_with_tools 
import pandas as pd 
from typing import Type 
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
import time


class FileReadArgs(BaseModel):
    filename: str = Field(..., description="The name of the file to read.")
    
class FileReadTool(BaseTool):
    name: str = "File Read"
    description: str = "A tool to read files from the local filesystem."
    args_schema: Type[BaseModel] = FileReadArgs


    def __init__(self):
        super().__init__()

    def _run(self, filename: str) -> str:
        print(f"Reading file: {filename}")
        if filename == "text1.txt":
            return "13 + 16"
        elif filename == "text2.txt":
            return "29 + 51"

class AdditionArgs(BaseModel):
    num1: str = Field(..., description="The first number.")
    num2: str = Field(..., description="The second number.")

class AdditionTool(BaseTool):
    name: str = "Addition Tool"
    description: str = "A tool to add two numbers."
    args_schema: Type[BaseModel] = AdditionArgs


    def __init__(self):
        super().__init__()

    def _run(self, num1: int, num2: int) -> str:
        print(f"Adding numbers: {num1} + {num2}")
        return str(int(num1) + int(num2))

lmwithtools = LMWithTools()

lotus.settings.configure(lm_with_tools=lmwithtools)
data = {
"File names": [
    "text1.txt",
    "text2.txt",
]
}

df = pd.DataFrame(data)
user_instruction = "Evaluate the mathametical expression in the file {File names} to a single integer. Give only the integer as output."
df = df.sem_map(user_instruction, tools=[FileReadTool(), AdditionTool()])
print(df)

(lotus) atharvachougule@Atharvas-MacBook-Pro-3 lotus % python3 test.py
Mapping:   0%|                                     0/2 Rows processed [00:00<?, ?it/s]2025-10-27 15:17:23,165 - INFO - OpenAI API usage: {'prompt_tokens': 402, 'completion_tokens': 47, 'total_tokens': 449}
Reading file: text1.txt
2025-10-27 15:17:30,578 - INFO - OpenAI API usage: {'prompt_tokens': 428, 'completion_tokens': 58, 'total_tokens': 486}
Adding numbers: 13 + 16
2025-10-27 15:17:32,052 - INFO - OpenAI API usage: {'prompt_tokens': 402, 'completion_tokens': 512, 'total_tokens': 914}
Reading file: text2.txt
2025-10-27 15:17:36,987 - INFO - OpenAI API usage: {'prompt_tokens': 470, 'completion_tokens': 15, 'total_tokens': 485}
Mapping:  50%|██████████████               1/2 Rows processed [00:15<00:15, 15.69s/it]2025-10-27 15:17:39,308 - INFO - OpenAI API usage: {'prompt_tokens': 428, 'completion_tokens': 64, 'total_tokens': 492}
Adding numbers: 29 + 51
2025-10-27 15:17:45,752 - INFO - OpenAI API usage: {'prompt_tokens': 476, 'completion_tokens': 14, 'total_tokens': 490}
Mapping: 100%|████████████████████████████ 2/2 Rows processed [00:24<00:00, 12.23s/it]
  File names _map
0  text1.txt   29
1  text2.txt   80