model.py

import os
import sys
import json
import time
import yaml
import joblib
import argparse

import jinja2
import anthropic
import pandas as pd
from tqdm import tqdm
from loguru import logger
from openai import OpenAI
from dotenv import load_dotenv
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold

from utils import parse_json_garbage, compose_query

try:
    logger.remove(0)
    logger.add(sys.stderr, level="INFO")
except ValueError:
    pass

load_dotenv()

def llm( provider, model, system_prompt, user_content, delay:int = 0):
    """Invoke LLM service
    Argument
    --------
    provider: str
        openai or anthropic
    model: str
        Model name for the API
    system_prompt: str
        System prompt for the API
    user_content: str
        User prompt for the API
    Return
    ------
    response: str
    """
    if delay:
        time.sleep(delay)

    if provider=='openai':
        client = OpenAI( organization = os.getenv('ORGANIZATION_ID'))
        chat_completion = client.chat.completions.create( 
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": user_content,
                }
            ],
            model = model, 
            response_format = {"type": "json_object"},
            temperature = 0,
            max_tokens = 4096,
            # stream = True
        )
        response = chat_completion.choices[0].message.content

    elif provider=='anthropic':
        client = anthropic.Client(api_key=os.getenv('ANTHROPIC_API_KEY'))
        response = client.messages.create(
            model= model,
            system= system_prompt,
            messages=[
                {"role": "user", "content": user_content} # <-- user prompt
            ],
            max_tokens = 4000
        )
        response = response.content[0].text

    elif provider=='google':
        genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
        model = genai.GenerativeModel(
            model_name = model,  
            system_instruction = system_prompt,
            generation_config={
            "temperature": 0,
            "max_output_tokens": 8192,
            "response_mime_type": "application/json"
        })
        safety_settings = {
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
        }
        messages = []
        # messages.append({
        #     'role':'user',
        #     'parts': [f"System instruction: {system_prompt}"]
        # })
        # response = model.generate_content(messages, safety_settings=safety_settings)
        # try:
        #     messages.append({
        #         'role': 'model',
        #         'parts': [response.text]
        #     })
        # except Exception as e:
        #     logger.error(f"response.candidates -> {response.candidates}")
        #     logger.error(f"error -> {e}")
        #     messages.append({
        #         'role': 'model',
        #         'parts': ["OK. I'm ready to help you."]
        #     })
        messages.append({
            'role': 'user',
            'parts': [user_content]
        })
        try:
            response = model.generate_content(messages, safety_settings=safety_settings, )
            response = response.text
        except Exception as e:
            logger.error(f"Error (will still return response) -> {e}")
            logger.error(f"response.candidates -> {response.candidates}")
            return response
    else:
        raise Exception("Invalid provider")
    
    return response

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument( "-c", "--config", type=str, default='config/config.yml', help="Path to the configuration file")
    parser.add_argument( "-t", "--task", type=str, default='prepare_batch', choices=['extract', 'classify'])
    parser.add_argument( "-i", "--input_path", type=str, default='', )
    parser.add_argument( "-o", "--output_path", type=str, default='', )
    parser.add_argument( "-topn", "--topn", type=int, default=None )
    args = parser.parse_args()
    # classes = ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', ]
    # backup_classes = [ '中式', '西式']

    assert os.path.exists(args.config), f"File not found: {args.config}"
    config = yaml.safe_load(open(args.config, "r").read())

    if args.task == 'extract':
        jenv = jinja2.Environment()
        template = jenv.from_string(config['extraction_prompt'])
        system_prompt = template.render( classes = config['classes'], traits = config['traits'])
        query = "山の迴饗"
        search_results = str([{"title": "山の迴饗", "snippet": "謝謝大家這麼支持山の迴饗 我們會繼續努力用心做出美味的料理 ————————— ⛰️ 山の迴饗地址：台東縣關山鎮中華路56號訂位專線：0975-957-056 · #山的迴饗 · #夢想起飛"}, {"title": "山的迴饗餐館- 店家介紹", "snippet": "營業登記資料 · 統一編號. 92433454 · 公司狀況. 營業中 · 公司名稱. 山的迴饗餐館 · 公司類型. 獨資 · 資本總額. 30000 · 所在地. 臺東縣關山鎮中福里中華路56號 · 使用發票."}, {"title": "關山漫遊| 💥山の迴饗x night bar", "snippet": "山の迴饗x night bar 即將在12/1號台東關山開幕! 別再煩惱池上、鹿野找不到宵夜餐酒館 各位敬請期待並關注我們✨ night bar❌山的迴饗 12/1 ..."}, {"title": "山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵 - 台灣美食網", "snippet": "山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食營業時間 ; 星期一, 休息 ; 星期二, 10:00–14:00 16:00–21:00 ; 星期三, 10:00–14:00 16:00– ..."}, {"title": "便當｜CP美食- 山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵", "snippet": "餐廳山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食google map 導航. 臺東縣關山鎮中華路56號 +886 975 957 056 ..."}, {"title": "山的迴饗餐館", "snippet": "山的迴饗餐館,統編:92433454,地址:臺東縣關山鎮中福里中華路56號,負責人姓名:周偉慈,設立日期:112年11月15日."}, {"title": "山的迴饗餐館", "snippet": "山的迴饗餐館. 資本總額(元), 30,000. 負責人, 周偉慈. 登記地址, 看地圖 臺東縣關山鎮中福里中華路56號 郵遞區號查詢. 設立日期, 2023-11-15. 資料管理 ..."}, {"title": "山的迴饗餐館, 公司統一編號92433454 - 食品業者登錄資料集", "snippet": "公司或商業登記名稱山的迴饗餐館的公司統一編號是92433454, 登錄項目是餐飲場所, 業者地址是台東縣關山鎮中福里中華路56號, 食品業者登錄字號是V-202257990-00001-5."}, {"title": "山的迴饗餐館, 公司統一編號92433454 - 食品業者登錄資料集", "snippet": "公司或商業登記名稱山的迴饗餐館的公司統一編號是92433454, 登錄項目是公司/商業登記, 業者地址是台東縣關山鎮中福里中華路56號, 食品業者登錄字號是V-202257990-00000-4 ..."}, {"title": "山的迴饗餐館", "snippet": "負責人, 周偉慈 ; 登記地址, 台東縣關山鎮中福里中華路56號 ; 公司狀態, 核准設立 「查詢最新營業狀況請至財政部稅務入口網 」 ; 資本額, 30,000元 ; 所在縣市 ..."}, {"title": "山的迴饗 | 關山美食｜焗烤飯｜酒吧｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食", "顧客評價": "324晚餐餐點豬排簡餐加白醬焗烤等等餐點。\t店家也提供免費的紅茶 綠茶 白開水 多種的調味料自取 總而言之 CP值真的很讚\t空間舒適涼爽，店員服務周到"}, {"title": "類似的店", "snippet": "['中國菜']\t['客家料理']\t['餐廳']\t['熟食店']\t['餐廳']"}, {"telephone_number": "0975 957 056"}])
        user_content = f'''
            `query`: `{query}`,
            `search_results`: {search_results}
        '''
        print(f"user_content -> {user_content}")
        resp = llm( config['provider'], config['model'], system_prompt, user_content)
        print(resp)

    elif args.task == 'classify':
        system_prompt = config['classification_prompt']

    else:
        raise Exception("Invalid task")