user-simulator 0.1.0__py3-none-any.whl

user_sim/handlers/image_recognition_module.py

@@ -0,0 +1,139 @@
+import re
+import logging
+from langchain.schema.messages import HumanMessage, SystemMessage
+from user_sim.utils.token_cost_calculator import calculate_cost, max_input_tokens_allowed, max_output_tokens_allowed
+from user_sim.utils import config
+from user_sim.utils.utilities import init_model
+from user_sim.utils.register_management import save_register, load_register, hash_generate
+
+
+logger = logging.getLogger('Info Logger')
+model = None
+llm = None
+
+
+image_register_name = "image_register.json"
+
+def init_vision_module():
+    global model
+    global llm
+    model, llm = init_model()
+
+
+
+def generate_image_description(image, url=True, detailed=False):
+
+    if not url:
+        image_parsed = f"data:image/png;base64,{image.decode('utf-8')}"
+    else:
+        image_parsed = image
+
+    if detailed:
+        prompt = ("""
+                  Describe in detail this image and its content. 
+                  If there's text, describe everything you read. don't give vague descriptions.
+                  If there is content listed, read it as it is.
+                  Be as detailed as possible.
+                  """)
+    else:
+        prompt = "briefly describe this image, don't over explain, just give a simple and fast explanation of the main characteristics."
+
+    if llm is None:
+        logger.error("vision module not initialized.")
+        return "Empty data"
+
+    if max_input_tokens_allowed(prompt, model, image=image):
+        logger.error(f"Token limit was surpassed")
+        return None
+
+    message = HumanMessage(
+                content=[
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": image_parsed,
+                            # "detail": "auto"
+                        }
+                    }
+                ]
+            )
+
+    try:
+        if config.token_count_enabled:
+            llm.max_tokens = max_output_tokens_allowed(model)
+            output = llm.invoke([message])
+        else:
+            output = llm.invoke([message])
+        output_text = f"(Image description: {output.content})"
+    except Exception as e:
+        logger.error(e)
+        logger.error("Couldn't get image description")
+        output_text = "Empty data"
+    logger.info(output_text)
+    if config.token_count_enabled:
+        calculate_cost(prompt, output_text, model=model, module="image recognition module", image=image)
+
+    return output_text
+
+def image_description(image, detailed=False, url=True):
+    if config.ignore_cache:
+        register = {}
+        logger.info("Cache will be ignored.")
+    else:
+        register = load_register(image_register_name)
+
+    image_hash = hash_generate(content=image)
+
+    if image_hash in register:
+        if config.update_cache:
+            description = generate_image_description(image, url, detailed)
+            register[image_hash] = description
+            logger.info("Cache updated!")
+        # description = register[image_hash]
+        logger.info("Retrieved information from cache.")
+        return register[image_hash]
+    else:
+        description = generate_image_description(image, url)
+        register[image_hash] = description
+
+    if config.ignore_cache:
+        logger.info("Images cache was ignored")
+    else:
+        save_register(register, image_register_name)
+        logger.info("Images cache was saved!")
+
+    return description
+
+
+def image_processor(text):
+
+    def get_images(phrase):
+        pattern = r"<image>(.*?)</image>"
+        matches = re.findall(pattern, phrase)
+        return matches
+
+    def replacer(match):
+        nonlocal replacement_index, descriptions
+        if replacement_index < len(descriptions):
+            original_image = match.group(1)
+            replacement = descriptions[replacement_index]
+            replacement_index += 1
+            return f"<image>{original_image}</image> {replacement}"
+        return match.group(0)  # If no more replacements, return the original match
+
+    if text is None:
+        return text
+    else:
+        images = get_images(text)
+        if images:
+            descriptions = []
+            for image in images:
+                descriptions.append(image_description(image))
+
+            replacement_index = 0
+
+            result = re.sub(r"<image>(.*?)</image>", replacer, text)
+            return result
+        else:
+            return text

user_sim/handlers/pdf_parser_module.py

@@ -0,0 +1,123 @@
+import fitz
+import base64
+import re
+import logging
+import os
+import requests
+from urllib.parse import urlparse
+from user_sim.utils import config
+from user_sim.utils.register_management import save_register, load_register, hash_generate
+from user_sim.handlers.image_recognition_module import image_description
+
+logger = logging.getLogger('Info Logger')
+current_script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.abspath(os.path.join(current_script_dir, "../../.."))
+pdf_register_name = "pdf_register.json"
+
+
+def pdf_reader(pdf):
+
+    if config.ignore_cache:
+        register = {}
+        logger.info("Cache will be ignored.")
+    else:
+        register = load_register(pdf_register_name)
+
+    pdf_hash = hash_generate(content_type="pdf", content=pdf)
+
+    def process_pdf(pdf_file):
+        doc = fitz.open(pdf_file)
+        plain_text = ""
+        for page_number in range(len(doc)):
+            page = doc.load_page(page_number)
+            plain_text += f"Page nª{page_number}: {page.get_text()} "
+
+            images = page.get_images(full=True)
+            if images:
+                plain_text += f"Images in this page: "
+                for img_index, img in enumerate(images):
+                    xref = img[0]
+                    base_image = doc.extract_image(xref)
+                    image_bytes = base_image["image"]
+                    image_base64 = base64.b64encode(image_bytes)
+                    description = image_description(image_base64, detailed=False, url=False)
+                    plain_text += f"Image description {img_index + 1}: {description}"
+        return f"(PDF content: {plain_text} >>)"
+
+    if pdf_hash in register:
+        if config.update_cache:
+            output_text = process_pdf(pdf)
+            register[pdf_hash] = output_text
+            logger.info("Cache updated!")
+        output_text = register[pdf_hash]
+        logger.info("Retrieved information from cache.")
+
+    else:
+        output_text = process_pdf(pdf)
+        register[pdf_hash] = output_text
+
+    if config.ignore_cache:
+        logger.info("PDF cache was ignored.")
+    else:
+        save_register(register, pdf_register_name)
+        logger.info("PDF cache was saved!")
+
+    logger.info(output_text)
+    return output_text
+
+
+def get_pdf(url):
+    # response = requests.get(url)
+    headers = {"User-Agent": "Mozilla/5.0"}
+    response = requests.get(url, headers=headers)
+    if response.status_code != 200:
+        return f"Error accessing the page: {response.status_code}"
+
+    response.encoding = response.apparent_encoding
+    content_type = response.headers.get("Content-Type", "")
+
+    filename = None
+    content_disposition = response.headers.get('Content-Disposition', '')
+
+    if 'application/pdf' in content_type:
+        extension = ".pdf"
+
+        pdfs_dir = os.path.join(project_root, "data/pdfs")
+
+        if not os.path.exists(pdfs_dir):
+            os.makedirs(pdfs_dir)
+
+        if 'filename=' in content_disposition:
+            filename_match = re.search(r'filename="?([^"]+)"?', content_disposition)
+            if filename_match:
+                filename = filename_match.group(1)
+
+        if not filename:
+            parsed_url = urlparse(url)
+            filename = os.path.basename(parsed_url.path)
+            if not filename:
+                filename = 'pdf_download'
+            if extension and not filename.endswith(extension):
+                filename += extension
+
+        full_path = os.path.join(pdfs_dir, filename)
+        content = response.content
+
+        with open(full_path, 'wb') as f:
+            f.write(content)
+
+        return full_path
+
+    else:
+        return None
+
+
+
+def pdf_processor(pdf_url):
+    if pdf_url is None:
+        return pdf_url
+    else:
+        pdf_path = get_pdf(pdf_url)
+        if pdf_path is not None:
+            description = pdf_reader(pdf_path)
+            return description

user_sim/utils/config.py

@@ -0,0 +1,47 @@
+# execution data
+errors = []
+conversation_name = ""
+serial = ""
+# model = "gpt-4o-mini"
+cost_ds_path = None
+test_name = ''
+ignore_cache = False
+update_cache = False
+clean_cache = False
+
+
+# project data
+root_path = ""
+project_folder_path = ""
+profiles_path = ""
+types_folder_path = ""
+test_cases_folder = ""
+types_dict = {}
+custom_personalities_folder = ""
+
+
+# cost metrics
+token_count_enabled = False
+limit_cost = 10000000000
+limit_individual_cost = 10000000000
+total_cost = 0
+total_individual_cost = 0
+
+#llm
+model = "gpt-4o-mini"
+model_provider = "openai"
+
+
+# context
+default_context = [
+    "You are a helpful user simulator that test chatbots.",
+    "Don't add starting sentences, for example 'Okay, here we go'. The first thing you say must be already in the role of a user"
+    "You must act like a user since the beginning of the conversation. "
+    "never recreate a whole conversation, just act like you're a user or client",
+    "never generate a message starting by 'user:'",
+    'Sometimes, interact with what the assistant just said.',
+    'Never act as the assistant, always behave as a user.',
+    "Don't end the conversation until you've asked everything you need.",
+    "you're testing a chatbot, so there can be random values or irrational things "
+    "in your requests"
+]

user_sim/utils/cost_tracker.py

@@ -0,0 +1,153 @@
+import os
+import csv
+import pandas as pd
+import tiktoken
+
+from user_sim.utils import config
+from datetime import datetime
+from typing import List, Union
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import LLMResult, ChatGeneration, ChatResult
+from langchain.chat_models import ChatModel
+import logging
+
+logger = logging.getLogger('Info Logger')
+
+cost_rates = {
+    "gpt-3.5-turbo": {"prompt": 0.0015, "completion": 0.0020},
+    "gpt-4":         {"prompt": 0.03,   "completion": 0.06},
+    # later, if you try e.g. Anthropic:
+    "claude-v1":     {"prompt": 0.0075, "completion": 0.015},
+}
+
+columns = ["Conversation",
+           "Test Name",
+           "Module",
+           "Model",
+           "Total Cost",
+           "Timestamp",
+           "Input Cost",
+           "Input Message",
+           "Output Cost",
+           "Output Message"]
+
+def create_cost_dataset(serial, test_cases_folder):
+    folder = f"{test_cases_folder}/reports/__cost_reports__"
+    file = f"cost_report_{serial}.csv"
+
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+        logger.info(f"Created cost report folder at: {folder}")
+
+    path = f"{folder}/{file}"
+    if not os.path.exists(path):
+        cost_df = pd.DataFrame(columns=columns)
+        cost_df.to_csv(path, index=False)
+        config.cost_ds_path = path
+        logger.info(f"Cost dataframe created at {path}.")
+
+    return path
+
+def count_message_tokens(messages: List[dict], model_name: str) -> int:
+    # """Rough token count for a list of role/content dicts via tiktoken."""
+    # enc = tiktoken.encoding_for_model(model_name)
+    # total = 0
+    # for msg in messages:
+    #     # include role token + content tokens
+    #     total += len(enc.encode(msg["role"]))
+    #     total += len(enc.encode(msg["content"]))
+    # return total
+    pass
+
+
+class CostTrackingCallback(BaseCallbackHandler):
+    """
+    A LangChain callback that, on each LLM invocation,
+      1. extracts token usage
+      2. computes cost via a per-model rate table
+      3. appends a row into a CSV in real time
+    """
+    def __init__(self):
+        """
+        csv_path: where to store the cost log
+        cost_rates: {
+          "<model_name>": {"prompt": <$/1K tokens>, "completion": <$/1K tokens>},
+          ...
+        }
+        """
+        path = create_cost_dataset(config.serial, )
+        self.serial = config.serial
+        self.csv_path = path
+        self.cost_rates = cost_rates
+
+        # If first time, write header
+        # if not os.path.exists(self.csv_path):
+        #     with open(self.csv_path, 'w', newline='') as f:
+        #         writer = csv.writer(f)
+        #         writer.writerow([
+        #             "timestamp",
+        #             "model_name",
+        #             "prompt_tokens",
+        #             "completion_tokens",
+        #             "total_tokens",
+        #             "total_cost_usd",
+        #         ])
+
+    def on_llm_end(self, result: LLMResult, **kwargs) -> None:
+        # 1. pull out usage (providers like OpenAI put it in result.llm_output["usage"])
+        usage = result.llm_output.get("usage", {})
+        p_tokens = usage.get("prompt_tokens", 0)
+        c_tokens = usage.get("completion_tokens", 0)
+        total = usage.get("total_tokens", p_tokens + c_tokens)
+
+        # 2. determine model name key
+        model_name = (
+            result.llm_output.get("model_name")
+            or getattr(result, "model_name", None)
+            or "unknown"
+        )
+
+        # 3. look up rates (per 1K tokens) and compute
+        rates = self.cost_rates.get(model_name, {})
+        cost = (p_tokens / 1000) * rates.get("prompt", 0.0) \
+             + (c_tokens / 1000) * rates.get("completion", 0.0)
+
+        # 4. append to CSV (real‐time)
+        with open(self.csv_path, 'a', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                datetime.utcnow().isoformat(),
+                model_name,
+                p_tokens,
+                c_tokens,
+                total,
+                f"{cost:.8f}",
+            ])
+
+    def budgeted_invoke(
+        self,
+        llm: ChatModel,
+        messages: List[dict],
+        **invoke_kwargs
+    ) -> ChatResult:
+        """
+        Checks your budget, then either:
+          • returns "" if spent up
+          • calls llm.invoke with max_tokens clipped to budget
+        """
+        model_name = llm.client.model_name  # or however your ChatModel exposes it
+        max_toks = self.get_max_completion_tokens(messages, model_name)
+
+        if max_toks is None:
+            # unknown model: just invoke normally
+            return llm.invoke(messages, callbacks=[self], **invoke_kwargs)
+
+        if max_toks <= 0:
+            # out of money
+            raise BudgetExceeded(f"No budget remaining for model {model_name}")
+
+        # enforce our budget
+        invoke_kwargs.setdefault("config", {})
+        invoke_kwargs["config"]["max_tokens"] = max_toks
+
+        return llm.invoke(messages, callbacks=[self], **invoke_kwargs)

user_sim/utils/cost_tracker_v2.py

@@ -0,0 +1,193 @@
+import os
+import csv
+import pandas as pd
+from datetime import datetime
+from typing import List, Union
+import tiktoken
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import LLMResult, ChatGeneration, ChatResult
+from user_sim.utils import config
+import logging
+
+logger = logging.getLogger('Info Logger')
+
+cost_rates = {
+    "gpt-3.5-turbo": {"prompt": 0.0015, "completion": 0.0020},
+    "gpt-4":         {"prompt": 0.03,   "completion": 0.06},
+    "gpt-4o-mini":   {"prompt": 0.03,   "completion": 0.06},
+    # later, if you try e.g. Anthropic:
+    "claude-v1":     {"prompt": 0.0075, "completion": 0.015},
+}
+
+columns = ["Conversation",
+           "Test Name",
+           "Module",
+           "Model",
+           "Total Cost",
+           "Timestamp",
+           "Input Cost",
+           "Input Message",
+           "Output Cost",
+           "Output Message"]
+
+def create_cost_dataset(serial, test_cases_folder):
+
+    if config.cost_ds_path is None:
+        folder = f"{test_cases_folder}/reports/__cost_reports__"
+        file = f"cost_report_{serial}.csv"
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+            logger.info(f"Created cost report folder at: {folder}")
+
+        path = f"{folder}/{file}"
+
+        cost_df = pd.DataFrame(columns=columns)
+        cost_df.to_csv(path, index=False)
+        config.cost_ds_path = path
+        logger.info(f"Cost dataframe created at {path}.")
+        return path
+    else:
+        return config.cost_ds_path
+
+
+
+class BudgetExceeded(Exception):
+    """Raised when there's no budget left for a new LLM call."""
+    pass
+
+def count_message_tokens(messages: List[dict], model_name: str) -> int:
+    """Rough token count for a list of role/content dicts via tiktoken."""
+    enc = tiktoken.encoding_for_model(model_name)
+    total = 0
+    for msg in messages:
+        # include role token + content tokens
+        total += len(enc.encode(msg["role"]))
+        total += len(enc.encode(msg["content"]))
+    return total
+
+class CostTrackingCallback(BaseCallbackHandler):
+    """
+    1. Logs every call’s token usage and cost to CSV.
+    2. Tracks cumulative spend against a total budget.
+    3. Provides a `budgeted_invoke` helper to enforce the budget.
+    """
+
+    def __init__(
+        self,
+        # csv_path: str,
+        # cost_rates: dict,
+        # total_budget_usd: float
+    ):
+        """
+        csv_path:        path to your log CSV
+        cost_rates:      {"model": {"prompt": $/1K, "completion": $/1K}, …}
+        total_budget_usd: how many dollars you're willing to spend in total
+        """
+        path = create_cost_dataset(config.serial, config.test_cases_folder)
+        self.serial = config.serial
+        self.csv_path = path
+        self.cost_rates = cost_rates
+        self.total_budget = 5
+        self.spent = 0.0
+
+
+        # self.csv_path = csv_path
+        # self.cost_rates = cost_rates
+        # self.total_budget = total_budget_usd
+        # self.spent = 0.0
+
+        # write header if new
+        # if not os.path.exists(self.csv_path):
+        #     with open(self.csv_path, "w", newline="") as f:
+        #         w = csv.writer(f)
+        #         w.writerow([
+        #             "timestamp","model_name",
+        #             "prompt_tokens","completion_tokens","total_tokens",
+        #             "call_cost_usd","cumulative_spent_usd"
+        #         ])
+
+    def on_llm_end(self, result: LLMResult, **kwargs) -> None:
+        # pull OpenAI‐style usage
+        usage = result.llm_output.get("usage", {})
+        p = usage.get("prompt_tokens", 0)
+        c = usage.get("completion_tokens", 0)
+        tot = usage.get("total_tokens", p + c)
+
+        # model lookup
+        model = result.llm_output.get("model_name") or getattr(result, "model_name", "unknown")
+        rates = self.cost_rates.get(model, {"prompt":0.0,"completion":0.0})
+        cost = (p/1000)*rates["prompt"] + (c/1000)*rates["completion"]
+
+        # update spend
+        self.spent += cost
+        print(self.spent)
+        # append row
+        with open(self.csv_path, "a", newline="") as f:
+            w = csv.writer(f)
+            w.writerow([
+                datetime.utcnow().isoformat(),
+                model, p, c, tot,
+                f"{cost:.8f}",
+                f"{self.spent:.8f}"
+            ])
+
+    def get_max_completion_tokens(self,
+                                  messages: List[dict],
+                                  model_name: str
+                                 ) -> Union[int, None]:
+        """
+        Based on your remaining budget, estimate how many completion tokens you can afford.
+        Returns None if model not in cost_rates.
+        """
+        if model_name not in self.cost_rates:
+            return None
+
+        rates = self.cost_rates[model_name]
+        # 1. count prompt tokens, cost them
+        p_tokens = count_message_tokens(messages, model_name)
+        cost_prompt = (p_tokens/1000) * rates["prompt"]
+        remaining = self.total_budget - self.spent - cost_prompt
+
+        # 2. if no budget left, signal
+        if remaining <= 0:
+            return 0
+
+        # 3. convert back into max tokens for completion
+        return int((remaining * 1000) / rates["completion"])
+
+
+
+def budgeted_invoke(
+    llm,
+    callbacks,
+    **invoke_kwargs
+) -> ChatResult:
+    """
+    Checks your budget, then either:
+      • returns "" if spent up
+      • calls llm.invoke with max_tokens clipped to budget
+    """
+
+
+    model_name = llm.model_name  # or however your ChatModel exposes it
+    max_toks = self.get_max_completion_tokens(messages, model_name)
+
+    invoke_kwargs.setdefault("config", {})
+    invoke_kwargs["config"].update({
+        "max_tokens": max_toks,
+        "callbacks": [self],  # ← put it here
+    })
+    if max_toks is None:
+        # unknown model: just invoke normally
+        return llm.invoke(messages, **invoke_kwargs)
+
+    if max_toks <= 0:
+        # out of money
+        raise BudgetExceeded(f"No budget remaining for model {model_name}")
+
+    # enforce our budget
+    invoke_kwargs.setdefault("config", {})
+    invoke_kwargs["config"]["max_tokens"] = max_toks
+
+    return llm.invoke(messages, callbacks=[self], **invoke_kwargs)

user_sim/utils/errors.py

@@ -0,0 +1,15 @@
+GOAL_NOT_COMPLETED = 1001
+NO_RESPONSE = 500
+TIMEOUT = 504
+EXCEEDED_LOOP_LIMIT = 1000
+MAX_COST_EXCEEDED = 2000
+MAX_INDIVIDUAL_COST_EXCEEDED = 2001
+
+all_errors = {
+    'goal_not_completed': GOAL_NOT_COMPLETED,
+    'no_response': NO_RESPONSE,
+    'timeout': TIMEOUT,
+    'exceeded_loop_limit': EXCEEDED_LOOP_LIMIT,
+    'max_cost_exceeded': MAX_COST_EXCEEDED,
+    'max_individual_cost_exceeded': MAX_INDIVIDUAL_COST_EXCEEDED
+}