Spaces:
Runtime error
Runtime error
File size: 3,160 Bytes
0fac726 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import argparse
import logging
import os
import sys
from dotenv import load_dotenv
from knowledgebase import Knowledgebase
from utils.constants import (
ENV_FILE,
ASSISTANT_TYPE_KEY,
AssistantType,
OPENAI_API_TOKEN_KEY,
OPENAI_KNOWLEDGEBASE_KEY,
HUGGINGFACEHUB_API_TOKEN_KEY,
HF_KNOWLEDGEBASE_KEY,
QUERY_TAG,
ANSWER_TAG,
SOURCES_TAG,
EMBEDDING_TYPE_KEY,
APIKeyType,
EmbeddingType,
)
from utils.llm import validate_api_token
# load the .env
load_dotenv(dotenv_path=os.path.join(os.getcwd(), ENV_FILE))
logger = logging.getLogger(__name__)
if __name__ == "__main__":
assistant_type = os.getenv(ASSISTANT_TYPE_KEY, AssistantType.HUGGINGFACE.value)
embedding_type = os.getenv(EMBEDDING_TYPE_KEY, EmbeddingType.HUGGINGFACE.value)
if assistant_type == AssistantType.OPENAI.value:
assistant_type = AssistantType.OPENAI
assistant_api_key = os.environ.get(OPENAI_API_TOKEN_KEY, None)
assistant_api_key_type = APIKeyType.OPENAI
knowledgebase_name = os.environ.get(OPENAI_KNOWLEDGEBASE_KEY, None)
if embedding_type == EmbeddingType.OPENAI.value:
embedding_type = EmbeddingType.OPENAI
embedding_api_key = assistant_api_key
embedding_api_key_type = APIKeyType.OPENAI
else:
embedding_type = EmbeddingType.HUGGINGFACE
embedding_api_key = os.getenv(HUGGINGFACEHUB_API_TOKEN_KEY, None)
embedding_api_key_type = APIKeyType.HUGGINGFACE
else:
assistant_type = AssistantType.HUGGINGFACE
assistant_api_key = os.environ.get(HUGGINGFACEHUB_API_TOKEN_KEY, None)
assistant_api_key_type = APIKeyType.HUGGINGFACE
knowledgebase_name = os.environ.get(HF_KNOWLEDGEBASE_KEY, None)
embedding_type = EmbeddingType.HUGGINGFACE
embedding_api_key = assistant_api_key
embedding_api_key_type = APIKeyType.HUGGINGFACE
logger.info("🗝️ Validating the API tokens...")
assistant_valid, assistant_err = validate_api_token(
api_key_type=assistant_api_key_type, api_key=assistant_api_key
)
if not assistant_valid:
logger.error(assistant_err)
sys.exit(1)
embedding_valid, embedding_err = validate_api_token(
api_key_type=embedding_api_key_type, api_key=embedding_api_key
)
if not embedding_valid:
logger.error(embedding_err)
sys.exit(1)
parser = argparse.ArgumentParser(description="LLM Website QA - CLI")
parser.add_argument(
QUERY_TAG, type=str, help="Question to be asked from the assistant"
)
args = parser.parse_args()
query = args.query
knowledgebase = Knowledgebase(
assistant_type=assistant_type,
embedding_type=embedding_type,
assistant_api_key=assistant_api_key,
embedding_api_key=embedding_api_key,
knowledgebase_name=knowledgebase_name,
)
result, metadata = knowledgebase.query_knowledgebase(query=query)
print(f"\nAnswer: \n{str(result.get(ANSWER_TAG, '').strip())}")
print(f"\nSources: \n{str(result.get(SOURCES_TAG, '').strip())}")
print(f"\nCost: \n{metadata}")
|