examples/test/test_assistant/test_assistant_retrieval.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import time

from openai import OpenAI

from project_settings import environment, project_path


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--openai_api_key",
        default=environment.get("openai_api_key", default=None, dtype=str),
        type=str
    )
    parser.add_argument(
        "--document_file",
        default=(project_path / "data/三国演义.txt").as_posix(),
        type=str
    )
    args = parser.parse_args()
    return args


def main():
    args = get_args()

    client = OpenAI(
        api_key=args.openai_api_key
    )

    document = client.files.create(
        file=open(args.document_file, "rb"),
        purpose='assistants'
    )

    assistant = client.beta.assistants.create(
        name="小说专家",
        instructions="根据小说内容回答问题。",
        tools=[{"type": "retrieval"}],
        model="gpt-4-1106-preview",
        file_ids=[document.id]

    )
    print(f"assistant.id: {assistant.id}")

    thread = client.beta.threads.create()
    print(f"thread.id: {thread.id}")

    message = client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content="刘备和张飞是什么关系。"
    )

    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant.id,
    )

    run = client.beta.threads.runs.retrieve(
        thread_id=thread.id,
        run_id=run.id
    )

    time.sleep(10)

    messages = client.beta.threads.messages.list(
        thread_id=thread.id
    )
    print(messages)

    return


if __name__ == '__main__':
    main()