diff --git a/README.md b/README.md index e0a72bd..722725a 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Pytest plugin for automatically mocking OpenAI requests. Powered by [RESPX](http - [Messages](https://platform.openai.com/docs/api-reference/messages) - [Runs](https://platform.openai.com/docs/api-reference/runs) - [Run Steps](https://platform.openai.com/docs/api-reference/run-steps) +- [Vector Stores](https://platform.openai.com/docs/api-reference/vector-stores) +- [Vector Store Files](https://platform.openai.com/docs/api-reference/vector-stores-files) View full support coverage [here](https://mharrisb1.github.io/openai-responses-python/coverage). diff --git a/docs/coverage.md b/docs/coverage.md index 36c9ea2..4dcdf00 100644 --- a/docs/coverage.md +++ b/docs/coverage.md @@ -77,16 +77,16 @@ The end-goal of this library is to eventually support all OpenAI API routes. See | List run steps | :material-check:{ .green } | - | Stateful | | Retrieve run step | :material-check:{ .green } | - | Stateful | | **Vector Stores** | -| Create vector store | :material-close:{ .red } | - | - | -| List vector stores | :material-close:{ .red } | - | - | -| Retrieve vector store | :material-close:{ .red } | - | - | -| Modify vector store | :material-close:{ .red } | - | - | -| Delete vector store | :material-close:{ .red } | - | - | +| Create vector store | :material-check:{ .green } | - | Stateful | +| List vector stores | :material-check:{ .green } | - | Stateful | +| Retrieve vector store | :material-check:{ .green } | - | Stateful | +| Modify vector store | :material-check:{ .green } | - | Stateful | +| Delete vector store | :material-check:{ .green } | - | Stateful | | **Vector Store Files** | -| Create vector store file | :material-close:{ .red } | - | - | -| List vector store files | :material-close:{ .red } | - | - | -| Retrieve vector store file | :material-close:{ .red } | - | - | -| Delete vector store file | :material-close:{ .red } | - | - | +| Create vector store file | :material-check:{ .green } | - | Stateful | +| List vector store files | :material-check:{ .green } | - | Stateful | +| Retrieve vector store file | :material-check:{ .green } | - | Stateful | +| Delete vector store file | :material-check:{ .green } | - | Stateful | | **Vector Store File Batches** | | Create vector store file batch | :material-close:{ .red } | - | - | | Retrieve vector store file batch | :material-close:{ .red } | - | - | diff --git a/examples/test_vector_store_files.py b/examples/test_vector_store_files.py new file mode 100644 index 0000000..bb79f1a --- /dev/null +++ b/examples/test_vector_store_files.py @@ -0,0 +1,110 @@ +import openai + +import openai_responses +from openai_responses import OpenAIMock + + +@openai_responses.mock() +def test_create_vector_store_file(openai_mock: OpenAIMock): + client = openai.Client(api_key="sk-fake123") + + vector_store = client.beta.vector_stores.create(name="Support FAQ") + file = client.files.create( + file=open("examples/example.json", "rb"), + purpose="assistants", + ) + + vector_store_file = client.beta.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + assert vector_store_file.vector_store_id == vector_store.id + assert vector_store_file.id == file.id + + 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.create.route.call_count == 1 + + +@openai_responses.mock() +def test_list_vector_store_files(openai_mock: OpenAIMock): + client = openai.Client(api_key="sk-fake123") + + vector_store = client.beta.vector_stores.create(name="Support FAQ") + + for _ in range(10): + file = client.files.create( + file=open("examples/example.json", "rb"), + purpose="assistants", + ) + + client.beta.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + vector_store_files = client.beta.vector_stores.files.list(vector_store.id) + + assert len(vector_store_files.data) == 10 + + assert openai_mock.files.create.route.call_count == 10 + assert openai_mock.beta.vector_stores.create.route.call_count == 1 + assert openai_mock.beta.vector_stores.files.create.route.call_count == 10 + assert openai_mock.beta.vector_stores.files.list.route.call_count == 1 + + +@openai_responses.mock() +def test_retrieve_vector_store_file(openai_mock: OpenAIMock): + client = openai.Client(api_key="sk-fake123") + + vector_store = client.beta.vector_stores.create(name="Support FAQ") + file = client.files.create( + file=open("examples/example.json", "rb"), + purpose="assistants", + ) + + vector_store_file = client.beta.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + found = client.beta.vector_stores.files.retrieve( + vector_store_file.id, + vector_store_id=vector_store.id, + ) + + assert found.id == vector_store_file.id + + 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.create.route.call_count == 1 + assert openai_mock.beta.vector_stores.files.retrieve.route.call_count == 1 + + +@openai_responses.mock() +def test_delete_vector_store_file(openai_mock: OpenAIMock): + client = openai.Client(api_key="sk-fake123") + + vector_store = client.beta.vector_stores.create(name="Support FAQ") + file = client.files.create( + file=open("examples/example.json", "rb"), + purpose="assistants", + ) + + vector_store_file = client.beta.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + deleted = client.beta.vector_stores.files.delete( + vector_store_file.id, + vector_store_id=vector_store.id, + ) + + assert deleted.deleted + + 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.create.route.call_count == 1 + assert openai_mock.beta.vector_stores.files.delete.route.call_count == 1 diff --git a/examples/test_vector_stores.py b/examples/test_vector_stores.py new file mode 100644 index 0000000..4049493 --- /dev/null +++ b/examples/test_vector_stores.py @@ -0,0 +1,70 @@ +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_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 diff --git a/pyproject.toml b/pyproject.toml index b1c96ed..2b32553 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,10 +20,11 @@ openai_responses = "openai_responses.plugin" [tool.poetry.dependencies] python = ">=3.9,<4.0" +faker = ">=24.2.0,<25.0.0" +faker-openai-api-provider = "0.2.0" openai = "^1.25" -respx = "^0.20.2" -faker-openai-api-provider = "^0.1.0" requests-toolbelt = "^1.0.0" +respx = "^0.20.2" [tool.poetry.group.dev.dependencies] black = "^24.2.0" diff --git a/src/openai_responses/_routes/__init__.py b/src/openai_responses/_routes/__init__.py index 84cb7c1..0498271 100644 --- a/src/openai_responses/_routes/__init__.py +++ b/src/openai_responses/_routes/__init__.py @@ -42,6 +42,19 @@ RunCancelRoute, ) from .run_steps import RunStepListRoute, RunStepRetrieveRoute +from .vector_stores import ( + VectorStoreCreateRoute, + VectorStoreListRoute, + VectorStoreRetrieveRoute, + VectorStoreUpdateRoute, + VectorStoreDeleteRoute, +) +from .vector_store_files import ( + VectorStoreFileCreateRoute, + VectorStoreFileListRoute, + VectorStoreFileRetrieveRoute, + VectorStoreFileDeleteRoute, +) __all__ = [ "BetaRoutes", @@ -85,6 +98,7 @@ class BetaRoutes: def __init__(self, router: respx.MockRouter, state: StateStore) -> None: self.assistants = AssistantsRoutes(router, state) self.threads = ThreadRoutes(router, state) + self.vector_stores = VectorStoreRoutes(router, state) class AssistantsRoutes: @@ -133,3 +147,22 @@ class RunStepRoutes: def __init__(self, router: respx.MockRouter, state: StateStore) -> None: self.list = RunStepListRoute(router, state) self.retrieve = RunStepRetrieveRoute(router, state) + + +class VectorStoreRoutes: + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + self.create = VectorStoreCreateRoute(router, state) + self.list = VectorStoreListRoute(router, state) + self.retrieve = VectorStoreRetrieveRoute(router, state) + self.update = VectorStoreUpdateRoute(router, state) + self.delete = VectorStoreDeleteRoute(router, state) + + self.files = VectorStoreFileRoutes(router, state) + + +class VectorStoreFileRoutes: + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + self.create = VectorStoreFileCreateRoute(router, state) + self.list = VectorStoreFileListRoute(router, state) + self.retrieve = VectorStoreFileRetrieveRoute(router, state) + self.delete = VectorStoreFileDeleteRoute(router, state) diff --git a/src/openai_responses/_routes/assistants.py b/src/openai_responses/_routes/assistants.py index 97ef9ce..5923648 100644 --- a/src/openai_responses/_routes/assistants.py +++ b/src/openai_responses/_routes/assistants.py @@ -12,11 +12,8 @@ from ._base import StatefulRoute from ..stores import StateStore -from .._types.partials.assistants import ( - PartialAssistant, - PartialAssistantList, - PartialAssistantDeleted, -) +from .._types.partials.sync_cursor_page import PartialSyncCursorPage +from .._types.partials.assistants import PartialAssistant, PartialAssistantDeleted from .._utils.faker import faker from .._utils.serde import json_loads, model_dict, model_parse @@ -62,7 +59,7 @@ def _build(partial: PartialAssistant, request: httpx.Request) -> Assistant: class AssistantListRoute( - StatefulRoute[SyncCursorPage[Assistant], PartialAssistantList] + StatefulRoute[SyncCursorPage[Assistant], PartialSyncCursorPage[PartialAssistant]] ): def __init__(self, router: respx.MockRouter, state: StateStore) -> None: super().__init__( @@ -96,7 +93,7 @@ def _handler(self, request: httpx.Request, route: respx.Route) -> httpx.Response @staticmethod def _build( - partial: PartialAssistantList, + partial: PartialSyncCursorPage[PartialAssistant], request: httpx.Request, ) -> SyncCursorPage[Assistant]: raise NotImplementedError diff --git a/src/openai_responses/_routes/messages.py b/src/openai_responses/_routes/messages.py index b48a612..07d6b68 100644 --- a/src/openai_responses/_routes/messages.py +++ b/src/openai_responses/_routes/messages.py @@ -12,11 +12,8 @@ from ._base import StatefulRoute from ..stores import StateStore -from .._types.partials.messages import ( - PartialMessage, - PartialMessageList, - PartialMessageDeleted, -) +from .._types.partials.sync_cursor_page import PartialSyncCursorPage +from .._types.partials.messages import PartialMessage, PartialMessageDeleted from .._utils.faker import faker from .._utils.serde import json_loads, model_dict, model_parse @@ -100,7 +97,9 @@ def _build(partial: PartialMessage, request: httpx.Request) -> Message: return model_parse(Message, defaults | partial | content) -class MessageListRoute(StatefulRoute[SyncCursorPage[Message], PartialMessageList]): +class MessageListRoute( + StatefulRoute[SyncCursorPage[Message], PartialSyncCursorPage[PartialMessage]] +): def __init__(self, router: respx.MockRouter, state: StateStore) -> None: super().__init__( route=router.get( @@ -153,7 +152,7 @@ def _handler( @staticmethod def _build( - partial: PartialMessageList, + partial: PartialSyncCursorPage[PartialMessage], request: httpx.Request, ) -> SyncCursorPage[Message]: raise NotImplementedError diff --git a/src/openai_responses/_routes/models.py b/src/openai_responses/_routes/models.py index f3028e6..fa94bcb 100644 --- a/src/openai_responses/_routes/models.py +++ b/src/openai_responses/_routes/models.py @@ -10,12 +10,15 @@ from ._base import StatefulRoute from ..stores import StateStore -from .._types.partials.models import PartialModel, PartialModelList +from .._types.partials.models import PartialModel +from .._types.partials.sync_cursor_page import PartialSyncCursorPage from .._utils.serde import model_dict -class ModelListRoute(StatefulRoute[SyncCursorPage[Model], PartialModelList]): +class ModelListRoute( + StatefulRoute[SyncCursorPage[Model], PartialSyncCursorPage[PartialModel]] +): def __init__(self, router: respx.MockRouter, state: StateStore) -> None: super().__init__( route=router.get(url__regex="/models"), @@ -32,7 +35,7 @@ def _handler(self, request: httpx.Request, route: respx.Route) -> httpx.Response @staticmethod def _build( - partial: PartialModelList, + partial: PartialSyncCursorPage[PartialModel], request: httpx.Request, ) -> SyncCursorPage[Model]: raise NotImplementedError diff --git a/src/openai_responses/_routes/run_steps.py b/src/openai_responses/_routes/run_steps.py index 4382223..b827cd2 100644 --- a/src/openai_responses/_routes/run_steps.py +++ b/src/openai_responses/_routes/run_steps.py @@ -10,7 +10,8 @@ from ._base import StatefulRoute from ..stores import StateStore -from .._types.partials.run_steps import PartialRunStep, PartialRunStepList +from .._types.partials.run_steps import PartialRunStep +from .._types.partials.sync_cursor_page import PartialSyncCursorPage from .._utils.serde import model_dict @@ -18,7 +19,9 @@ __all__ = ["RunStepListRoute", "RunStepRetrieveRoute"] -class RunStepListRoute(StatefulRoute[SyncCursorPage[RunStep], PartialRunStepList]): +class RunStepListRoute( + StatefulRoute[SyncCursorPage[RunStep], PartialSyncCursorPage[PartialRunStep]] +): def __init__(self, router: respx.MockRouter, state: StateStore) -> None: super().__init__( route=router.get( @@ -75,7 +78,7 @@ def _handler( @staticmethod def _build( - partial: PartialRunStepList, + partial: PartialSyncCursorPage[PartialRunStep], request: httpx.Request, ) -> SyncCursorPage[RunStep]: raise NotImplementedError diff --git a/src/openai_responses/_routes/runs.py b/src/openai_responses/_routes/runs.py index 4d33332..4a879fb 100644 --- a/src/openai_responses/_routes/runs.py +++ b/src/openai_responses/_routes/runs.py @@ -17,7 +17,8 @@ from ..helpers.builders.threads import thread_from_create_request from ..stores import StateStore -from .._types.partials.runs import PartialRun, PartialRunList +from .._types.partials.sync_cursor_page import PartialSyncCursorPage +from .._types.partials.runs import PartialRun from .._utils.copy import model_copy from .._utils.faker import faker @@ -148,7 +149,9 @@ def _build(partial: PartialRun, request: httpx.Request) -> Run: return RunCreateRoute._build(partial, request) -class RunListRoute(StatefulRoute[SyncCursorPage[Run], PartialRunList]): +class RunListRoute( + StatefulRoute[SyncCursorPage[Run], PartialSyncCursorPage[PartialRun]] +): def __init__(self, router: respx.MockRouter, state: StateStore) -> None: super().__init__( route=router.get(url__regex=r"/threads/(?P[a-zA-Z0-9\_]+)/runs"), @@ -196,7 +199,10 @@ def _handler( ) @staticmethod - def _build(partial: PartialRunList, request: httpx.Request) -> SyncCursorPage[Run]: + def _build( + partial: PartialSyncCursorPage[PartialRun], + request: httpx.Request, + ) -> SyncCursorPage[Run]: raise NotImplementedError diff --git a/src/openai_responses/_routes/vector_store_files.py b/src/openai_responses/_routes/vector_store_files.py new file mode 100644 index 0000000..bb37e07 --- /dev/null +++ b/src/openai_responses/_routes/vector_store_files.py @@ -0,0 +1,235 @@ +from typing import Any, Literal +from typing_extensions import override + +import httpx +import respx + +from openai.pagination import SyncCursorPage +from openai.types.beta.vector_stores.vector_store_file import VectorStoreFile +from openai.types.beta.vector_stores.vector_store_file_deleted import ( + VectorStoreFileDeleted, +) + +from ._base import StatefulRoute + +from ..stores import StateStore +from .._types.partials.deleted import PartialResourceDeleted +from .._types.partials.sync_cursor_page import PartialSyncCursorPage +from .._types.partials.vector_store_files import PartialVectorStoreFile + +from .._utils.serde import json_loads, model_dict, model_parse +from .._utils.time import utcnow_unix_timestamp_s + +__all__ = [ + "VectorStoreFileCreateRoute", + "VectorStoreFileListRoute", + "VectorStoreFileRetrieveRoute", + "VectorStoreFileDeleteRoute", +] + + +class VectorStoreFileCreateRoute( + StatefulRoute[VectorStoreFile, PartialVectorStoreFile] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.post( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)/files" + ), + status_code=201, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + + vector_store_id = kwargs["vector_store_id"] + found_vector_store = self._state.beta.vector_stores.get(vector_store_id) + if not found_vector_store: + return httpx.Response(404) + + model = self._build({"vector_store_id": vector_store_id}, request) + found_file = self._state.files.get(model.id) + if not found_file: + return httpx.Response(404) + + self._state.beta.vector_stores.files.put(model) + return httpx.Response(status_code=self._status_code, json=model_dict(model)) + + @staticmethod + def _build( + partial: PartialVectorStoreFile, + request: httpx.Request, + ) -> VectorStoreFile: + content = json_loads(request.content) + defaults: PartialVectorStoreFile = { + "id": content["file_id"], + "created_at": utcnow_unix_timestamp_s(), + "object": "vector_store.file", + "status": "completed", + "usage_bytes": 0, + } + return model_parse(VectorStoreFile, defaults | partial) + + +class VectorStoreFileListRoute( + StatefulRoute[ + SyncCursorPage[VectorStoreFile], PartialSyncCursorPage[PartialVectorStoreFile] + ] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.get( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)/files" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + + vector_store_id = kwargs["vector_store_id"] + found_vector_store = self._state.beta.vector_stores.get(vector_store_id) + if not found_vector_store: + return httpx.Response(404) + + limit = request.url.params.get("limit") + order = request.url.params.get("order") + after = request.url.params.get("after") + before = request.url.params.get("before") + filter = request.url.params.get("filter") + + data = self._state.beta.vector_stores.files.list( + vector_store_id, + limit, + order, + after, + before, + filter, + ) + result_count = len(data) + total_count = len(self._state.beta.vector_stores.files.list(vector_store_id)) + has_data = bool(result_count) + has_more = total_count != result_count + first_id = data[0].id if has_data else None + last_id = data[-1].id if has_data else None + model = SyncCursorPage[VectorStoreFile](data=data) + return httpx.Response( + status_code=200, + json=model_dict(model) + | {"first_id": first_id, "last_id": last_id, "has_more": has_more}, + ) + + @staticmethod + def _build( + partial: PartialSyncCursorPage[PartialVectorStoreFile], + request: httpx.Request, + ) -> SyncCursorPage[VectorStoreFile]: + raise NotImplementedError + + +class VectorStoreFileRetrieveRoute( + StatefulRoute[VectorStoreFile, PartialVectorStoreFile] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.get( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)/files/(?P[a-zA-Z0-9\-]+)" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + vector_store_id = kwargs["vector_store_id"] + found_vector_store = self._state.beta.vector_stores.get(vector_store_id) + if not found_vector_store: + return httpx.Response(404) + + file_id = kwargs["file_id"] + found = self._state.beta.vector_stores.files.get(file_id) + if not found: + return httpx.Response(404) + + return httpx.Response(status_code=self._status_code, json=model_dict(found)) + + @staticmethod + def _build( + partial: PartialVectorStoreFile, + request: httpx.Request, + ) -> VectorStoreFile: + raise NotImplementedError + + +class VectorStoreFileDeleteRoute( + StatefulRoute[ + VectorStoreFileDeleted, + PartialResourceDeleted[Literal["vector_store.file.deleted"]], + ] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.delete( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)/files/(?P[a-zA-Z0-9\-]+)" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + vector_store_id = kwargs["vector_store_id"] + found_vector_store = self._state.beta.vector_stores.get(vector_store_id) + if not found_vector_store: + return httpx.Response(404) + + file_id = kwargs["file_id"] + found = self._state.beta.vector_stores.files.get(file_id) + if not found: + return httpx.Response(404) + + deleted = self._state.beta.vector_stores.files.delete(file_id) + + return httpx.Response( + status_code=200, + json=model_dict( + VectorStoreFileDeleted( + id=file_id, + deleted=deleted, + object="vector_store.file.deleted", + ) + ), + ) + + @staticmethod + def _build( + partial: PartialResourceDeleted[Literal["vector_store.file.deleted"]], + request: httpx.Request, + ) -> VectorStoreFileDeleted: + raise NotImplementedError diff --git a/src/openai_responses/_routes/vector_stores.py b/src/openai_responses/_routes/vector_stores.py new file mode 100644 index 0000000..def6ee5 --- /dev/null +++ b/src/openai_responses/_routes/vector_stores.py @@ -0,0 +1,214 @@ +from typing import Any +from typing_extensions import override + +import httpx +import respx + +from openai.pagination import SyncCursorPage +from openai.types.beta.vector_store import VectorStore +from openai.types.beta.vector_store_update_params import VectorStoreUpdateParams +from openai.types.beta.vector_store_deleted import VectorStoreDeleted + +from ._base import StatefulRoute + +from ..stores import StateStore +from .._types.partials.sync_cursor_page import PartialSyncCursorPage +from .._types.partials.vector_stores import ( + PartialVectorStore, + PartialVectorStoreDeleted, +) + +from .._utils.faker import faker +from .._utils.serde import json_loads, model_dict, model_parse +from .._utils.time import utcnow_unix_timestamp_s + +__all__ = [ + "VectorStoreCreateRoute", + "VectorStoreListRoute", + "VectorStoreRetrieveRoute", + "VectorStoreUpdateRoute", + "VectorStoreDeleteRoute", +] + + +class VectorStoreCreateRoute(StatefulRoute[VectorStore, PartialVectorStore]): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.post(url__regex="/vector_stores"), + status_code=201, + state=state, + ) + + @override + def _handler(self, request: httpx.Request, route: respx.Route) -> httpx.Response: + self._route = route + model = self._build({}, request) + self._state.beta.vector_stores.put(model) + return httpx.Response(status_code=self._status_code, json=model_dict(model)) + + @staticmethod + def _build(partial: PartialVectorStore, request: httpx.Request) -> VectorStore: + content = json_loads(request.content) + defaults: PartialVectorStore = { + "id": faker.beta.vector_store.id(), + "created_at": utcnow_unix_timestamp_s(), + "file_counts": { + "cancelled": 0, + "completed": 0, + "failed": 0, + "in_progress": 0, + "total": 0, + }, + "name": "", + "object": "vector_store", + "status": "completed", + "usage_bytes": 0, + } + return model_parse(VectorStore, defaults | partial | content) + + +class VectorStoreListRoute( + StatefulRoute[ + SyncCursorPage[VectorStore], PartialSyncCursorPage[PartialVectorStore] + ] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.get(url__regex="/vector_stores"), + status_code=200, + state=state, + ) + + @override + def _handler(self, request: httpx.Request, route: respx.Route) -> httpx.Response: + self._route = route + + limit = request.url.params.get("limit") + order = request.url.params.get("order") + after = request.url.params.get("after") + before = request.url.params.get("before") + + data = self._state.beta.vector_stores.list(limit, order, after, before) + result_count = len(data) + total_count = len(self._state.beta.vector_stores.list()) + has_data = bool(result_count) + has_more = total_count != result_count + first_id = data[0].id if has_data else None + last_id = data[-1].id if has_data else None + model = SyncCursorPage[VectorStore](data=data) + return httpx.Response( + status_code=200, + json=model_dict(model) + | {"first_id": first_id, "last_id": last_id, "has_more": has_more}, + ) + + @staticmethod + def _build( + partial: PartialSyncCursorPage[PartialVectorStore], + request: httpx.Request, + ) -> SyncCursorPage[VectorStore]: + raise NotImplementedError + + +class VectorStoreRetrieveRoute(StatefulRoute[VectorStore, PartialVectorStore]): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.get( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + vector_store_id = kwargs["vector_store_id"] + found = self._state.beta.vector_stores.get(vector_store_id) + if not found: + return httpx.Response(404) + + return httpx.Response(status_code=self._status_code, json=model_dict(found)) + + @staticmethod + def _build(partial: PartialVectorStore, request: httpx.Request) -> VectorStore: + raise NotImplementedError + + +class VectorStoreUpdateRoute(StatefulRoute[VectorStore, PartialVectorStore]): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.post( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + vector_store_id = kwargs["vector_store_id"] + found = self._state.beta.vector_stores.get(vector_store_id) + if not found: + return httpx.Response(404) + + content: VectorStoreUpdateParams = json_loads(request.content) + deserialized = model_dict(found) + updated = model_parse(VectorStore, deserialized | content) + return httpx.Response(status_code=self._status_code, json=model_dict(updated)) + + @staticmethod + def _build(partial: PartialVectorStore, request: httpx.Request) -> VectorStore: + raise NotImplementedError + + +class VectorStoreDeleteRoute( + StatefulRoute[VectorStoreDeleted, PartialVectorStoreDeleted] +): + def __init__(self, router: respx.MockRouter, state: StateStore) -> None: + super().__init__( + route=router.delete( + url__regex=r"/vector_stores/(?P[a-zA-Z0-9\_]+)" + ), + status_code=200, + state=state, + ) + + @override + def _handler( + self, + request: httpx.Request, + route: respx.Route, + **kwargs: Any, + ) -> httpx.Response: + self._route = route + vector_store_id = kwargs["vector_store_id"] + deleted = self._state.beta.vector_stores.delete(vector_store_id) + return httpx.Response( + status_code=200, + json=model_dict( + VectorStoreDeleted( + id=vector_store_id, + deleted=deleted, + object="vector_store.deleted", + ) + ), + ) + + @staticmethod + def _build( + partial: PartialVectorStoreDeleted, + request: httpx.Request, + ) -> VectorStoreDeleted: + raise NotImplementedError diff --git a/src/openai_responses/_types/partials/assistants.py b/src/openai_responses/_types/partials/assistants.py index ab7af10..7bfc210 100644 --- a/src/openai_responses/_types/partials/assistants.py +++ b/src/openai_responses/_types/partials/assistants.py @@ -3,7 +3,7 @@ from openai._utils._transform import PropertyInfo -__all__ = ["PartialAssistant", "PartialAssistantList", "PartialAssistantDeleted"] +__all__ = ["PartialAssistant", "PartialAssistantDeleted"] class PartialAssistantResponseFormat(TypedDict): @@ -71,14 +71,6 @@ class PartialAssistant(TypedDict): top_p: NotRequired[float] -class PartialAssistantList(TypedDict): - object: NotRequired[Literal["list"]] - data: NotRequired[List[PartialAssistant]] - first_id: NotRequired[str] - last_id: NotRequired[str] - has_more: NotRequired[bool] - - class PartialAssistantDeleted(TypedDict): id: NotRequired[str] object: NotRequired[Literal["assistant.deleted"]] diff --git a/src/openai_responses/_types/partials/deleted.py b/src/openai_responses/_types/partials/deleted.py new file mode 100644 index 0000000..92dd428 --- /dev/null +++ b/src/openai_responses/_types/partials/deleted.py @@ -0,0 +1,12 @@ +from typing import Generic, TypeVar +from typing_extensions import NotRequired, TypedDict + +S = TypeVar("S") + +__all__ = ["PartialResourceDeleted"] + + +class PartialResourceDeleted(TypedDict, Generic[S]): + id: NotRequired[str] + object: NotRequired[S] + deleted: NotRequired[bool] diff --git a/src/openai_responses/_types/partials/messages.py b/src/openai_responses/_types/partials/messages.py index df0bf21..4601d43 100644 --- a/src/openai_responses/_types/partials/messages.py +++ b/src/openai_responses/_types/partials/messages.py @@ -3,7 +3,7 @@ from openai._utils import PropertyInfo -__all__ = ["PartialMessage", "PartialMessageList", "PartialMessageDeleted"] +__all__ = ["PartialMessage", "PartialMessageDeleted"] class PartialFileSearchTool(TypedDict): @@ -91,14 +91,6 @@ class PartialMessage(TypedDict): thread_id: NotRequired[str] -class PartialMessageList(TypedDict): - object: NotRequired[Literal["list"]] - data: NotRequired[List[PartialMessage]] - first_id: NotRequired[str] - last_id: NotRequired[str] - has_more: NotRequired[bool] - - class PartialMessageDeleted(TypedDict): id: NotRequired[str] object: NotRequired[Literal["thread.message.deleted"]] diff --git a/src/openai_responses/_types/partials/models.py b/src/openai_responses/_types/partials/models.py index 2c37885..171acd8 100644 --- a/src/openai_responses/_types/partials/models.py +++ b/src/openai_responses/_types/partials/models.py @@ -1,4 +1,4 @@ -from typing import List, Literal, TypedDict +from typing import Literal, TypedDict from typing_extensions import NotRequired @@ -7,8 +7,3 @@ class PartialModel(TypedDict): created: NotRequired[int] object: NotRequired[Literal["model"]] owned_by: NotRequired[str] - - -class PartialModelList(TypedDict): - object: NotRequired[Literal["list"]] - data: List[PartialModel] diff --git a/src/openai_responses/_types/partials/run_steps.py b/src/openai_responses/_types/partials/run_steps.py index 0ef39a7..1a5d2bf 100644 --- a/src/openai_responses/_types/partials/run_steps.py +++ b/src/openai_responses/_types/partials/run_steps.py @@ -3,7 +3,7 @@ from openai._utils import PropertyInfo -__all__ = ["PartialRunStep", "PartialRunStepList"] +__all__ = ["PartialRunStep"] class PartialLastError(TypedDict): @@ -111,11 +111,3 @@ class PartialRunStep(TypedDict): thread_id: str type: Literal["message_creation", "tool_calls"] usage: NotRequired[PartialUsage] - - -class PartialRunStepList(TypedDict): - object: NotRequired[Literal["list"]] - data: NotRequired[List[PartialRunStep]] - first_id: NotRequired[str] - last_id: NotRequired[str] - has_more: NotRequired[bool] diff --git a/src/openai_responses/_types/partials/runs.py b/src/openai_responses/_types/partials/runs.py index aea8f4b..2709195 100644 --- a/src/openai_responses/_types/partials/runs.py +++ b/src/openai_responses/_types/partials/runs.py @@ -3,7 +3,7 @@ from openai._utils import PropertyInfo -__all__ = ["PartialRun", "PartialRunList"] +__all__ = ["PartialRun"] class PartialIncompleteDetails(TypedDict): @@ -131,11 +131,3 @@ class PartialRun(TypedDict): usage: NotRequired[PartialUsage] temperature: NotRequired[float] top_p: NotRequired[float] - - -class PartialRunList(TypedDict): - object: NotRequired[Literal["list"]] - data: NotRequired[List[PartialRun]] - first_id: NotRequired[str] - last_id: NotRequired[str] - has_more: NotRequired[bool] diff --git a/src/openai_responses/_types/partials/sync_cursor_page.py b/src/openai_responses/_types/partials/sync_cursor_page.py new file mode 100644 index 0000000..8556509 --- /dev/null +++ b/src/openai_responses/_types/partials/sync_cursor_page.py @@ -0,0 +1,14 @@ +from typing import Generic, List, Literal, TypeVar +from typing_extensions import NotRequired, TypedDict + +M = TypeVar("M") + +__all__ = ["PartialSyncCursorPage"] + + +class PartialSyncCursorPage(TypedDict, Generic[M]): + object: NotRequired[Literal["list"]] + data: NotRequired[List[M]] + first_id: NotRequired[str] + last_id: NotRequired[str] + has_more: NotRequired[bool] diff --git a/src/openai_responses/_types/partials/vector_store_files.py b/src/openai_responses/_types/partials/vector_store_files.py new file mode 100644 index 0000000..3998a85 --- /dev/null +++ b/src/openai_responses/_types/partials/vector_store_files.py @@ -0,0 +1,25 @@ +from typing import Literal, TypedDict + +from typing_extensions import NotRequired + +__all__ = ["PartialVectorStoreFile"] + + +class PartialLastError(TypedDict): + code: Literal[ + "internal_error", + "file_not_found", + "parsing_error", + "unhandled_mime_type", + ] + message: str + + +class PartialVectorStoreFile(TypedDict): + id: NotRequired[str] + created_at: NotRequired[int] + last_error: NotRequired[PartialLastError] + object: NotRequired[Literal["vector_store.file"]] + status: NotRequired[Literal["in_progress", "completed", "cancelled", "failed"]] + usage_bytes: NotRequired[int] + vector_store_id: NotRequired[str] diff --git a/src/openai_responses/_types/partials/vector_stores.py b/src/openai_responses/_types/partials/vector_stores.py new file mode 100644 index 0000000..af9869a --- /dev/null +++ b/src/openai_responses/_types/partials/vector_stores.py @@ -0,0 +1,37 @@ +from typing import Dict, Literal, TypedDict +from typing_extensions import NotRequired + +__all__ = ["PartialVectorStore", "PartialVectorStoreDeleted"] + + +class PartialFileCount(TypedDict): + cancelled: int + completed: int + failed: int + in_progress: int + total: int + + +class PartialExpiresAfter(TypedDict): + anchor: Literal["last_active_at"] + days: int + + +class PartialVectorStore(TypedDict): + id: NotRequired[str] + created_at: NotRequired[int] + file_counts: NotRequired[PartialFileCount] + last_active_at: NotRequired[int] + metadata: NotRequired[Dict[str, str]] + name: NotRequired[str] + object: NotRequired[Literal["vector_store"]] + status: NotRequired[Literal["expired", "in_progress", "completed"]] + usage_bytes: NotRequired[int] + expires_after: NotRequired[PartialExpiresAfter] + expires_at: NotRequired[int] + + +class PartialVectorStoreDeleted(TypedDict): + id: NotRequired[str] + object: NotRequired[Literal["vector_store.deleted"]] + deleted: NotRequired[bool] diff --git a/src/openai_responses/stores/state_store.py b/src/openai_responses/stores/state_store.py index a308382..f90bbf6 100644 --- a/src/openai_responses/stores/state_store.py +++ b/src/openai_responses/stores/state_store.py @@ -1,32 +1,36 @@ from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union from openai.types import FileObject, Model + from openai.types.beta.assistant import Assistant from openai.types.beta.thread import Thread from openai.types.beta.threads.message import Message from openai.types.beta.threads.run import Run from openai.types.beta.threads.runs.run_step import RunStep +from openai.types.beta.vector_store import VectorStore +from openai.types.beta.vector_stores.vector_store_file import VectorStoreFile + from .content_store import ContentStore from .._constants import SYSTEM_MODELS from .._utils.serde import model_parse __all__ = ["StateStore"] -M = TypeVar( - "M", - bound=Union[ - FileObject, - Model, - Assistant, - Thread, - Message, - Run, - RunStep, - ], -) +AnyModel = Union[ + FileObject, + Assistant, + Thread, + Message, + Run, + RunStep, + Model, + VectorStore, + VectorStoreFile, +] -Resource = Union[FileObject, Assistant, Thread, Message, Run, RunStep, Model, Any] + +M = TypeVar("M", bound=AnyModel) class StateStore: @@ -35,7 +39,7 @@ def __init__(self) -> None: self.models = ModelStore() self.beta = Beta() - def _blind_put(self, resource: Resource) -> None: + def _blind_put(self, resource: Union[AnyModel, Any]) -> None: if isinstance(resource, FileObject): self.files.put(resource) elif isinstance(resource, Assistant): @@ -50,6 +54,10 @@ def _blind_put(self, resource: Resource) -> None: self.beta.threads.runs.steps.put(resource) elif isinstance(resource, Model): self.models.put(resource) + elif isinstance(resource, VectorStore): + self.beta.vector_stores.put(resource) + elif isinstance(resource, VectorStoreFile): + self.beta.vector_stores.files.put(resource) else: raise TypeError(f"Cannot put object of type {type(resource)} in store") @@ -58,6 +66,7 @@ class Beta: def __init__(self) -> None: self.assistants = AssistantStore() self.threads = ThreadStore() + self.vector_stores = VectorStoreStore() class BaseStore(Generic[M]): @@ -242,3 +251,71 @@ def list( objs = objs[start_ix:end_ix] return objs[: int(limit)] + + +class VectorStoreStore(BaseStore[VectorStore]): + def __init__(self) -> None: + super().__init__() + self.files = VectorStoreFileStore() + + def list( + self, + limit: Optional[str] = None, + order: Optional[str] = None, + after: Optional[str] = None, + before: Optional[str] = None, + ) -> List[VectorStore]: + limit = limit or "20" + objs = list(self._data.values()) + objs = list(reversed(objs)) if (order or "desc") == "desc" else objs + + start_ix = 0 + if after: + obj = self._data.get(after) + if obj: + start_ix = objs.index(obj) + 1 + + end_ix = None + if before: + obj = self._data.get(before) + if obj: + end_ix = objs.index(obj) + + objs = objs[start_ix:end_ix] + return objs[: int(limit)] + + +class VectorStoreFileStore(BaseStore[VectorStoreFile]): + def list( + self, + vector_store_id: str, + limit: Optional[str] = None, + order: Optional[str] = None, + after: Optional[str] = None, + before: Optional[str] = None, + filter: Optional[ + Literal["in_progress", "completed", "failed", "cancelled"] + ] = None, + ) -> List[VectorStoreFile]: + limit = limit or "20" + objs = [ + m for m in list(self._data.values()) if m.vector_store_id == vector_store_id + ] + if filter: + objs = [obj for obj in objs if obj.status == filter] + objs = list(reversed(objs)) if (order or "desc") == "desc" else objs + + start_ix = 0 + if after: + obj = self._data.get(after) + if obj: + start_ix = objs.index(obj) + 1 + + end_ix = None + if before: + obj = self._data.get(before) + if obj: + end_ix = objs.index(obj) + + objs = objs[start_ix:end_ix] + return objs[: int(limit)]