model.py

import os
import argparse
import time

from dotenv import load_dotenv
import anthropic
from openai import OpenAI

from utils import parse_json_garbage

load_dotenv()

def llm( provider, model, system_prompt, user_content, delay:int = 10):
    """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,
            # stream = True
        )
        response = chat_completion.choices[0].message.content

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


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--provider", type=str, default='anthropic', help="openai or anthropic")
    parser.add_argument("--model", type=str, default='claude-3-sonnet-20240229', help="Model name for the API", 
                        choices = ["claude-3-sonnet-20240229", "claude-3-haiku-20240307", "gpt-3.5-turbo-0125", "gpt-4-0125-preview"])
    parser.add_argument("--classes", type=list, default=['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', '西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)', '西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)', '早餐'])
    parser.add_argument("--task", type=list, default='extract', choices=['extract', 'classify'])
    args = parser.parse_args()



    classes = ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', ]
    backup_classes = [ '中式', '西式']
    
    extraction_prompt = '''
            As a helpful and rigorous retail analyst, given the provided query and a list of search results for the query, 
            your task is to first identify relevant information of the identical store based on store name and proxmity of address if known. After that, extract `store_name`, `address`, `description`, `category` and `phone_number` from the found relevant information, where `category` can only be `小吃店`, `日式料理(含居酒屋，串燒)`, `火(鍋／爐)`, `東南亞料理(不含日韓)`, `海鮮熱炒`,  `特色餐廳(含雞、鵝、牛、羊肉)`, `傳統餐廳`, `燒烤`, `韓式料理(含火鍋，烤肉)`, `西餐廳(含美式，義式，墨式)`, `西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)`, `西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)` or `早餐`. 
            It's very important to omit unrelated results. Do not make up any assumption.
            Please think step by step, and output in json format. An example output json is like {"store_name": "...", "address": "...", "description": "... products, service or highlights ...", "category": "...", "phone_number": "..."}
            If no relevant information has been found, simply output json with empty values.
            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
        '''
    classification_prompt = f"""
            As a helpful and rigorous retail analyst, given the provided information about a store, 
            your task is two-fold. First, classify provided evidence below into the mostly relevant category from the following: {classes}. 
            Second, if no relevant information has been found, classify the evidence into the mostly relevant supercategory from the following: {backup_classes}.
            It's very important to omit unrelated piece of evidence and don't make up any assumption.
            Please think step by step, and must output in json format. An example output json is like {{"category": "..."}}
            If no relevant piece of information can ever be found at all, simply output json with empty string "".
            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
    """
    
    if args.task == 'extract':
        system_prompt = extraction_prompt
    elif args.task == 'classify':
        system_prompt = classification_prompt
    else:
        raise Exception("Invalid task")
    
    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"}])

    # query = "大吃一斤泰國蝦麻辣牛肉爐"
    # search_results = str([{"title": "大吃一斤泰國蝦麻辣牛肉爐", "snippet": "... 一支、本店特賣價600元免費代料理、 保證、活的!歡迎來電預定0975-147-848大吃一斤活蝦料理店新北市三重區自強路一段222號泰國蝦活蝦現場料理不漲價一斤維持一斤480元."}, {"title": "大吃一斤泰國蝦麻辣牛肉爐", "snippet": "... 一支、本店特賣價600元免費代料理、 保證、活的!歡迎來電預定0975-147-848大吃一斤活蝦料理店新北市三重區自強路一段222號泰國蝦活蝦現場料理不漲價一斤維持一斤480元."}, {"title": "大吃一斤", "snippet": "大吃一斤在foodpanda點的到，更多New Taipei City 推薦美食，線上訂立即送，下載foodpanda APP，20分鐘外送上門！瀏覽菜單和獨家優惠折扣."}, {"title": "大吃一斤(新北板橋店)菜單", "snippet": "大吃一斤(新北板橋店) 在foodpanda點的到，更多New Taipei City 推薦美食，線上訂立即送，下載foodpanda APP，20分鐘外送上門！"}, {"title": "大吃一斤活蝦餐廳- 店家介紹", "snippet": "大吃一斤活蝦餐廳. 資本總額. 200000. 代表人. 李錦鴻. 所在區域. 新北市. 所在地. 新北市三重區自強路1段222號(1樓). 商業類型. 獨資. 異動紀錄. 1111108. 營業狀態為: ..."}, {"title": "新北市| 三重區大吃一斤(泰國蝦牛肉料理店)", "snippet": "大吃一斤(泰國蝦牛肉料理店) 餐廳介紹 ; phone icon 電話, 0975 147 848 ; 營業時間, 星期一17:00–04:00 星期二17:00–04:00 星期三17:00–04:00 星期四17:00– ..."}, {"title": "大吃一斤活蝦餐廳", "snippet": "大吃一斤活蝦餐廳. 負責人姓名, 李錦鴻. 地址, 新北市三重區自強路1段222號(1樓). 現況, 核准設立. 資本額(元), 200,000. 組織類型, 獨資. 登記機關, 新北市政府經濟發展局."}, {"title": "【大吃一斤(泰國蝦牛肉料理店)】網友評價- 新北三重區合菜餐廳", "snippet": "大吃一斤(泰國蝦牛肉料理店) - 網友評論、最新食記(132則) 評分: 4.4分。大吃一斤(泰國蝦牛肉料理店)是位於新北三重區的餐廳，地址: 新北市 ... 生猛活海鮮."}, {"title": "大吃一斤生猛海鮮/活魚料理超值優惠方案", "snippet": "大吃一斤生猛海鮮/活魚料理. 電話：0975-147-848. 地址：新北市三重區自強路一段222號. 營業時間：週一至週日17: ..."}, {"title": "大吃一斤三重店 (泰國蝦料理.平價快炒熱炒.各式海鮮)", "顧客評價": "塔香蛤蜊、胡椒蝦、檸檬蝦、胡椒鳳螺 口味不錯食材新鮮 拍照時蛤蜊已經快被小孩吃光\t蝦子不大，店面不大，魚腥味很重，廁所很多蚊子，連菜裡面也有蚊子🦟，根本吃不下去\t新鮮好吃😋老闆人很Nice 推薦鹽烤蝦以及蒜味奶油蝦👍👍👍"}, {"title": "類似的店", "snippet": "['海鮮']\t['海鮮']\t['海鮮']\t['海鮮']"}, {"telephone_number": "0975 147 848"}])

    if args.provider == "openai":
        client = OpenAI( organization = os.getenv('ORGANIZATION_ID'))
        # categories = ", ".join([ "`"+x+"`" for x in args.classes if x!='早餐' ])+ " or " + "`早餐`"
        user_content = f'''
            `query`: `{query}`,
            `search_results`: {search_results}
        '''
        resp = llm( args.provider, args.model, system_prompt, user_content)
        print(f"resp -> {resp}")


    elif args.provider == "anthropic":
        client = anthropic.Client(api_key=os.getenv('ANTHROPIC_APIKEY'))
        user_content = f'''
            `query`: `{query}`,
            `search_results`: {search_results}
        '''
        print(f"user_content -> {user_content}")
        resp = llm( args.provider, args.model, system_prompt, user_content)
        print(resp)