-
Notifications
You must be signed in to change notification settings - Fork 3
/
test_vector_stores.py
160 lines (122 loc) · 5.35 KB
/
test_vector_stores.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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import openai
import openai_responses
from openai_responses import OpenAIMock
@openai_responses.mock()
def test_create_vector_store(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
vector_store = client.beta.vector_stores.create(name="Support FAQ")
assert vector_store.name == "Support FAQ"
assert openai_mock.beta.vector_stores.create.route.call_count == 1
@openai_responses.mock()
def test_create_vector_store_with_file_ids(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
file = client.files.create(
file=open("examples/example.json", "rb"),
purpose="assistants",
)
vector_store = client.beta.vector_stores.create(
name="Support FAQ", file_ids=[file.id]
)
vector_store_files = client.beta.vector_stores.files.list(vector_store.id)
assert len(vector_store_files.data) == 1
assert openai_mock.files.create.route.call_count == 1
assert openai_mock.beta.vector_stores.create.route.call_count == 1
assert openai_mock.beta.vector_stores.files.list.route.call_count == 1
@openai_responses.mock()
def test_create_vector_store_from_assistant_create(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
file = client.files.create(
file=open("examples/example.json", "rb"),
purpose="assistants",
)
assistant = client.beta.assistants.create(
model="gpt-4o",
tool_resources={
"file_search": {
"vector_stores": [
{
"file_ids": [file.id],
}
]
}
},
)
assert assistant.tool_resources
assert assistant.tool_resources.file_search
assert assistant.tool_resources.file_search.vector_store_ids
vector_store = client.beta.vector_stores.retrieve(
assistant.tool_resources.file_search.vector_store_ids[0]
)
vector_store_files = client.beta.vector_stores.files.list(vector_store.id)
assert len(vector_store_files.data) == 1
assert openai_mock.files.create.route.call_count == 1
assert openai_mock.beta.assistants.create.route.call_count == 1
assert openai_mock.beta.vector_stores.retrieve.route.call_count == 1
assert openai_mock.beta.vector_stores.files.list.route.call_count == 1
@openai_responses.mock()
def test_create_vector_store_from_thread_create(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
file = client.files.create(
file=open("examples/example.json", "rb"),
purpose="assistants",
)
thread = client.beta.threads.create(
tool_resources={
"file_search": {
"vector_stores": [
{
"file_ids": [file.id],
}
]
}
}
)
assert thread.tool_resources
assert thread.tool_resources.file_search
assert thread.tool_resources.file_search.vector_store_ids
vector_store = client.beta.vector_stores.retrieve(
thread.tool_resources.file_search.vector_store_ids[0]
)
vector_store_files = client.beta.vector_stores.files.list(vector_store.id)
assert len(vector_store_files.data) == 1
assert openai_mock.files.create.route.call_count == 1
assert openai_mock.beta.threads.create.route.call_count == 1
assert openai_mock.beta.vector_stores.retrieve.route.call_count == 1
assert openai_mock.beta.vector_stores.files.list.route.call_count == 1
@openai_responses.mock()
def test_list_vector_stores(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
for i in range(10):
client.beta.vector_stores.create(name=f"vector-store-{i}")
vector_stores = client.beta.vector_stores.list()
assert len(vector_stores.data) == 10
assert openai_mock.beta.vector_stores.create.route.call_count == 10
assert openai_mock.beta.vector_stores.list.route.call_count == 1
@openai_responses.mock()
def test_retrieve_vector_store(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
vector_store = client.beta.vector_stores.create(name="Support FAQ")
found = client.beta.vector_stores.retrieve(vector_store.id)
assert vector_store.name == "Support FAQ"
assert found.name == vector_store.name
assert openai_mock.beta.vector_stores.create.route.call_count == 1
assert openai_mock.beta.vector_stores.retrieve.route.call_count == 1
@openai_responses.mock()
def test_update_vector_store(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
vector_store = client.beta.vector_stores.create(name="Support FAQ")
updated = client.beta.vector_stores.update(
vector_store.id,
name="Support FAQ Updated",
)
assert updated.id == vector_store.id
assert updated.name == "Support FAQ Updated"
assert openai_mock.beta.vector_stores.create.route.call_count == 1
assert openai_mock.beta.vector_stores.update.route.call_count == 1
@openai_responses.mock()
def test_delete_vector_store(openai_mock: OpenAIMock):
client = openai.Client(api_key="sk-fake123")
vector_store = client.beta.vector_stores.create(name="Support FAQ")
assert client.beta.vector_stores.delete(vector_store.id).deleted
assert openai_mock.beta.vector_stores.create.route.call_count == 1
assert openai_mock.beta.vector_stores.delete.route.call_count == 1