from openai import OpenAI
import requests
from langchain_openai import ChatOpenAI  # pip install -U langchain_openai

url = "https://api.avalai.ir/v1"
key = 'aa-4tvAEazUBovEN1i7i7tdl1PR93OaWXs6hMflR4oQbIIA4K7Z'
api_key = f"Bearer {key}"

with open('models.txt', 'r') as modelfiles:
    models = modelfiles.read().splitlines()

messages = [
    {"role": "system", "content": "تو یک دستیار حقوقی در زمینه قانون گذاری مجلس هستی و بر قوانین مجلس شورای اسلامی کشور ایران به خوبی احاطه داری."},
    {"role": "user", "content": """کلیدواژه های مهم این متن قانونی را استخراج کن و در قالب یک لیست در خروجی نمایش بده: 
     متن قانونی:«ماده واحده – به وزارت دارايي اجازه داده مي شود مبلغ 166848180 ريال از محل درآمد عمومي كل كشور براي تشكيل دادگاه هاي منحل شده و ‌احتياجات ضروري به اختيار وزارت دادگستري بگذارد تا طبق صورت ريزي كه با نظر وزارت دارايي تنظيم خواهد گرديد به مصرف برسانند.
اين قانون كه مشتمل بر ماده واحده است در جلسه پنجشنبه بيست و ششم اسفند ماه يك هزار و سيصد و سي و سه به تصويب مجلس شوراي ملي‌رسيد.»
"""},
]
# استفاده از مدل gpt-4o-mini
model_names = "gpt-4o", "gpt-4o-mini", "deepseek-chat", "anthropic.claude-3-5-haiku-20241022-v1:0", "cohere.command-r-v1:0", "gemini-2.0-flash", "meta.llama3-1-405b-instruct-v1:0", "mistral.mistral-large-2407-v1:0"

results = ''
for model in model_names:
    #if model.startswith('cohere'):
        # model_name = model.split('---')[1]
        model_name = model
        
        llm = ChatOpenAI(
            model=model_name, base_url=url , api_key=api_key
        )

        try:
            result = llm.invoke(messages)

            print(f'OK  --- {model_name}')
            
            results += model_name + " *** \n" +  result.content + '\n\n'
        except:
            print(f'FAIL  --- {model_name}')
            
            results += model_name + " *** \n" +  "error" + '\n\n'

with open('result.txt', mode='w', encoding='utf-8') as file:
    file.write(results.strip())  
    
"""
models_list = []
for item in response.json()['data']:
    models_list.append(item['owned_by'] + "---" + item['id'])

models_list.sort()

models_text = ''
for item in models_list:
    models_text += item + '\n'
    
with open('models.txt', mode='w') as file:
    file.write(models_text.strip())  
"""