forked from shadowaxe99/poddyx
-
Notifications
You must be signed in to change notification settings - Fork 0
/
openai_function_call.py
53 lines (46 loc) · 1.48 KB
/
openai_function_call.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from typing import List
from openai.api_resources.abstract.api_resource import APIResource
from pydantic import BaseModel
class OpenAISchema(BaseModel):
pass
class User(OpenAISchema):
name: str
age: int
class MultiTask(OpenAISchema):
def __init__(self, schema):
self.schema = schema
self.openai_schema = {"name": schema}
@staticmethod
def from_streaming_response(response):
for message in response['choices'][0]['message']['role']:
yield User(**message)
def stream_extract(input: str) -> List[User]:
openai = APIResource(api_key="your-api-key")
MultiUser = MultiTask(User)
completion = openai.ChatCompletion.create(
model="gpt-4-0613",
temperature=0.1,
stream=True,
functions=[MultiUser.openai_schema],
function_call={"name": MultiUser.openai_schema["name"]},
messages=[
{
"role": "system",
"content": "You are a perfect entity extraction system",
},
{
"role": "user",
"content": (
f"Consider the data below:\n{input}\n"
"Correctly segment it into entities\n"
"Make sure the JSON is correct"
),
},
],
max_tokens=1000,
)
users = []
for user in MultiUser.from_streaming_response(completion):
assert isinstance(user, User)
users.append(user)
return users