dam_helper/repository/repository.py

import os
from pathlib import Path

from repository import ModelRoles, Model
from utils.env_utils import in_hf

if not in_hf():
    from repository.intel_npu import IntelNpuRepository
    from repository.ollama import OllamaRepository
from repository.ondemand import OndemandRepository
from repository.testing_repo import TestingRepository


def build_repo_from_environment(system_prompt: str):
    implementation = os.getenv("implementation")
    model_name = os.getenv("model_name")

    if implementation:
        return get_repository(implementation, Model(model_name, ModelRoles("system",
                                                                           "user",
                                                                           "assistant")),
                              system_prompt)
    else:
        return None


def get_repository(implementation: str, model: Model, system_msg: str = None, log_to_file: Path = None):
    known_implementations = ["ollama", "intel_npu", "testing", "ondemand"]
    if not implementation or implementation.lower() not in known_implementations:
        raise ValueError(f"Unknown implementation {implementation}. Known implementations: {known_implementations}")
    if not in_hf():
        if "intel_npu" == implementation:
            return IntelNpuRepository(model, system_msg, log_to_file)
        if "ollama" == implementation:
            return OllamaRepository(model, system_msg)
    if "ondemand" == implementation:
        return OndemandRepository(model, system_msg, log_to_file)
    if "testing" == implementation:
        return TestingRepository(prompts_answers=[
            {
                "role": "assistant",
                "content": "What is my full name?\n\nnull\n\nWhat is the nature of the work I need to do?\n\nPest control\n\nIn which community is the work taking place?\n\nJBR\n\nIn which building?\n\nnull\n\nIn which unit/apartment number?\n\nnull\n\nAm I the owner or the tenant?\n\nTenant\n\nIn which date is the work taking place?\n\n12/09/2024\n\nIn which date will the work finish?\n\n12/09/2024\n\nWhat is my contact number?\n\nnull\n\nWhat is the name of the contracting company?\n\nnull\n\nWhat is the contact number of the contracting company?\n\nnull\n\nWhat is the email of the contracting company?\n\nnull"
            },
            {
                "role": "assistant",
                "content": "pest_control"
            }
        ])