Simulate A Perspective

import openai
import instructor
from pydantic import BaseModel, Field
from typing import Iterable

client = instructor.from_openai(openai.OpenAI())


class KnownFact(BaseModel):
    fact: str = Field(description="A fact that the given entity would know")


class Response(BaseModel):
    location: str


def generate_known_facts(entity, context, query) -> Iterable[KnownFact]:
    return client.chat.completions.create(
        model="gpt-4o",
        response_model=Iterable[KnownFact],
        messages=[
            {
                "role": "user",
                "content": f"""Given the following context, list
                the facts that {entity} would know:

                Context:
                {context}
                {query}

                List only the facts relevant to {entity}.
                """,
            }
        ],
    )


def answer_question_based_on_facts(entity, query, known_facts) -> Response:
    return client.chat.completions.create(
        model="gpt-4o",
        response_model=Response,
        messages=[
            {
                "role": "system",
                "content": f"""You are {entity}. Answer the following question
                based only on these facts you know:
                {" ".join([str(fact) for fact in known_facts])}""",
            },
            {
                "role": "user",
                "content": f"Question: {query}",
            },
        ],
    )


if __name__ == "__main__":
    entity = "Alice"
    context = """Alice puts the book on the table.
        Alice leaves the room.
        Bob moves the book to the shelf.
        """
    query = f"Where does {entity} think the book is?"

    known_facts = generate_known_facts(entity, context, query)
    response = answer_question_based_on_facts(entity, query, known_facts)

    for fact in known_facts:
        print(fact)
        #> fact='Alice puts the book on the table.'
        #> fact='Alice leaves the room.'
    print(response.location)
    #> on the table