#!/usr/bin/env python3
# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import unittest

from unifree.llms.code_extrators import extract_first_source_code, extract_header_implementation


class TestCodeExtractors(unittest.TestCase):

    def test_extract_first_source_code(self):
        target_code = "some code here!"
        responses = [
            f"""
```
{target_code}
```
""",
            f"""
```
{target_code}
```
Some ending comment
""",
            f"""
Some beginning comment
```
{target_code}
```
Some ending comment
""",
            f"""
Some beginning comment
```cpp
{target_code}
```
Some ending comment
""",
            f"""
        Some beginning comment
        ```h
        {target_code}
        ```
        ```cpp
        code we don't care about
        ```
        """,
            f"""
{target_code}
""",
        ]

        for response in responses:
            extracted_code = extract_first_source_code(response)
            self.assertEqual(target_code, extracted_code)

    def test_extract_header_implementation(self):
        target_header = "header here!"
        target_implementation = "implementation here!"
        responses = [
            f"""
    ```cpp
    {target_header}
    ```
    ```cpp
    {target_implementation}
    ```
    """,
            f"""
// Header:
    ```
    {target_header}
    ```
// Implementation:
    ```
    {target_implementation}
    ```
    Some ending comment
    """,
            f"""
    // Header.h
    ```
    {target_header}
    ```
    // Implementation.cpp
    ```
    {target_implementation}
    ```
    """,
            f"""
    
    // Header.h
    
    {target_header}
    
    // Implementation.cpp
    
    {target_implementation}
    """,
        ]

        for response in responses:
            extracted_header, extracted_implementation = extract_header_implementation(response)
            self.assertEqual(target_header, extracted_header, f"Response: {response}")
            self.assertEqual(target_implementation, extracted_implementation, f"Response: {response}")


if __name__ == '__main__':
    unittest.main()

tests/test_code_extrators.py
ProjectUnifree/unifree
#!/usr/bin/env python3
# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)
import os
import re
import unittest
from typing import Dict, TypeVar, Optional, List, Any

import unifree
from unifree import LLM, QueryHistoryItem
from unifree.csharp_migration_strategies import CSharpCompilationUnitMigrationStrategy, CSharpCompilationUnitMigrationWithLLM, CSharpCompilationUnitToSingleFileWithLLM
from unifree.llms import TrivialLLM
from unifree.utils import to_default_dict


class CSharpCompilationUnitMigrationWithLLMProxy(CSharpCompilationUnitMigrationWithLLM):
    def execute(self) -> None:
        pass

    def load_llm(self) -> unifree.LLM:
        return TrivialLLM(self.config)


class TestCSharpCompilationUnitMigrationWithLLM(unittest.TestCase):
    def test_create_code_prompt(self):
        strategy = CSharpCompilationUnitMigrationWithLLMProxy(unifree.FileMigrationSpec('', '', ''), to_default_dict({
            "prompts": {
                "test": "This is a ${CODE}"
            },
            "llm": {
                "class": "TrivialLLM"
            }
        }))

        self.assertTrue("This is a test with 01010", strategy.create_code_prompt('test', "01010"))

    def test_create_prompt(self):
        strategy = CSharpCompilationUnitMigrationWithLLMProxy(unifree.FileMigrationSpec('', '', ''), to_default_dict({
            "prompts": {
                "test": "This is a ${TEST} with ${MACRO}"
            },
            "llm": {
                "class": "TrivialLLM"
            }
        }))

        self.assertTrue("This is a test with ${MACRO}", strategy.create_prompt('test', {"TEST": "test"}))
        self.assertTrue("This is a test with 111", strategy.create_prompt('test', {"TEST": "test", "MACRO": "111"}))

    def test_create_prompt_missing(self):
        strategy = CSharpCompilationUnitMigrationWithLLMProxy(unifree.FileMigrationSpec('', '', ''), to_default_dict({
            "prompts": {
                "test": "This is a ${TEST} with ${MACRO}"
            },
            "llm": {
                "class": "TrivialLLM"
            }

        }))

        try:
            strategy.create_prompt('missing', {})
            self.fail("Exception should have been thrown")
        except RuntimeError:
            pass  # expected


class CSharpCompilationUnitMigrationStrategyProxy(CSharpCompilationUnitMigrationStrategy):
    def execute(self) -> None:
        pass


class TestCSharpCompilationUnitMigrationStrategy(unittest.TestCase):
    config: Dict

    def test_source_text(self):
        for source_file_name in ['MalformedShortClass.cs', 'ShortClassWithNamespace.cs']:
            strategy = CSharpCompilationUnitMigrationStrategyProxy(_load_file_migration_spec(source_file_name), self.config)
            source_text = _load_source_file(source_file_name)

            self.assertEqual(_normalize_definition(source_text), _normalize_definition(strategy.source_text), f"Source: {source_file_name}")

    def test_method_definitions(self):
        reference_methods = [
            """
                public void setRevenue(double amount, string currency)
                {
                    this.revenue = amount;
                    this.currency = currency;
                }
            """,
            """            
                public void setAdImpressionsCount(int adImpressionsCount)
                {
                    this.adImpressionsCount = adImpressionsCount;
                }
            """,
            """            
                public void setAdRevenueNetwork(string adRevenueNetwork)
                {
                    this.adRevenueNetwork = adRevenueNetwork;
                }
            """]
        norm_reference_methods = [_normalize_definition(r) for r in reference_methods]

        strategy = CSharpCompilationUnitMigrationStrategyProxy(_load_file_migration_spec('ShortClassNoNamespace.cs'), self.config)
        method_definitions = strategy.method_declarations

        for method_definition in method_definitions:
            norm_method_definition = _normalize_definition(method_definition)
            self.assertTrue(norm_method_definition in norm_reference_methods, f"Missing {method_definition}")

    def test_method_definitions_from_malformed(self):
        reference_methods = [
            """
                public void CloseBracketMissing(string key, string value)
                {
                    partnerList.Add(key);
                    partnerList.Add(value);
            """,
        ]
        norm_reference_methods = [_normalize_definition(r) for r in reference_methods]

        strategy = CSharpCompilationUnitMigrationStrategyProxy(_load_file_migration_spec('MalformedShortClass.cs'), self.config)
        method_definitions = strategy.method_declarations

        for method_definition in method_definitions:
            norm_method_definition = _normalize_definition(method_definition)
            self.assertTrue(norm_method_definition in norm_reference_methods, f"Missing {method_definition}")

    def test_create_destination_file_path(self):
        strategy = CSharpCompilationUnitMigrationStrategyProxy(unifree.FileMigrationSpec(
            source_file_path="/Test/Folder/source/Assets/sub/SomeFile.cs",
            source_project_path="/Test/Folder/source/",
            destination_project_path="/Test/Folder/destination",

        ), self.config)
        self.assertEqual("/Test/Folder/destination/Assets/sub/SomeFile.gd", strategy.create_destination_file_path('.gd'))

        strategy._config = to_default_dict({
            "target": {
                "convert_filename_to_camelcase": True
            }
        })
        self.assertEqual("/Test/Folder/destination/Assets/sub/SomeFile.gd", strategy.create_destination_file_path('.gd'))

        strategy._config = to_default_dict({
            "target": {
                "convert_filename_to_snake_case": True
            }
        })
        self.assertEqual("/Test/Folder/destination/Assets/sub/some_file.gd", strategy.create_destination_file_path('.gd'))

        strategy._config = to_default_dict({
            "target": {
                "lower_folder_names": True
            }
        })
        self.assertEqual("/Test/Folder/destination/assets/sub/SomeFile.gd", strategy.create_destination_file_path('.gd'))

    def test_class_definition(self):
        class_definitions = {
            'ShortClassNoNamespace.cs': """
                using System;
                using System.Collections.Generic;
                
                public class AdjustAdRevenue 
                {
                    internal string source;
                    internal double? revenue;
                    internal string currency;
                    internal int? adImpressionsCount;
                    internal string adRevenueNetwork;
                    internal string adRevenueUnit;
                    internal string adRevenuePlacement;
                    internal List<string> partnerList;
                    internal List<string> callbackList;
                public AdjustAdRevenue(string source)
                    {
                        this.source = source;
                // #if UNITY_EDITOR
                        this.callbackList = []
                // #endif
                
                    } // #region Revenue // #endregion 
                }
            """,
            'ShortClassWithNamespace.cs': """
                        using System;
                        using System.Collections.Generic;

                        namespace com.adjust.sdk
                        {
                            public class AdjustAdRevenue
                            {
                                internal string source;
                                internal double? revenue;
                                internal string currency;
                                internal int? adImpressionsCount;
                                internal string adRevenueNetwork;
                                internal string adRevenueUnit;
                                internal string adRevenuePlacement;
                                internal List<string> partnerList;
                                internal List<string> callbackList;

                                public AdjustAdRevenue(string source)
                                {
                                    this.source = source;
                                }
                            }
                        }
            """
        }
        norm_class_definitions = {n: _normalize_definition(d) for n, d in class_definitions.items()}

        for source_file, norm_class_definition in norm_class_definitions.items():
            strategy = CSharpCompilationUnitMigrationStrategyProxy(_load_file_migration_spec(source_file), self.config)
            class_definition = strategy.everything_except_method_declarations

            self.assertEqual(norm_class_definition, _normalize_definition(class_definition), f"Source: {source_file}")

    @classmethod
    def setUpClass(cls) -> None:
        _setup_test_config(cls)


class CSharpCompilationUnitToSingleFileWithLLMProxy(CSharpCompilationUnitToSingleFileWithLLM):
    saved_content: str
    saved_path: str

    batch_count: int = 0

    ResultType = TypeVar('ResultType')

    def save_content(self, content: str, target_file_path: str):
        self.saved_content = content
        self.saved_path = target_file_path

    def load_llm(self) -> LLM:
        class LocalLLM(LLM):
            _parent: Any

            def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
                self._parent.batch_count += 1
                return f"TRANSLATED {user} {self._parent.batch_count}"

            def fits_in_one_prompt(self, token_count: int) -> bool:
                return token_count < 1_000

            def count_tokens(self, source_text: str) -> int:
                return len(source_text)

            def initialize(self) -> None:
                pass

        local_llm = LocalLLM(self.config)
        local_llm._parent = self

        return local_llm


class TestCSharpCompilationUnitToSingleFileWithLLM(unittest.TestCase):
    config: Dict

    def test_execute_small_file(self):
        strategy = CSharpCompilationUnitToSingleFileWithLLMProxy(_load_file_migration_spec('ShortClassNoNamespace.cs'), self.config)
        strategy.execute()

        self.assertEqual("TRANSLATED full 1", strategy.saved_content)
        self.assertEqual("na/resources/ShortClassNoNamespace.gd", strategy.saved_path)

    def test_execute_large_file(self):
        expected_class = "TRANSLATED class_only 1"

        for ix in range(2, 45):
            expected_class += f"\nTRANSLATED methods_only {ix}\n"

        strategy = CSharpCompilationUnitToSingleFileWithLLMProxy(_load_file_migration_spec('LongClassWithNamespace.cs'), self.config)
        strategy.execute()

        self.assertEqual(_normalize_definition(expected_class), _normalize_definition(strategy.saved_content))
        self.assertEqual("na/resources/LongClassWithNamespace.gd", strategy.saved_path)

    @classmethod
    def setUpClass(cls) -> None:
        _setup_test_config(cls)


def _load_file_migration_spec(source_file_name: str) -> unifree.FileMigrationSpec:
    dir_path = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(dir_path, 'resources', source_file_name)

    if not os.path.exists(file_path):
        raise RuntimeError(f"'{source_file_name}' not found")

    return unifree.FileMigrationSpec(
        source_file_path=file_path,
        source_project_path=dir_path,
        destination_project_path="na"
    )


def _load_source_file(source_file_name: str) -> str:
    file_path = _load_file_migration_spec(source_file_name)

    with open(file_path.source_file_path, 'r') as source_file:
        return source_file.read()


def _normalize_definition(s: str) -> str:
    return re.sub(r'\s+', ' ', s).lstrip().rstrip()


def _setup_test_config(target) -> None:
    target.config = to_default_dict({
        "source": {
            "csharp": {
                "convert_macros_to_comments": True
            }
        },
        "target": {
            "convert_tabs_to_spaces": True,
            "extension": ".gd"
        },
        "prompts": {
            "system": "System Prompt",
            "full": "full",
            "class_only": "class_only",
            "methods_only": "methods_only",
        },
        "llm": {
            "class": "TrivialLLM"
        }
    })


if __name__ == '__main__':
    unittest.main()

tests/test_csharp_migration_strategies.py
ProjectUnifree/unifree
#!/usr/bin/env python3
import unittest

from unifree.known_translations_db import KnownTranslationsDb, KnownTranslation
from unifree.utils import to_default_dict


# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)


class KnownTranslationsDbProxy(KnownTranslationsDb):
    def _upsert_translations_from_yaml(self, language: str) -> None:
        # Do not lookup yaml file, instead upsert manually
        self._upsert_translations([
            KnownTranslation(
                source="Hello, world!",
                target="Sentence One",
            ),
            KnownTranslation(
                source="Sator arepo tenet opera rotas",
                target="Sentence Two",
            ),
        ])


class TestKnownTranslationsDb(unittest.TestCase):
    _db: KnownTranslationsDbProxy

    def test_fetch_nearest(self):
        self._db.initialize()

        result = self._db.fetch_nearest_known_translations("Sator arepo", count=1)
        self.assertEqual(1, len(result))
        self.assertEqual("Sator arepo tenet opera rotas", result[0].source)
        self.assertEqual("Sentence Two", result[0].target)

    def test_fetch_nearest_as_query_history(self):
        self._db.initialize()

        result = self._db.fetch_nearest_as_query_history("Hello")
        self.assertEqual(4, len(result))

        self.assertEqual("user", result[0].role)
        self.assertEqual("'Hello, world!' is 'Sentence One'", result[0].content)

        self.assertEqual("assistant", result[1].role)
        self.assertEqual("!OK!", result[1].content)

        self.assertEqual("user", result[2].role)
        self.assertEqual("'Sator arepo tenet opera rotas' is 'Sentence Two'", result[2].content)

        self.assertEqual("assistant", result[3].role)
        self.assertEqual("!OK!", result[3].content)

    @classmethod
    def setUpClass(cls) -> None:
        cls._db = KnownTranslationsDbProxy(to_default_dict({
            "known_translations": {
                "language": "trivial",
                "embedding_function": "all-MiniLM-L6-v2",
                "result_count": 2,
                "assistant_response": "!OK!",
                "user_request": "'${SOURCE}' is '${TARGET}'",
            }
        }))

tests/test_known_translations_db.py
ProjectUnifree/unifree
#!/usr/bin/env python3
# Copyright (c) AppLovin. and its affiliates. All rights reserved.
import unittest
from concurrent.futures import ThreadPoolExecutor

from unifree.llms import MultiprocessLocalLLM


class TestMultiprocessLocalLLM(unittest.TestCase):

    def test_translate(self):
        config = {
            "class": "MultiprocessLocalLLM",
            "llm_config": {
                "class": "TrivialLLM",
                "config": {
                    "key": "value"
                }
            },

            "wrapper_config": {
                "num_workers": 5,
                "query_timeout_sec": 1,
            }
        }

        mp_llm = MultiprocessLocalLLM(config)
        mp_llm.initialize()

        count = mp_llm.count_tokens("some text")
        self.assertEqual(9, count)

        self.assertTrue(mp_llm.fits_in_one_prompt(123))

        def map_func(query_ix_loc):
            nonlocal mp_llm
            return query_ix_loc, mp_llm.query(f"QUERY {query_ix_loc}")

        inputs = range(200)
        with ThreadPoolExecutor(max_workers=4) as e:
            results = e.map(map_func, inputs)

        results_count = 0
        for query_ix, response in results:
            self.assertEqual(f"QUERY {query_ix}", response)
            results_count += 1

        self.assertEqual(200, results_count)

tests/test_multiprocess_local_llm.py
ProjectUnifree/unifree
#!/usr/bin/env python3
import pathlib
# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Dict, List

# =====================
# Overall Configuration
# =====================

log_level: Optional[str] = 'info'
"""Log level to use. Options are 'debug', 'info', 'warn' and 'error'"""

supress_warnings: bool = False
"""If true 'warn'-level messages would be logged into stderr"""

project_root: str = str(pathlib.Path(__file__).parent.parent.resolve())
"""This is the _root_ folder of the project. All resource files (config, known translations, etc.) will be relative to this folder"""


class MigrationStrategy(ABC):
    _config: Dict

    def __init__(self, config: Dict) -> None:
        self._config = config

    @property
    def config(self):
        return self._config

    def execute(self) -> None:
        raise NotImplementedError


class DummyMigrationStrategy(MigrationStrategy):
    def execute(self) -> None:
        pass


@dataclass
class FileMigrationSpec:
    source_file_path: str
    source_project_path: str
    destination_project_path: str


@dataclass
class QueryHistoryItem:
    role: str
    content: str


class LLM(ABC):
    _config: Dict

    def __init__(self, config: Dict) -> None:
        self._config = config

    @abstractmethod
    def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
        """
        Query LLM with the provided input.

        :param user:     The actual query.
        :param system:   "System" query. Some LLMs support instructions that should be respected by the LLM and treated with higher regard
                         then the user query
        :param history:  User conversation history. This could be helpful to convey knowledge about language/system.
        :return: Response from the LLM
        """
        raise NotImplementedError

    def initialize(self) -> None:
        """
        Initialize the LLM
        """
        raise NotImplementedError

    def fits_in_one_prompt(self, token_count: int) -> bool:
        """
        Check if the given token count fits into one prompt
        :param token_count: Token count to check
        :return: True if this token count fits into one prompt
        """
        raise NotImplementedError

    def count_tokens(self, source_text: str) -> int:
        """
        Count number of tokens for the given source text
        :param source_text: Text to count tokens in
        :return: Number of tokens
        """
        raise NotImplementedError

    @property
    def config(self) -> Dict:
        return self._config

unifree/__init__.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import os
from abc import ABC
from typing import Dict, Optional, Callable, TypeVar, List

import tree_sitter

from unifree import log, MigrationStrategy, FileMigrationSpec, utils, LLM, QueryHistoryItem
from unifree.llms.code_extrators import extract_first_source_code, extract_header_implementation
from unifree.source_code_parsers import CSharpCodeParser
from unifree.utils import load_llm


class CSharpCompilationUnitMigrationStrategy(MigrationStrategy, ABC):
    """
    Strategy to migrate a compilation unit
    """

    _tree: Optional[tree_sitter.Tree]
    _file_migration_spec: FileMigrationSpec

    def __init__(self, file_migration_spec: FileMigrationSpec, config: Dict) -> None:
        super().__init__(config)
        self._file_migration_spec = file_migration_spec
        self._tree = None

    @property
    def source_file_path(self) -> str:
        return self._file_migration_spec.source_file_path

    @property
    def source_project_path(self) -> str:
        return self._file_migration_spec.source_project_path

    @property
    def destination_project_path(self) -> str:
        return self._file_migration_spec.destination_project_path

    @property
    def tree(self) -> tree_sitter.Tree:
        if self._tree is None:
            self._tree = CSharpCodeParser(self.config).parse(self.source_file_path)

        return self._tree

    @property
    def source_text(self) -> str:
        return self.tree.text.decode('utf-8')

    @property
    def everything_except_method_declarations(self) -> str:
        recurse_into_types = ['compilation_unit', 'namespace_declaration', 'class_declaration', 'declaration_list']
        ignore_types = ['method_declaration']
        result = ''

        def everything_except_method_declarations(node):
            nonlocal recurse_into_types
            nonlocal result

            if node.type in recurse_into_types:
                for child in node.children:
                    everything_except_method_declarations(child)

            elif node.type in ignore_types:
                pass
            elif node.type == "{" or node.type == "}":
                result += f"\n{node.type}\n"
            else:
                text = node.text.decode('utf-8')
                text += "\n" if text.endswith(";") else " "

                result += text

        everything_except_method_declarations(self.tree.root_node)

        return result

    @property
    def method_declarations(self) -> [str]:
        result = []

        def extract_methods(node):
            nonlocal result
            if node.type == "method_declaration":
                result.append(node.text.decode('utf-8'))

            for child in node.children:
                extract_methods(child)

        extract_methods(self.tree.root_node)

        return result

    def create_destination_file_path(self, extension: str) -> str:
        relative_path = os.path.relpath(self.source_file_path, self.source_project_path)
        relative_folder_path = os.path.dirname(relative_path)

        if self.config["target"]["lower_folder_names"]:
            relative_folder_path = relative_folder_path.lower()

        target_folder = os.path.join(self.destination_project_path, relative_folder_path)
        try:
            os.makedirs(target_folder, exist_ok=True)
        except Exception:
            pass  # Ignore

        source_filename = os.path.basename(relative_path)
        file_name, file_extension = os.path.splitext(os.path.basename(source_filename))

        if self.config["target"]["convert_filename_to_camelcase"]:
            file_name = utils.snake_to_camel(file_name)
        elif self.config["target"]["convert_filename_to_snake_case"]:
            file_name = utils.camel_to_snake(file_name)

        return os.path.join(target_folder, file_name + extension)

    def maybe_convert_tabs_and_spaces(self, source: str) -> str:
        result = source

        if self.config["target"]["convert_tabs_to_spaces"]:
            result = self.replace_tabs_with_spaces(result)

        if self.config["target"]["convert_spaces_to_tabs"]:
            result = self.replace_spaces_with_tabs(result)

        return result

    def save_content(self, content: str, target_file_path: str):
        if len(target_file_path) > 0:
            log.debug(f"Saving {len(content):,} bytes to '{target_file_path}'...")

            with open(target_file_path, 'w') as output_file:
                output_file.write(content)

    @staticmethod
    def replace_tabs_with_spaces(s: str) -> str:
        return s.replace("\t", "    ")

    @staticmethod
    def replace_spaces_with_tabs(s: str) -> str:
        return s.replace("    ", "\t")


class CSharpCompilationUnitMigrationWithLLM(CSharpCompilationUnitMigrationStrategy, ABC):
    _llm: LLM

    def __init__(self, file_migration_spec: FileMigrationSpec, config: Dict) -> None:
        super().__init__(file_migration_spec, config)

        if "llm" not in config:
            raise RuntimeError(f"No 'llm' key found in the tool configuration")

        self._llm = self.load_llm()
        self._llm.initialize()

    @property
    def llm(self) -> LLM:
        return self._llm

    ResultType = TypeVar('ResultType')

    def translate_code(self, code: str, prompt_type: str, system: str, extractor_fn: Callable[[str], ResultType]) -> ResultType:
        user = self.create_code_prompt(prompt_type, code)
        history = self.load_translation_history(code)
        response = self.llm.query(user, system, history)
        return extractor_fn(response)

    def create_code_prompt(self, prompt_type: str, code: str) -> str:
        return self.create_prompt(prompt_type, {"CODE": code})

    def create_prompt(self, prompt_type: str, replacements: Dict[str, str]) -> str:
        text: str = self.config['prompts'][prompt_type]

        if not text:
            raise RuntimeError(f"Prompt '{prompt_type}' is not defined. Please confirm it is set in the destination YAML file under 'prompts'")

        for key, value in replacements.items():
            text = text.replace('${' + key + '}', value)

        return text

    def load_translation_history(self, code: str) -> List[QueryHistoryItem]:
        """
        Load translation history that is relevant to the given bit of translated code.

        :param code:    Code that will be translated

        :return: Documents relevant to the translation history
        """
        from unifree.known_translations_db import KnownTranslationsDb
        if KnownTranslationsDb.is_instance_initialized():
            return KnownTranslationsDb.instance().fetch_nearest_as_query_history(
                query=code
            )
        else:
            return []

    def load_llm(self) -> LLM:
        """
        Load target LLM.

        Please note: this method is separate for unit tests to be able to override it

        :return: Newly loaded LLM
        """
        return load_llm(self.config["llm"])


class CSharpCompilationUnitToSingleFileWithLLM(CSharpCompilationUnitMigrationWithLLM):
    def __init__(self, file_migration_spec: FileMigrationSpec, destination: Dict) -> None:
        super().__init__(file_migration_spec, destination)

    def execute(self) -> None:
        system = self.config['prompts']['system']

        # LLMs is sometimes not very good at handling large input source code. So if a code is
        # beyond a certain threshold, translate each method individually
        token_count = self.llm.count_tokens(self.source_text)
        if self.llm.fits_in_one_prompt(token_count):
            response = self.translate_code(self.source_text, 'full', system, extract_first_source_code)
        else:
            translated_class_only = self.translate_code(self.everything_except_method_declarations, 'class_only', system, extract_first_source_code)

            current_methods = ''
            translated_methods = ''
            for method_declaration in self.method_declarations:
                updated_current_methods = current_methods + "\n\n" + method_declaration

                token_count = self.llm.count_tokens(updated_current_methods)
                if self.llm.fits_in_one_prompt(token_count):
                    current_methods = updated_current_methods
                else:
                    translated_methods += "\n\n" + self.translate_code(current_methods, 'methods_only', system, extract_first_source_code)
                    current_methods = method_declaration

            if len(current_methods) > 0:
                translated_methods += "\n\n" + self.translate_code(current_methods, 'methods_only', system, extract_first_source_code)

            if "${METHODS}" in translated_class_only:
                response = translated_class_only.replace("${METHODS}", translated_methods)
            else:
                response = translated_class_only + translated_methods

        response = self.maybe_convert_tabs_and_spaces(response)

        output_file_name = self.create_destination_file_path(self.config["target"]["extension"])
        self.save_content(content=response, target_file_path=output_file_name)

    def __str__(self) -> str:
        output_file_name = self.create_destination_file_path(self.config["target"]["extension"])
        return f"[Migrate '{self.source_file_path}' to '{output_file_name}']"


class CSharpCompilationUnitToInterfaceImplementationWithLLM(CSharpCompilationUnitMigrationWithLLM):
    def execute(self) -> None:
        system = self.config['prompts']['system']

        # Chat GPT is sometimes not very good at handling large input source code. So if a code is
        # beyond a certain threshold, translate each method individually
        token_count = self.llm.count_tokens(self.source_text)
        if self.llm.fits_in_one_prompt(token_count):
            header, implementation = self.translate_code(self.source_text, 'full', system, extract_header_implementation)
        else:
            class_header, class_implementation = self.translate_code(self.everything_except_method_declarations, 'class_only', system, extract_header_implementation)

            current_methods = ''
            method_headers, method_implementations = '', ''

            def translate_current_methods_and_append():
                nonlocal current_methods, method_headers, method_implementations
                translated_header, translated_implementation = self.translate_code(current_methods, 'methods_only', system, extract_header_implementation)
                method_headers += "\n\n" + translated_header
                method_implementations += "\n\n" + translated_implementation

            for method_declaration in self.method_declarations:
                updated_current_methods = current_methods + "\n\n" + method_declaration

                token_count = self.llm.count_tokens(updated_current_methods)
                if self.llm.fits_in_one_prompt(token_count):
                    current_methods = updated_current_methods
                else:
                    translate_current_methods_and_append()
                    current_methods = method_declaration

            if len(current_methods) > 0:
                translate_current_methods_and_append()

            if "${METHODS}" in class_header:
                header = class_header.replace("${METHODS}", method_headers)
            else:
                header = class_header + method_headers

            implementation = class_implementation + "\n\n" + method_implementations

        header = self.maybe_convert_tabs_and_spaces(header)
        implementation = self.maybe_convert_tabs_and_spaces(implementation)

        self.save_content(content=header, target_file_path=self.create_destination_file_path(self.config["target"]["header_extension"]))
        self.save_content(content=implementation, target_file_path=self.create_destination_file_path(self.config["target"]["implementation_extension"]))

    def __str__(self) -> str:
        output_file_name = self.create_destination_file_path(self.config["target"]["header_extension"])
        implementation_ext = self.config["target"]["implementation_extension"]
        return f"[Migrate '{self.source_file_path}' to '{output_file_name}/{implementation_ext}']"

unifree/csharp_migration_strategies.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import argparse
import os.path
import platform
import sys
from typing import Optional

import unifree
from unifree import utils, log


def run_migration(
        config: str,
        source: str,
        destination: str,
        llm_secret_key: Optional[str] = None,
        verbose: bool = False,
):
    if verbose:
        unifree.log_level = 'debug'

    try:
        config = utils.load_config(config)
        config["verbose"] = verbose

        if llm_secret_key and "llm" in config and "config" in config["llm"]:
            config["llm"]["config"]["secret_key"] = llm_secret_key

    except Exception as e:
        log.error(f"Unable to start: {config} is invalid: {e}", exc_info=e)
        return 78 # os.EX_CONFIG

    if not os.path.exists(source):
        log.error(f"Unable to start: source folder does not exist ('{source}')")
        return 78 # os.EX_CONFIG

    if not os.path.isdir(source):
        log.error(f"Unable to start: source folder is not a folder ('{source}')")
        return 78 # os.EX_CONFIG

    if os.path.exists(destination):
        log.info(f"Using existing destination folder '{destination}'")
    else:
        log.info(f"Creating destination folder '{destination}'...")

        try:
            os.makedirs(destination)
        except Exception as e:
            log.error(f"Unable to create destination folder '{destination}': {e}", exc_info=e)
            return 74 # os.EX_IOERR

    from unifree.project_migration_strategies import CreateMigrations, ExecuteMigrations

    try:
        create_migrations = CreateMigrations(source, destination, config)
        create_migrations.execute()

        execute_migrations = ExecuteMigrations(create_migrations.migrations(), config)
        execute_migrations.execute()

        log.info("Migration completed successfully")

        return os.EX_OK
    except Exception as e:
        log.error(f"Unable to migrate: {e}", exc_info=e)
        return 70 # os.EX_SOFTWARE


def migrate():
    system_platform = platform.system()
    args_parser = argparse.ArgumentParser(
        description="Run a migration of a Unity project",
        usage=f"""\npython3 unifree/free.py -c Config_Name -k ChatGPT_Key -s Source_Location  -d Destination_Location [-v]
            \nExample call: python3 unifree/free.py -c godot_with_gds -k sk-5X8...3L2a -s /Users/john/Unity/FlappyBird -s /Users/john/Godot/FlappyBird
        """
    )
    args_parser.add_argument(
        '--config', '-c',
        required=True,
        type=str,
        help=f"Name of the migration configuration. Must match file name in the ./configs folder",
    )
    args_parser.add_argument(
        '--source', '-s',
        required=True,
        type=str,
        help=f"Path to the source Unity project",
    )
    args_parser.add_argument(
        '--destination', '-d',
        required=True,
        type=str,
        help=f"Path to the destination project folder. Migrated files would be stored there",
    )
    args_parser.add_argument(
        '--llm_secret_key', '-k',
        required=False,
        type=str,
        help=f"Secret key for the current LLM (optional)",
    )
    args_parser.add_argument(
        '--verbose', '-v',
        required=False,
        default=False,
        action='store_true',
        help=f"Print verbose information about the migration process")

    try:
        args, _ = args_parser.parse_known_args()
    except SystemExit:
        args_parser.print_usage()
        sys.exit(78) # os.EX_CONFIG

    args = vars(args)

    status_code = run_migration(
        **args
    )
    sys.exit(status_code)


if __name__ == '__main__':
    migrate()

unifree/free.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from __future__ import annotations

import os
from typing import Dict, Optional, List, Iterable

import chromadb
import yaml
from attr import dataclass
from sentence_transformers import SentenceTransformer

import unifree
from unifree import log


@dataclass
class KnownTranslation:
    source: str
    target: str

    @classmethod
    def from_dict(cls, input_dict: Dict) -> KnownTranslation:
        return KnownTranslation(
            source=input_dict['source'],
            target=input_dict['target'],
        )


class KnownTranslationsDb:
    _class_instance: Optional[KnownTranslationsDb] = None

    _config: Dict

    _chroma_client: Optional[chromadb.Client]
    _collection: Optional[chromadb.Collection]
    _sentence_transformer: Optional[SentenceTransformer]
    _default_n_results: Optional[int]

    def __init__(self, config: Dict) -> None:
        self._config = config
        self._collection = None
        self._sentence_transformer = None
        self._default_n_results = None

    def fetch_nearest_as_query_history(self, query: str, count: Optional[int] = None) -> List[unifree.QueryHistoryItem]:
        result = []
        known_translations = self.fetch_nearest_known_translations(query, count)
        if len(known_translations) > 0:
            translations_config = self._config["known_translations"]
            assistant_response = translations_config["assistant_response"]
            user_request_template = translations_config["user_request"]

            for known_translation in known_translations:
                user_request = user_request_template \
                    .replace("${SOURCE}", known_translation.source) \
                    .replace("${TARGET}", known_translation.target)

                result.append(unifree.QueryHistoryItem(
                    role="user",
                    content=user_request
                ))
                result.append(unifree.QueryHistoryItem(
                    role="assistant",
                    content=assistant_response
                ))

        return result

    def fetch_nearest_known_translations(self, query: str, count: Optional[int] = None) -> List[KnownTranslation]:
        if self._collection:
            if not count:
                count = self._default_n_results

            results = self._collection.query(
                query_embeddings=[self._sentence_transformer.encode(query).tolist()],
                n_results=count,
            )

            documents = results['documents']
            metadatas = results['metadatas']

            if len(metadatas) > 0 and len(documents) > 0:
                return [KnownTranslation(
                    source=document,
                    target=metadata['target'],

                ) for document, metadata in zip(documents[0], metadatas[0])]

        log.debug("Known translations DB was not initialized, not returning anything")
        return []

    def initialize(self) -> None:
        if self._collection is not None:
            return  # Already initialized

        if self._config["known_translations"]:
            translations_config = self._config["known_translations"]
            target_engine = translations_config["target_engine"]

            self._default_n_results = translations_config["result_count"]

            log.debug("Loading sentence transformer...")
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
            self._sentence_transformer = self._create_sentence_transformer(translations_config)

            log.debug("Creating and filling up sentence DB...")
            self._chroma_client = chromadb.Client()
            self._collection = self._chroma_client.create_collection(
                name=f"{target_engine}-known-translations",
                embedding_function=self._sentence_transformer
            )
            self._upsert_translations_from_yaml(target_engine)

    def _create_sentence_transformer(self, translations_config: Dict) -> SentenceTransformer:
        embedding_function_name = translations_config["embedding_function"]

        return SentenceTransformer(embedding_function_name)

    def _upsert_translations_from_yaml(self, target_engine: str) -> None:
        known_translations_file_path = os.path.join(unifree.project_root, "known_translations", f"{target_engine}.yaml")
        if not os.path.exists(known_translations_file_path) or not os.path.isfile(known_translations_file_path):
            log.warn(f"No requested known translations found at '{known_translations_file_path}'")
            return

        with open(known_translations_file_path, 'r') as known_translations_file:
            known_translations = yaml.safe_load(known_translations_file)

        if "translations" in known_translations:
            translations = map(KnownTranslation.from_dict, known_translations['translations'])
            self._upsert_translations(translations)

        else:
            log.warn(f"'{known_translations_file_path}' is malformed: no root node called 'translations' found")

    def _upsert_translations(self, translations: Iterable[KnownTranslation]) -> None:
        documents = []
        metadatas = []
        ids = []

        for translation in translations:
            documents.append(translation.source)
            metadatas.append({
                'target': translation.target,
            })
            ids.append(str(len(ids) + 1))

        embeddings = self._sentence_transformer.encode(documents).tolist()

        self._collection.upsert(
            documents=documents,
            embeddings=embeddings,
            metadatas=metadatas,
            ids=ids,
        )
        log.debug(f"Inserted {len(ids):,} known translations")

    @classmethod
    def instance(cls) -> KnownTranslationsDb:
        if not cls.is_instance_initialized():
            raise RuntimeError(f"Known translations DB is not initialized")

        return cls._class_instance

    @classmethod
    def is_instance_initialized(cls) -> bool:
        return cls._class_instance is not None

    @classmethod
    def initialize_instance(cls, config: Dict) -> None:
        cls._class_instance = KnownTranslationsDb(config)
        cls._class_instance.initialize()

unifree/known_translations_db.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)


from .chatgpt_llm import ChatGptLLM
from .huggingface_llm import HuggingfaceLLM

from .trivial_llm import TrivialLLM
from .multiprocess_local_llm import MultiprocessLocalLLM

unifree/llms/__init__.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from typing import Dict, Optional, List

import openai
import tiktoken

from unifree import LLM, log, QueryHistoryItem


class ChatGptLLM(LLM):
    """
    This class adds a capability to query chat GPT for a completion.

    The configuration looks like:

    ```
    llm:
      class: ChatGptLLM
      config:
          secret_key: <INSERTED DYNAMICALLY BY free.py>
          model: gpt-4
          max_tokens: 4000
    ```

    """

    def __init__(self, config: Dict) -> None:
        super().__init__(config)

    def initialize(self) -> None:
        pass

    def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
        openai.api_key = self.config["config"]['secret_key']

        if log.is_debug():
            short_user_query = user[:50].replace("\n", " ")
            token_count = self.count_tokens(user)
            log.debug(f"Requesting running a query for {token_count:,} tokens ('{short_user_query}...')...")

        try:
            messages = []

            if system:
                messages.append({
                    "role": "system",
                    "content": system
                })

            if history:
                for history_item in history:
                    messages.append({
                        "role": history_item.role,
                        "content": history_item.content
                    })

            messages.append({
                "role": "user",
                "content": user
            })

            completion = openai.ChatCompletion.create(
                model=self.config["config"]["model"],
                messages=messages
            )

            if len(completion.choices) < 1 or len(completion.choices[0].message.content) < 1:
                raise RuntimeError(f"ChatGPT returned malformed response: {completion}")

            response = completion.choices[0].message.content

            if log.is_debug():
                messages_str = [f"> {m['role']}: {m['content']}" for m in messages]
                messages_str = "\n\n".join(messages_str)

                log.debug(f"\n==== GPT REQUEST ====\n{messages_str}\n\n==== GPT RESPONSE ====\n{response}\n")

            return response
        except Exception as e:
            raise RuntimeError(f"ChatGPT query failed: {e}")

    def fits_in_one_prompt(self, token_count: int) -> bool:
        return token_count < self.config["config"]["max_tokens"]

    def count_tokens(self, source_text: str) -> int:
        encoding = tiktoken.encoding_for_model(self.config["config"]["model"])
        num_tokens = len(encoding.encode(source_text))

        return num_tokens

unifree/llms/chatgpt_llm.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from typing import Tuple


def extract_first_source_code(response: str, code_delimiter: str = "```") -> str:
    """
    Parse response and find the first occurrence of text separated with `code_delimiter`
    :param response:        Response from the model
    :param code_delimiter:  Delimiter that is used to separate the code. Default is ````
    :return: String with the first occurrence or whole response if no delimiters found
    """
    result = ''

    starting_ix = response.find(code_delimiter)
    if starting_ix >= 0:
        starting_ix = response.find("\n", starting_ix)  # Move to the new line
        ending_ix = response.find(code_delimiter, starting_ix)
        if ending_ix > 0 and starting_ix > 0:
            result = response[starting_ix:ending_ix]

    if len(result) == 0:
        result = response

    return result.rstrip().lstrip()


def extract_header_implementation(response: str, header_ext: str = '.h', implementation_ext='.cpp', code_delimiter: str = "```") -> Tuple[str, str]:
    """
    This is a hacky attempt to parse ChatGPT response that contain both header an implementation. Sometimes the result
    would be something like:
        Player.h:

        ```cpp
                ... header code ...
        ```

        Player.cpp:

        ```cpp
                ... implementation code ...
        ```

    But other times it is:
        // Player.h
                ... header code ...

        // Player.cpp:
                ... implementation code ...

    :param response:            Full ChatGPT response
    :param header_ext:          Extension of the expected header file (default '.h')
    :param implementation_ext:  Extension of the implementation file (default '.cpp')
    :param code_delimiter:  Delimiter that is used to separate the code. Default is ````

    :return: Best attempt at parsed header/implementation
    """
    header = ''
    is_in_header = False

    implementation = ''
    is_in_implementation = False

    # True in case header and implementation are separated by '// Filename.h' comment
    is_comment_based_delimiter = False

    for line in response.splitlines():
        if line.lstrip().startswith(code_delimiter):
            if is_comment_based_delimiter:
                continue

            is_opening_delimiter = not is_in_header and not is_in_implementation
            if is_opening_delimiter:
                is_header_present = len(header) > 0
                if is_header_present:
                    is_in_implementation = True
                    is_in_header = False
                else:
                    is_in_header = True
                    is_in_implementation = False
            else:  # Closing delimiter
                if is_in_header:
                    is_in_header = False
                if is_in_implementation:
                    is_in_implementation = False
        elif line.lstrip().startswith("//") and line.rstrip().endswith(header_ext):
            is_in_header = True
            is_in_implementation = False
            is_comment_based_delimiter = True
        elif line.lstrip().startswith("//") and line.rstrip().endswith(implementation_ext):
            is_in_header = False
            is_in_implementation = True
            is_comment_based_delimiter = True
        else:
            if is_in_header:
                header += line + "\n"
            elif is_in_implementation:
                implementation += line + "\n"

    # If we couldn't parse the response, just return everything in the header
    if len(header) < 1 and len(implementation) < 1:
        header = response
        implementation = "SEE HEADER FOR FULL RESPONSE"

    return header.rstrip().lstrip(), implementation.rstrip().lstrip()

unifree/llms/code_extrators.py
ProjectUnifree/unifree
#!/usr/bin/env python3
from typing import Optional, List, Dict

from unifree import LLM, QueryHistoryItem, log
from unifree.utils import get_or_create_global_instance


# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)


class HuggingfaceLLM(LLM):
    """
    This is a model that represents Huggingface 'AutoModelForCausalLM' transformer model.

      llm:
        class: HuggingfaceLLM
        config:
          checkpoint: TheBloke/CodeLlama-34B-Instruct-GGUF
          context_length: 4096
          model_type: llama
          gpu_layers: 50

    """
    _model: Optional[any]

    def __init__(self, config: Dict) -> None:
        super().__init__(config)

        self._model = None

    def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
        prompt = ''
        if system:
            prompt += self._to_user_prompt(f"Remember these rules:\n{system}\n")
            prompt += "\nCertainly, I will remember and follow these rules.\n"

        if history:
            for item in history:
                if item.role == "user":
                    prompt += self._to_user_prompt(f"\n{item.content}\n")
                else:
                    prompt += f"\n{item.content}\n"

        prompt += self._to_user_prompt(f"\n{user}\n")

        log.debug(f"\n==== LLM REQUEST ====\n{prompt}\n")

        response = self._model(prompt)

        log.debug(f"\n==== LLM RESPONSE ====\n{response}\n")

        return response

    def initialize(self) -> None:
        from ctransformers import AutoModelForCausalLM

        llm_config = self.config["config"]
        checkpoint = llm_config["checkpoint"]

        self._model = get_or_create_global_instance(checkpoint, lambda: AutoModelForCausalLM.from_pretrained(
            checkpoint,
            model_type=llm_config["model_type"],
            gpu_layers=llm_config["gpu_layers"],
            context_length=llm_config["context_length"],
        ))

    def fits_in_one_prompt(self, token_count: int) -> bool:
        return token_count < self.config["config"]["context_length"]

    def count_tokens(self, source_text: str) -> int:
        return len(self._model.tokenize(source_text))

    def _to_user_prompt(self, user: str) -> str:
        prompt_template = self.config["config"]["prompt_template"]
        return prompt_template.replace("${PROMPT}", user)

unifree/llms/huggingface_llm.py
ProjectUnifree/unifree
#!/usr/bin/env python3
import sys
# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from typing import Optional, List, Dict

from unifree import LLM, QueryHistoryItem, log
from unifree.utils import load_llm


class MultiprocessLocalLLM(LLM):
    """
    This class wraps an LLM implementation and run inference (aka 'query' method) in separate processes on the local box.

    The configuration would look like:

    ```
    llm:
      class: MultiprocessLocalLLM
      llm_config:
          class: <wrapped LLM class>
          config <wrapped LLM config>

        wrapper_config:
          num_workers: 5
          query_timeout_sec: 10
    ```
    """
    _shared_executor: Optional[ProcessPoolExecutor] = None
    _query_timeout_sec: int = 20

    _local_model: Optional[LLM]

    def __init__(self, config: Dict) -> None:
        super().__init__(config)

        self._local_model = None

    def initialize(self) -> None:
        llm_config = self.config["llm_config"]

        self._local_model = load_llm(llm_config)
        self._local_model.initialize()

    def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
        self.maybe_initialize_shared_executor(self.config)

        result_future = self._shared_executor.submit(_multi_process_worker_translate, _QueryRequest(
            user=user,
            system=system,
            history=history
        ))
        try:
            result = result_future.result(timeout=self._query_timeout_sec)
            if result.is_success():
                return result.response
            else:
                raise RuntimeError(f"LocalLLM query failed: {result.exception}")
        except Exception as e:
            raise RuntimeError(f"LocalLLM query failed: {e}")

    def fits_in_one_prompt(self, token_count: int) -> bool:
        assert self._local_model is not None
        return self._local_model.fits_in_one_prompt(token_count)

    def count_tokens(self, source_text: str) -> int:
        assert self._local_model is not None
        return self._local_model.count_tokens(source_text)

    @classmethod
    def maybe_initialize_shared_executor(cls, config: Dict):
        if not cls._shared_executor:
            wrapper_config = config["wrapper_config"]
            llm_config = config["llm_config"]
            # cls._shared_executor = ProcessPoolExecutor(
            cls._shared_executor = ProcessPoolExecutor(
                max_workers=wrapper_config["num_workers"],
                initializer=_multi_process_worker_init,
                initargs=(llm_config,)
            )
            cls._query_timeout_sec = wrapper_config["query_timeout_sec"]


@dataclass
class _QueryRequest:
    """
    This class represents a request ot make a query to the LLM.
    """
    user: str
    system: Optional[str] = None
    history: Optional[List[QueryHistoryItem]] = None


@dataclass
class _QueryResult:
    """
    This object represents response to a given query
    """
    response: str
    exception: Optional[Exception] = None

    def is_success(self) -> bool:
        return self.exception is None


_process_local_llm: Optional[LLM] = None


def _multi_process_worker_init(config: Dict):
    global _process_local_llm

    try:
        _process_local_llm = load_llm(config)
        _process_local_llm.initialize()
    except Exception as e:
        log.error(f"Failed to load local LLM model with: {e}", exc_info=e)
        _process_local_llm = None


def _multi_process_worker_translate(query: _QueryRequest) -> _QueryResult:
    global _process_local_llm

    try:
        return _QueryResult(
            response=_process_local_llm.query(
                user=query.user,
                system=query.system,
                history=query.history,
            )
        )
    except Exception as e:
        import traceback
        traceback.print_exc(file=sys.stderr)

        log.error(f"Failed to query local LLM: {e}")
        return _QueryResult(
            response='',
            exception=e
        )
    except:
        log.error(f"Failed to query local LLM with unknown exception")
        return _QueryResult(
            response='',
            exception=Exception("Unknown exception")
        )

unifree/llms/multiprocess_local_llm.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

from typing import Optional, List, Dict

from unifree import LLM, QueryHistoryItem


class TrivialLLM(LLM):
    """
    This is a trivial implementation of an LLM used for unit tests and an example for other implementations.

    The configuration looks like:

    ```
    llm:
      class: TrivialLLM
      config:
    ```
    """

    def __init__(self, config: Dict) -> None:
        super().__init__(config)

    def query(self, user: str, system: Optional[str] = None, history: Optional[List[QueryHistoryItem]] = None) -> str:
        return user

    def fits_in_one_prompt(self, token_count: int) -> bool:
        return True

    def count_tokens(self, source_text: str) -> int:
        return len(source_text)

    def initialize(self) -> None:
        pass

unifree/llms/trivial_llm.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import logging
import sys
import traceback

import unifree

py_logger = logging.getLogger("lion")


def debug(message: str) -> None:
    if unifree.log_level == "debug":
        print(message, file=sys.stdout)

    py_logger.debug(message)


def info(message: str) -> None:
    if unifree.log_level in ["debug", "info"]:
        print(message, file=sys.stdout)

    py_logger.info(message)


def warn(message: str, exc_info=False) -> None:
    if not unifree.supress_warnings:
        print(message, file=sys.stderr)
        if exc_info:
            traceback.print_exc(file=sys.stderr)

    py_logger.warning(message, exc_info=exc_info)


def error(message: str, exc_info=False) -> None:
    print(message, file=sys.stderr)
    if exc_info:
        traceback.print_exc(file=sys.stderr)

    py_logger.error(message, exc_info=exc_info)


def is_debug():
    return unifree.log_level == 'debug'

unifree/log.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)


from .csharp_migration_strategies import \
    CSharpCompilationUnitToSingleFileWithLLM, \
    CSharpCompilationUnitToInterfaceImplementationWithLLM

unifree/migration_strategies.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import os.path
import traceback
from abc import ABC
from typing import List, Union, Dict, Optional, Iterable

from tqdm.contrib.concurrent import thread_map

from unifree import log, MigrationStrategy, utils, FileMigrationSpec


class ConcurrentMigrationStrategy(MigrationStrategy, ABC):

    def __init__(self, config: Dict) -> None:
        super().__init__(config)

    def map_concurrently(self, fn, iterables, **tqdm_kwargs) -> Iterable:
        return thread_map(fn, iterables, **tqdm_kwargs)


class CreateMigrations(ConcurrentMigrationStrategy):
    _source_path: str
    _destination_path: str

    _migrations: List[MigrationStrategy]
    _errors: List[str]

    def __init__(
            self,
            source_path: str,
            destination_path: str,
            config: Dict,
    ) -> None:
        super().__init__(config)

        self._source_path = source_path
        self._destination_path = destination_path

        self._migrations = []
        self._errors = []

        if not os.path.exists(source_path) or not os.path.isdir(source_path):
            log.error(f"Invalid source path: {source_path}")
            raise RuntimeError("Source path invalid. Please select a folder with Unity project")

        if not os.path.exists(destination_path):
            os.makedirs(destination_path, exist_ok=True)

        self._initialize_shared_objects()

    def migrations(self) -> List[MigrationStrategy]:
        return self._migrations

    def execute(self) -> None:
        log.info(f"Loading source files from '{self._source_path}'...")

        self._check_if_source_is_unity_project()

        project_files = self._load_source_file_paths()

        log.info(f"Computing migration strategies for {len(project_files):,} files...")
        results = self.map_concurrently(
            self._map_file_path_to_migration,
            project_files,
            max_workers=self.config["concurrency"]["create_strategy_worker"] if self.config["concurrency"]["create_strategy_worker"] else 1,
            unit='path',
            chunksize=1,
        )

        warnings = []
        for result in results:
            if isinstance(result, str):
                warnings.append(result)
            elif result is None:
                pass
            elif isinstance(result, MigrationStrategy):
                self._migrations.append(result)
            else:
                raise RuntimeError(f"Programming error: unknown migration {result}")

        for warning in warnings:
            log.warn(warning)

    def _load_source_file_paths(self) -> List[str]:
        absolute_paths = []

        ignored_locations = self.config["source"]["ignore_locations"]
        ignored_locations = [os.path.join(self._source_path, il) for il in ignored_locations] if ignored_locations else []

        # Walk through the folder and its sub-folders
        for root, _, files in os.walk(self._source_path):
            for file in files:
                should_include = True
                for ignored_location in ignored_locations:
                    if root.startswith(ignored_location):
                        should_include = False
                        break

                if should_include:
                    file_path = os.path.join(root, file)
                    absolute_paths.append(file_path)

        return absolute_paths

    def _map_file_path_to_migration(self, file_path: str) -> Union[MigrationStrategy, str, None]:
        try:
            if os.path.isfile(file_path):
                _, file_extension = os.path.splitext(file_path)
                strategy_name = self.config["strategies"][file_extension]
                if strategy_name:
                    spec = FileMigrationSpec(
                        source_file_path=file_path,
                        source_project_path=self._source_path,
                        destination_project_path=self._destination_path
                    )

                    return self._create_migration_strategy(strategy_name, spec)
            return None
        except Exception as e:
            traceback.print_exc()
            return f"'{file_path}' failed to create strategy: {e}"

    def _create_migration_strategy(self, strategy_name: str, spec: FileMigrationSpec) -> MigrationStrategy:
        loaded_class = utils.load_class(strategy_name, 'migration_strategies')
        result = loaded_class(spec, self.config)
        if not isinstance(result, MigrationStrategy):
            raise RuntimeError(f"Loaded {strategy_name} is not a migration strategy")

        return result

    def _initialize_shared_objects(self):
        from unifree.source_code_parsers import CSharpCodeParser
        CSharpCodeParser.initialize()

        from unifree.known_translations_db import KnownTranslationsDb
        KnownTranslationsDb.initialize_instance(self.config)

    def _check_if_source_is_unity_project(self):
        files = os.listdir(self._source_path)
        for required_file in ['Assets', 'ProjectSettings']:
            if required_file not in files:
                log.error(f"Unable to find '{required_file}' folder in '{self._source_path}'")
                raise RuntimeError(f"Source is not a valid Unity project")


class ExecuteMigrations(ConcurrentMigrationStrategy):
    _strategies: List[MigrationStrategy]

    def __init__(self, strategies: List[MigrationStrategy], config: Dict) -> None:
        super().__init__(config)

        self._strategies = strategies

    def execute(self) -> None:
        log.info(f"Executing {len(self._strategies):,} migration strategies...")

        results = self.map_concurrently(
            self._execute_strategy, self._strategies,
            max_workers=self.config["concurrency"]["execute_strategy_workers"] if self.config["concurrency"]["execute_strategy_workers"] else 1,
            unit='file',
            chunksize=1,
        )

        warnings = []
        for result in results:
            if isinstance(result, str):
                warnings.append(result)
            elif result is None:
                pass
            else:
                raise RuntimeError(f"Programming error: unknown migration {result}")

        for warning in warnings:
            log.warn(warning)

    def _execute_strategy(self, strategy: MigrationStrategy) -> Optional[str]:
        try:
            strategy.execute()
        except Exception as e:
            return f"Failed to execute {strategy}: {e}"

unifree/project_migration_strategies.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import os.path
from typing import Dict

import tree_sitter

import unifree
from unifree import log


class CSharpCodeParser:
    _parser: tree_sitter.Parser
    _config: Dict

    def __init__(self, config: Dict) -> None:
        self._config = config
        self.initialize()

        self._parser = tree_sitter.Parser()
        self._parser.set_language(tree_sitter.Language(self._c_sharp_library_path(), "c_sharp"))

    @classmethod
    def initialize(cls) -> None:
        c_sharp_library_path = cls._c_sharp_library_path()
        if not os.path.exists(c_sharp_library_path):
            c_sharp_library_folder_path = os.path.dirname(c_sharp_library_path)
            log.debug(f"C# library not found at '{c_sharp_library_path}', building...")

            tree_sitter_csharp_folder_path = os.path.join(unifree.project_root, 'vendor', 'tree-sitter-c-sharp')
            if not os.path.exists(tree_sitter_csharp_folder_path):
                log.error(f"CSharp implementation of the tree sitter is not found at {tree_sitter_csharp_folder_path}")
                raise RuntimeError(f"CSharp tree sitter not found in {tree_sitter_csharp_folder_path}")

            os.makedirs(c_sharp_library_folder_path, exist_ok=True)
            tree_sitter.Language.build_library(c_sharp_library_path, [tree_sitter_csharp_folder_path])

    def parse(self, file_path: str) -> tree_sitter.Tree:
        if not os.path.exists(file_path) or not os.path.isfile(file_path):
            raise RuntimeError(f"File {file_path} does not exist")

        with open(file_path, 'r') as source_file:
            source_str = source_file.read()
            if len(source_str) < 1:
                raise RuntimeError(f"File {file_path} is empty")

            if self._config["source"]["csharp"]["convert_macros_to_comments"]:
                source_str = self._replace_macros_with_comments(source_str)

            source_bytes = bytes(source_str, "utf8")

            try:

                result: tree_sitter.Tree = self._parser.parse(source_bytes)
                if result.root_node:
                    return result
                else:
                    raise RuntimeError(f"Failed to parse '{file_path}': no root node found")

            except Exception as ex:
                raise RuntimeError(f"Failed to parse '{file_path}': threw exception while parsing", ex)

    @classmethod
    def _c_sharp_library_path(cls) -> str:
        c_sharp_library_folder_path = os.path.join(unifree.project_root, 'vendor', 'build', 'libraries')
        return os.path.join(c_sharp_library_folder_path, 'c-sharp.so')

    def _replace_macros_with_comments(self, source_str):
        result = ''
        for line in source_str.split("\n"):
            stripped_line = line.rstrip().lstrip()
            if stripped_line.startswith("#"):
                result += "// " + line + "\n"
            else:
                result += line + "\n"

        return result

unifree/source_code_parsers.py
ProjectUnifree/unifree
#!/usr/bin/env python3

# Copyright (c) Unifree
# This code is licensed under MIT license (see LICENSE.txt for details)

import importlib
import os
import re
import threading
from collections import defaultdict
from typing import Type, Dict, Any, TypeVar, Callable

import yaml

from unifree import LLM


def load_config(config_name: str) -> Dict[str, Any]:
    from unifree import log
    import unifree

    config_folder_path = os.path.join(unifree.project_root, "configs")
    config_file_name = os.path.join(config_folder_path, f"{config_name}.yaml")
    if not os.path.exists(config_file_name) or not os.path.isfile(config_file_name):
        supported_destinations = []
        for config_file_name in os.listdir(config_folder_path):
            if config_file_name.endswith(".yaml"):
                supported_destinations.append(config_file_name.replace(".yaml", ""))

        log.error(f"{config_name} is not a supported config. All supported configs: " + ", ".join(supported_destinations))
        raise RuntimeError("Config name not supported")

    with open(config_file_name, 'r') as config_file:
        result = yaml.safe_load(config_file)
        result = to_default_dict(result)

        log.debug(f"Loaded config from '{config_file_name}'")

    return result


def load_class(class_name: str, module: str) -> Type:
    module = importlib.import_module('.' + module, package='unifree')

    # Load the class dynamically
    loaded_class = getattr(module, class_name)
    if not loaded_class:
        raise RuntimeError(f"Class {class_name} not found")

    return loaded_class


def load_llm(config: Dict) -> LLM:
    llm_class = load_class(config["class"], "llms")
    return llm_class(config)


def camel_to_snake(camel_case_str: str) -> str:
    snake_case_str = re.sub(r'([a-z])([A-Z])', r'\1_\2', camel_case_str)
    return snake_case_str.lower()


def snake_to_camel(snake_case_str: str) -> str:
    words = snake_case_str.split('_')
    camel_case_str = words[0] + ''.join(word.capitalize() for word in words[1:])
    return camel_case_str


def to_default_dict(d):
    if isinstance(d, dict):
        return defaultdict(_return_none, {k: to_default_dict(v) for k, v in d.items()})
    elif isinstance(d, list):
        return [to_default_dict(e) for e in d]
    else:
        return d


def _return_none():
    return None


InstanceType = TypeVar('InstanceType')

_global_instances: Dict[str, Any] = {}
_global_instances_lock: threading.Lock = threading.Lock()


def get_or_create_global_instance(name: str, new_instance_creator: Callable[[], InstanceType]) -> InstanceType:
    global _global_instances
    global _global_instances_lock

    if name not in _global_instances:
        with _global_instances_lock:
            if name not in _global_instances:
                _global_instances[name] = new_instance_creator()

    return _global_instances[name]

unifree/utils.py
ProjectUnifree/unifree