reggpt/controller/router.py

"""
 /*************************************************************************
 * 
 * CONFIDENTIAL
 * __________________
 * 
 *  Copyright (2023-2024) AI Labs, IronOne Technologies, LLC
 *  All Rights Reserved
 *
 *  Author  : Theekshana Samaradiwakara
 *  Description :Python Backend API to chat with private data  
 *  CreatedDate : 14/11/2023
 *  LastModifiedDate : 18/03/2024
 *************************************************************************/
 """

import time
import logging
logger = logging.getLogger(__name__)

from reggpt.chains.llmChain import get_router_chain

from reggpt.configs.model import ROUTER_MODEL_TYPE


router_model_type=ROUTER_MODEL_TYPE

router_chain= get_router_chain(router_model_type)

def run_router_chain(query):
    try:
        logger.info(f"run_router_chain : Question: {query}")
        # Get the answer from the chain
        start = time.time()
        chain_type = router_chain.invoke(query)['text']
        end = time.time()

        # log the result
        logger.info(f"Answer (took {round(end - start, 2)} s.) chain_type: {chain_type}")

        return chain_type
    
    except Exception as e:
        logger.exception(e)
        raise e