Entity recognition

The following example, which closely inspired by the Claude documentation, shows how .chat_structured() can be used to perform entity recognition.

from chatlas import ChatOpenAI
from pydantic import BaseModel, Field
import pandas as pd

# | warning: false
text = "John works at Google in New York. He met with Sarah, the CEO of Acme Inc., last week in San Francisco."


class NamedEntity(BaseModel):
    """Named entity in the text."""

    name: str = Field(description="The extracted entity name")

    type_: str = Field(description="The entity type, e.g. 'person', 'location', 'organization'")

    context: str = Field(description="The context in which the entity appears in the text.")


class NamedEntities(BaseModel):
    """Named entities in the text."""

    entities: list[NamedEntity] = Field(description="Array of named entities")


chat = ChatOpenAI()
data = chat.chat_structured(text, data_model=NamedEntities)
pd.DataFrame([e.model_dump() for e in data.entities])

	name	type_	context
0	John	person	Employee at Google
1	Google	organization	Employer of John, located in New York
2	New York	location	Location of Google's office where John works
3	Sarah	person	CEO of Acme Inc., met with John
4	Acme Inc.	organization	Sarah is the CEO
5	San Francisco	location	Place where John and Sarah met last week