utils.py

import base64
import hashlib

import contextlib

from functools import wraps
from typing import Callable, Any
from inspect import signature, Parameter, Signature


def extract_substring(text, start_tag, end_tag):
    start_position = text.lower().find(start_tag.lower())
    end_position = text.lower().find(end_tag.lower())

    if start_position == -1 or end_position == -1:
        return None

    return text[start_position + len(start_tag):end_position].strip()

def compute_hash(s):
    # Hash a string using SHA-1 and return the base64 encoded result

    m = hashlib.sha1()
    m.update(s.encode())

    b = m.digest()

    return base64.b64encode(b).decode('ascii')

def add_params_and_annotations(name: str, description: str, params: dict, return_type: Any):
    """
    A decorator to add parameters and a return type annotation to a function.

    :param params: A dictionary where the keys are parameter names and values are their types.
    :param return_type: The return type to add to the function's signature.
    """
    def decorator(func: Callable):
        # Create new parameters based on the provided params dict
        new_params = [
            Parameter(name, Parameter.POSITIONAL_OR_KEYWORD, annotation=type_)
            for name, (type_, _) in params.items()
        ]
        
        # Get the existing signature and parameters
        original_sig = signature(func)
        original_params = list(original_sig.parameters.values())

        # Combine new parameters with the existing ones
        combined_params = new_params# + original_params

        # Create a new signature with updated parameters and return annotation
        new_sig = Signature(parameters=combined_params, return_annotation=return_type)
        
        # Define the wrapper function
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        docstring = description

        if len(params) > 0:
            docstring += '\n\nArgs:'
            for param_name, (type_, param_description) in params.items():
                docstring += f'\n  {param_name}: {param_description}'
        
        docstring += f'\n\nReturns:\n  {return_type.__name__}'

        print('Input params:', params)


        # Set the new signature on the wrapper
        wrapper.__name__ = name
        wrapper.__signature__ = new_sig
        wrapper.__annotations__.update({ k: v[0] for k, v in params.items() })
        wrapper.__annotations__['return'] = return_type
        #wrapper.__annotations__['name'] = str
        wrapper.__doc__ = docstring

        print(docstring)

        return wrapper
    return decorator

_cost_tracker = {}

@contextlib.contextmanager
def use_cost_tracker():
    _cost_tracker.clear()
    _cost_tracker['conversation'] = 0
    _cost_tracker['negotiation'] = 0
    _cost_tracker['programming'] = 0
    yield

def register_cost(category, cost):
    _cost_tracker[category] += cost

def get_costs():
    return _cost_tracker