import os

from trallie.schema_generation.schema_generator import SchemaGenerator
from trallie.data_extraction.data_extractor import DataExtractor

# os.environ["GROQ_API_KEY"] = None #ENTER GROQ KEY HERE
# os.environ["OPENAI_API_KEY"] = None #ENTER OPENAI KEY HERE

os.environ["GROQ_API_KEY"] = "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"
os.environ["OPENAI_API_KEY"] = "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"

# Define the path to a set of documents/a data collection for inference
# records = [
#     "data/use-cases/EO_papers/pdf_0808.3837.pdf",
#     "data/use-cases/EO_papers/pdf_1001.4405.pdf",
#     "data/use-cases/EO_papers/pdf_1002.3408.pdf",
# ]

records = [
    "Marie a acheté 3 billets pour le concert de Stromae à Lyon.",
    "Le 14 juillet 2022, Jacques a quitté Marseille pour s’installer à Montréal.",
    "Amazon a livré la commande de Sophie en 48 heures."
]

# Provide a description of the data collection
description = "Un ensemble de phrases aléatoires"

# Initialize the schema generator with a provider and model
schema_generator = SchemaGenerator(provider="openai", model_name="gpt-4o", language="fr")
# Feed records to the LLM and discover schema
print("SCHEMA GENERATION IN ACTION ...")
schema = schema_generator.discover_schema(description, records, from_text=True)
print("Inferred schema", schema)

# Initialize data extractor with a provider and model
data_extractor = DataExtractor(provider="openai", model_name="gpt-4o", language="fr")
# Extract values from the text based on the schema
print("SCHEMA COMPLETION IN ACTION ...")
for record in records:
    extracted_json = data_extractor.extract_data(schema, record, from_text=True)
    print("Extracted attributes:", extracted_json)
