How to extract entities from text with Python

import os
from openai import OpenAI
import json

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

def extract_entities_stream(text: str):
    prompt = (
        f"Extract all named entities from the following text as a JSON list of objects with 'entity' and 'type' keys.\n"
        f"Text: {text}\n"
        f"Output format example: [{'{"entity": "Apple", "type": "Organization"}'}]"
    )

    stream = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        stream=True
    )

    full_response = ""
    for chunk in stream:
        delta = chunk.choices[0].delta.content or ""
        full_response += delta
    try:
        entities = json.loads(full_response)
    except json.JSONDecodeError:
        entities = []
    return entities

if __name__ == "__main__":
    sample_text = "Amazon's headquarters are in Seattle, and Jeff Bezos founded it in 1994."
    extracted = extract_entities_stream(sample_text)
    print("Extracted entities:")
    for entity in extracted:
        print(f"- {entity['entity']} ({entity['type']})")

import os
import asyncio
from openai import OpenAI
import json

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

async def extract_entities_async(text: str):
    prompt = (
        f"Extract all named entities from the following text as a JSON list of objects with 'entity' and 'type' keys.\n"
        f"Text: {text}\n"
        f"Output format example: [{'{"entity": "Apple", "type": "Organization"}'}]"
    )

    response = await client.chat.completions.acreate(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}]
    )

    content = response.choices[0].message.content
    try:
        entities = json.loads(content)
    except json.JSONDecodeError:
        entities = []
    return entities

async def main():
    sample_text = "No entities here, just plain text."
    extracted = await extract_entities_async(sample_text)
    print("Extracted entities:")
    for entity in extracted:
        print(f"- {entity['entity']} ({entity['type']})")

if __name__ == "__main__":
    asyncio.run(main())

import os
from anthropic import Anthropic
import json

client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

def extract_entities_claude(text: str):
    system_prompt = "You are a helpful assistant that extracts named entities as JSON objects with 'entity' and 'type'."
    user_message = f"Extract entities from this text:\n{text}"

    response = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": user_message}]
    )

    content = response.content[0].text
    try:
        entities = json.loads(content)
    except json.JSONDecodeError:
        entities = []
    return entities

if __name__ == "__main__":
    sample_text = "Google was founded by Larry Page and Sergey Brin."
    extracted = extract_entities_claude(sample_text)
    print("Extracted entities:")
    for entity in extracted:
        print(f"- {entity['entity']} ({entity['type']})")