|
|
|
|
|
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'] |
|
|
|
|
|
|
|
|
|
|
|
class Person(BaseModel): |
|
"""Information about people to extract.""" |
|
|
|
name: str |
|
age: Optional[int] = None |
|
|
|
|
|
def extract_information(): |
|
|
|
llm = ChatOpenAI(model="gpt-3.5-turbo-1106") |
|
|
|
return create_extraction_chain_pydantic(Person, llm) |
|
|
|
|
|
if __name__ == "__main__": |
|
text = "My name is John and I am 20 years old. My name is sally and I am 30 years old." |
|
chain = extract_information() |
|
print(chain.invoke({"input": text})["text"]) |
|
|
|
async def extract_information_async(message: str): |
|
return chain.invoke({"input": message})["text"] |
|
|
|
async def main(): |
|
res = await extract_information_async(text) |
|
print(res) |
|
|
|
asyncio.run(main()) |