# -*- coding: utf-8 -*-
# Imports
import asyncio
import os
import openai

from typing import List, Optional
from pydantic import BaseModel, Field

from langchain.chains.openai_functions.extraction import create_extraction_chain_pydantic
from langchain.chat_models import ChatOpenAI
# from langchain.prompts import ChatPromptTemplate
from langchain.pydantic_v1 import BaseModel
# from langchain.utils.openai_functions import convert_pydantic_to_openai_function

from dotenv import load_dotenv
load_dotenv()

openai.api_key = os.environ['OPENAI_API_KEY']


# App

# Pydantic is an easy way to define a schema
class Person(BaseModel):
    """Information about people to extract."""

    name: str
    age: Optional[int]

# Main function to extract information
def extract_information():
    # Make sure to use a recent model that supports tools
    llm = ChatOpenAI(model="gpt-3.5-turbo-1106")

    return create_extraction_chain_pydantic(Person, llm)