openaivec 0.13.2__tar.gz → 0.13.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. openaivec-0.13.4/.github/copilot-instructions.md +199 -0
  2. {openaivec-0.13.2 → openaivec-0.13.4}/PKG-INFO +4 -2
  3. {openaivec-0.13.2 → openaivec-0.13.4}/README.md +1 -1
  4. {openaivec-0.13.2 → openaivec-0.13.4}/mkdocs.yml +1 -0
  5. {openaivec-0.13.2 → openaivec-0.13.4}/pyproject.toml +22 -0
  6. openaivec-0.13.4/src/openaivec/__init__.py +9 -0
  7. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/di.py +3 -3
  8. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/embeddings.py +5 -4
  9. openaivec-0.13.4/src/openaivec/optimize.py +108 -0
  10. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/pandas_ext.py +129 -21
  11. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/prompt.py +34 -13
  12. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/provider.py +3 -3
  13. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/proxy.py +207 -86
  14. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/responses.py +6 -5
  15. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/serialize.py +1 -1
  16. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/spark.py +8 -7
  17. openaivec-0.13.4/src/openaivec/task/customer_support/__init__.py +26 -0
  18. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/customer_sentiment.py +12 -4
  19. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/inquiry_classification.py +11 -4
  20. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/inquiry_summary.py +8 -3
  21. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/intent_analysis.py +10 -4
  22. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/response_suggestion.py +10 -4
  23. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/customer_support/urgency_analysis.py +8 -3
  24. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/dependency_parsing.py +4 -2
  25. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/keyword_extraction.py +3 -2
  26. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/morphological_analysis.py +4 -2
  27. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/named_entity_recognition.py +4 -2
  28. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/sentiment_analysis.py +7 -2
  29. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/translation.py +1 -1
  30. openaivec-0.13.4/src/openaivec/task/table/__init__.py +3 -0
  31. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/table/fillna.py +4 -3
  32. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/util.py +0 -1
  33. openaivec-0.13.4/tests/test_optimize.py +318 -0
  34. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_pandas_ext.py +4 -2
  35. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_prompt.py +44 -0
  36. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_provider.py +1 -0
  37. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_proxy.py +254 -6
  38. openaivec-0.13.4/tests/test_proxy_suggester.py +201 -0
  39. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_util.py +2 -1
  40. {openaivec-0.13.2 → openaivec-0.13.4}/uv.lock +1258 -1207
  41. openaivec-0.13.2/src/openaivec/__init__.py +0 -9
  42. openaivec-0.13.2/src/openaivec/task/customer_support/__init__.py +0 -32
  43. openaivec-0.13.2/src/openaivec/task/table/__init__.py +0 -3
  44. {openaivec-0.13.2 → openaivec-0.13.4}/.env.example +0 -0
  45. {openaivec-0.13.2 → openaivec-0.13.4}/.github/workflows/python-mkdocs.yml +0 -0
  46. {openaivec-0.13.2 → openaivec-0.13.4}/.github/workflows/python-package.yml +0 -0
  47. {openaivec-0.13.2 → openaivec-0.13.4}/.github/workflows/python-test.yml +0 -0
  48. {openaivec-0.13.2 → openaivec-0.13.4}/.github/workflows/python-update.yml +0 -0
  49. {openaivec-0.13.2 → openaivec-0.13.4}/.gitignore +0 -0
  50. {openaivec-0.13.2 → openaivec-0.13.4}/CODE_OF_CONDUCT.md +0 -0
  51. {openaivec-0.13.2 → openaivec-0.13.4}/LICENSE +0 -0
  52. {openaivec-0.13.2 → openaivec-0.13.4}/SECURITY.md +0 -0
  53. {openaivec-0.13.2 → openaivec-0.13.4}/SUPPORT.md +0 -0
  54. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/di.md +0 -0
  55. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/embeddings.md +0 -0
  56. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/pandas_ext.md +0 -0
  57. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/prompt.md +0 -0
  58. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/proxy.md +0 -0
  59. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/responses.md +0 -0
  60. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/spark.md +0 -0
  61. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/task.md +0 -0
  62. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/customer_sentiment.md +0 -0
  63. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/inquiry_classification.md +0 -0
  64. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/inquiry_summary.md +0 -0
  65. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/intent_analysis.md +0 -0
  66. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/response_suggestion.md +0 -0
  67. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/customer_support/urgency_analysis.md +0 -0
  68. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/dependency_parsing.md +0 -0
  69. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/keyword_extraction.md +0 -0
  70. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/morphological_analysis.md +0 -0
  71. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/named_entity_recognition.md +0 -0
  72. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/sentiment_analysis.md +0 -0
  73. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/tasks/nlp/translation.md +0 -0
  74. {openaivec-0.13.2 → openaivec-0.13.4}/docs/api/util.md +0 -0
  75. {openaivec-0.13.2 → openaivec-0.13.4}/docs/index.md +0 -0
  76. {openaivec-0.13.2 → openaivec-0.13.4}/docs/robots.txt +0 -0
  77. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/log.py +0 -0
  78. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/model.py +0 -0
  79. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/__init__.py +0 -0
  80. {openaivec-0.13.2 → openaivec-0.13.4}/src/openaivec/task/nlp/__init__.py +3 -3
  81. {openaivec-0.13.2 → openaivec-0.13.4}/tests/__init__.py +0 -0
  82. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_di.py +0 -0
  83. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_embeddings.py +0 -0
  84. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_responses.py +0 -0
  85. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_serialize.py +0 -0
  86. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_spark.py +4 -4
  87. {openaivec-0.13.2 → openaivec-0.13.4}/tests/test_task.py +0 -0
@@ -0,0 +1,199 @@
1
+ # Copilot instructions for openaivec
2
+
3
+ This repository-wide guide tells GitHub Copilot how to propose code that fits our architecture, APIs, style, tests, and docs. Prefer these rules when completing or generating code.
4
+
5
+ ## Project overview
6
+
7
+ - Goal: Provide a vectorized (batched) interface to OpenAI/Azure OpenAI so pandas/Spark can process large text corpora with high throughput.
8
+ - Public API exports (`src/openaivec/__init__.py`):
9
+ - Responses: `BatchResponses`, `AsyncBatchResponses`
10
+ - Embeddings: `BatchEmbeddings`, `AsyncBatchEmbeddings`
11
+ - First-class pandas extensions (`.ai` / `.aio`) and Spark UDF builders
12
+ - Azure OpenAI is supported with the same APIs (use deployment name as the “model” for Azure)
13
+
14
+ ## Architecture and roles
15
+
16
+ - `src/openaivec/proxy.py`
17
+ - Core batching, deduplication, order preservation, and caching
18
+ - `BatchingMapProxy[S, T]` (sync) / `AsyncBatchingMapProxy[S, T]` (async)
19
+ - The map_func contract is strict: return a list of the same length and order as the inputs
20
+ - Progress bars only in notebook environments via `tqdm.auto`, gated by `show_progress=True`
21
+ - `src/openaivec/responses.py`
22
+ - Batched wrapper over OpenAI Responses JSON-mode API
23
+ - `BatchResponses` / `AsyncBatchResponses` use the proxy internally
24
+ - Retries via `backoff`/`backoff_async` for transient errors (RateLimit, 5xx)
25
+ - Reasoning models (o1/o3 family) must use `temperature=None`; helpful guidance on errors
26
+ - `src/openaivec/embeddings.py`
27
+ - Batched embeddings (sync/async)
28
+ - `src/openaivec/pandas_ext.py`
29
+ - `Series.ai` / `Series.aio` entry points for responses/embeddings
30
+ - Uses DI container (`provider.CONTAINER`) to get client and model names
31
+ - Supports batch size, progress, and cache sharing (`*_with_cache`)
32
+ - `src/openaivec/spark.py`
33
+ - UDF builders: `responses_udf` / `task_udf` / `embeddings_udf` / `count_tokens_udf` / `split_to_chunks_udf`
34
+ - Per-partition duplicate caching to reduce API calls
35
+ - Pydantic → Spark StructType schema conversion
36
+ - `src/openaivec/provider.py`
37
+ - DI container and automatic OpenAI/Azure OpenAI client provisioning
38
+ - Warns if Azure base URL isn’t v1 format
39
+ - `src/openaivec/util.py`
40
+ - `backoff` / `backoff_async` and `TextChunker`
41
+ - Additional modules from CLAUDE.md
42
+ - `src/openaivec/di.py`: lightweight DI container
43
+ - `src/openaivec/log.py`: logging/observe helpers
44
+ - `src/openaivec/prompt.py`: few-shot prompt building
45
+ - `src/openaivec/serialize.py`: Pydantic schema (de)serialization
46
+ - `src/openaivec/task/`: pre-built, structured task library
47
+
48
+ ## Dev commands (uv)
49
+
50
+ ```bash
51
+ # Install all dependencies (dev + extras)
52
+ uv sync --all-extras --dev
53
+
54
+ # Install in editable mode
55
+ uv pip install -e .
56
+
57
+ # Lint and format
58
+ uv run ruff check . --fix && uv run ruff format .
59
+
60
+ # Run tests
61
+ uv run pytest
62
+
63
+ # Build/serve docs
64
+ uv run mkdocs serve
65
+ ```
66
+
67
+ ## Coding standards (Ruff/types/style)
68
+
69
+ - Python ≥ 3.10
70
+ - Lint/format via Ruff (line-length=120, target=py310)
71
+ - Imports: absolute only (enforced by TID252), except `__init__.py` may re-export relatively
72
+ - Type hints required using modern syntax (`str | None` over `Optional[str]`)
73
+ - Public APIs should document return values and exceptions (Google-style docstrings preferred)
74
+ - Favor `@dataclass` for simple data contracts; separate mutable state cleanly
75
+ - Don’t swallow errors broadly; raise `ValueError` etc. on contract violations
76
+
77
+ ## API contracts and critical rules
78
+
79
+ - Proxy (BatchingMapProxy / AsyncBatchingMapProxy)
80
+ - map_func must return a list with the same length and order as inputs; on mismatch, release events and raise ValueError
81
+ - Inputs are de-duplicated while preserving first-occurrence order; outputs are restored to the original order
82
+ - Progress is only shown in notebooks when `show_progress=True`
83
+ - Async version enforces `max_concurrency` via `asyncio.Semaphore`
84
+ - Responses
85
+ - Use OpenAI Responses JSON mode (`responses.parse`)
86
+ - For reasoning models (o1/o3 families), you MUST set `temperature=None`; helpful error messaging is built-in
87
+ - Strongly prefer structured outputs with Pydantic models
88
+ - Retries with exponential backoff for RateLimit/5xx
89
+ - Embeddings
90
+ - Return NumPy float32 arrays
91
+ - pandas extensions
92
+ - `.ai.responses` / `.ai.embeddings` strictly preserve Series index and length
93
+ - `.aio` provides async variants; tune with `max_concurrency` and `batch_size`
94
+ - `*_with_cache` variants let callers share external caches across ops
95
+ - Spark UDFs
96
+ - Cache duplicates within each partition to minimize API cost
97
+ - Convert Pydantic models to Spark schemas; treat Enum/Literal as strings
98
+ - Reasoning models require `temperature=None`
99
+ - Provide token counting and text chunking helpers
100
+ - Provider/DI and Azure
101
+ - Auto-detect OpenAI vs Azure OpenAI from env vars
102
+ - Azure requires v1 base URL (warn otherwise) and uses deployment name as the “model”
103
+
104
+ ## Preferred patterns (Do) and Avoid (Don’t)
105
+
106
+ Do
107
+
108
+ - Batch through the Proxy rather than per-item loops
109
+ - Attach `backoff`/`backoff_async` to external API calls (RateLimit, 5xx)
110
+ - Preserve index/order/schema for pandas/Spark APIs
111
+ - Clarify Azure specifics (“deployment name” vs “model name”); respect `_check_azure_v1_api_url`
112
+ - When changing public APIs, update `__all__` and docs in `docs/`
113
+
114
+ Don’t
115
+
116
+ - Break the Proxy contract (same-length, ordered result)
117
+ - Fire one API request per item—always batch via the Proxy
118
+ - Show progress outside notebook contexts or ignore `show_progress`
119
+ - Use relative imports (except `__init__.py` re-exports)
120
+ - Hit real external APIs in unit tests (prefer mocks/stubs)
121
+
122
+ ## Performance guidance
123
+
124
+ - Typical batch size ranges: Responses 32–128, Embeddings 64–256 (defaults are 128 in code)
125
+ - Async `max_concurrency` commonly 4–12 per process/partition; scale with rate limits in mind
126
+ - Partition-level caching (Spark) and cross-op cache sharing (pandas `*_with_cache`) greatly reduce costs
127
+
128
+ ## Testing strategy (pytest)
129
+
130
+ - Tests live in `tests/`; cover both sync and async where applicable
131
+ - Prefer mocks/stubs for external API calls; keep data small and deterministic
132
+ - Focus areas:
133
+ - Order/length preservation
134
+ - Deduplication and cache reuse
135
+ - Event release on exceptions (deadlock prevention)
136
+ - `max_concurrency` is not exceeded
137
+ - Reasoning model guidance (`temperature=None`)
138
+ - Use `asyncio.run` in async tests (mirrors existing tests)
139
+ - Optional integration tests can run with valid API keys; keep unit tests independent of network
140
+
141
+ ## Documentation (MkDocs)
142
+
143
+ - For new developer-facing APIs, update `docs/api/` and consider a short example under `docs/examples/`
144
+ - Keep pandas/Spark examples concise to minimize learning curve
145
+ - Update `mkdocs.yml` navigation when adding modules or examples
146
+
147
+ ## PR checklist (pre-merge)
148
+
149
+ - [ ] Ruff check/format passes (line-length 120, absolute imports)
150
+ - [ ] Public API contracts (order/length/types) are satisfied
151
+ - [ ] Large-scale processing is batched via the Proxy
152
+ - [ ] Reasoning models use `temperature=None` where applicable
153
+ - [ ] Tests added/updated without calling live external APIs
154
+ - [ ] Docs updated if needed (`docs/` and/or examples)
155
+
156
+ ## Common snippets (what to suggest)
157
+
158
+ - New batched API wrapper (sync)
159
+
160
+ ```python
161
+ @observe(_LOGGER)
162
+ @backoff(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
163
+ def _unit_of_work(self, xs: list[str]) -> list[TOut]:
164
+ resp = self.client.api(xs) # real API call
165
+ return convert(resp) # same length/order as xs
166
+
167
+ def create(self, inputs: list[str]) -> list[TOut]:
168
+ return self.cache.map(inputs, self._unit_of_work)
169
+ ```
170
+
171
+ - Reasoning model temperature
172
+ ```python
173
+ # o1/o3 and similar reasoning models must use None
174
+ temperature=None
175
+ ```
176
+ - pandas `.ai` with shared cache
177
+
178
+ ```python
179
+ from openaivec.proxy import BatchingMapProxy
180
+
181
+ shared = BatchingMapProxy[str, str](batch_size=64)
182
+ df["text"].ai.responses_with_cache("instructions", cache=shared)
183
+ ```
184
+
185
+ - Spark UDF (structured output)
186
+
187
+ ```python
188
+ from pydantic import BaseModel
189
+ from openaivec.spark import responses_udf
190
+
191
+ class R(BaseModel):
192
+ value: str
193
+
194
+ udf = responses_udf("do something", response_format=R, batch_size=64, max_concurrency=8)
195
+ ```
196
+
197
+ ---
198
+
199
+ By following this guide, Copilot suggestions will match the repository’s design, performance goals, and testing standards. When in doubt, read the implementations in `proxy.py`, `responses.py`, `pandas_ext.py`, and `spark.py`, and the tests under `tests/`.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: openaivec
3
- Version: 0.13.2
3
+ Version: 0.13.4
4
4
  Summary: Generative mutation for tabular calculation
5
5
  Project-URL: Homepage, https://microsoft.github.io/openaivec/
6
6
  Project-URL: Repository, https://github.com/microsoft/openaivec
@@ -15,9 +15,11 @@ Classifier: Programming Language :: Python :: 3.10
15
15
  Classifier: Programming Language :: Python :: 3.11
16
16
  Classifier: Programming Language :: Python :: 3.12
17
17
  Requires-Python: >=3.10
18
+ Requires-Dist: ipywidgets>=8.1.7
18
19
  Requires-Dist: openai>=1.74.0
19
20
  Requires-Dist: pandas>=2.2.3
20
21
  Requires-Dist: tiktoken>=0.9.0
22
+ Requires-Dist: tqdm>=4.67.1
21
23
  Provides-Extra: spark
22
24
  Requires-Dist: pyspark>=3.5.5; extra == 'spark'
23
25
  Description-Content-Type: text/markdown
@@ -590,7 +592,7 @@ improved_prompt: str = (
590
592
  .example("Apple", "Color")
591
593
  .example("Apple", "Animal")
592
594
  # improve the prompt with OpenAI's API
593
- .improve(client, model_name)
595
+ .improve()
594
596
  .build()
595
597
  )
596
598
  print(improved_prompt)
@@ -566,7 +566,7 @@ improved_prompt: str = (
566
566
  .example("Apple", "Color")
567
567
  .example("Apple", "Animal")
568
568
  # improve the prompt with OpenAI's API
569
- .improve(client, model_name)
569
+ .improve()
570
570
  .build()
571
571
  )
572
572
  print(improved_prompt)
@@ -61,6 +61,7 @@ nav:
61
61
  - Async Workflows: examples/aio.ipynb
62
62
  - Prompt Engineering: examples/prompt.ipynb
63
63
  - FAQ Generation: examples/generate_faq.ipynb
64
+ - Token Count and Processing Time: examples/batch_size.ipynb
64
65
  - API Reference:
65
66
  - di: api/di.md
66
67
  - pandas_ext: api/pandas_ext.md
@@ -26,9 +26,11 @@ classifiers = [
26
26
 
27
27
  requires-python = ">=3.10"
28
28
  dependencies = [
29
+ "ipywidgets>=8.1.7",
29
30
  "openai>=1.74.0",
30
31
  "pandas>=2.2.3",
31
32
  "tiktoken>=0.9.0",
33
+ "tqdm>=4.67.1",
32
34
  ]
33
35
 
34
36
  [dependency-groups]
@@ -62,6 +64,26 @@ spark = [
62
64
  line-length = 120
63
65
  target-version = "py310"
64
66
 
67
+ [tool.ruff.lint]
68
+ select = [
69
+ "E", # pycodestyle errors
70
+ "W", # pycodestyle warnings
71
+ "F", # pyflakes
72
+ "I", # isort
73
+ "TID", # flake8-tidy-imports
74
+ ]
75
+ # ignore = [] # グローバルではE501を有効化
76
+
77
+ [tool.ruff.lint.flake8-tidy-imports]
78
+ # Enforce absolute imports - ban relative imports (except in __init__.py files)
79
+ ban-relative-imports = "all"
80
+
81
+ [tool.ruff.lint.per-file-ignores]
82
+ # Allow relative imports in __init__.py files
83
+ "**/__init__.py" = ["TID252"]
84
+ # Test files contain long test data - ignore line length
85
+ "tests/**/*.py" = ["E501"]
86
+
65
87
  [project.urls]
66
88
  Homepage = "https://microsoft.github.io/openaivec/"
67
89
  Repository = "https://github.com/microsoft/openaivec"
@@ -0,0 +1,9 @@
1
+ from .embeddings import AsyncBatchEmbeddings, BatchEmbeddings
2
+ from .responses import AsyncBatchResponses, BatchResponses
3
+
4
+ __all__ = [
5
+ "BatchResponses",
6
+ "AsyncBatchResponses",
7
+ "BatchEmbeddings",
8
+ "AsyncBatchEmbeddings",
9
+ ]
@@ -11,14 +11,14 @@ are created once and reused across multiple resolve calls.
11
11
  Example:
12
12
  ```python
13
13
  from openaivec.di import Container
14
-
14
+
15
15
  class DatabaseService:
16
16
  def __init__(self):
17
17
  self.connection = "database://localhost"
18
-
18
+
19
19
  container = Container()
20
20
  container.register(DatabaseService, lambda: DatabaseService())
21
-
21
+
22
22
  db1 = container.resolve(DatabaseService)
23
23
  db2 = container.resolve(DatabaseService)
24
24
  print(db1 is db2) # True - same instance
@@ -6,9 +6,9 @@ import numpy as np
6
6
  from numpy.typing import NDArray
7
7
  from openai import AsyncOpenAI, InternalServerError, OpenAI, RateLimitError
8
8
 
9
- from .log import observe
10
- from .proxy import AsyncBatchingMapProxy, BatchingMapProxy
11
- from .util import backoff, backoff_async
9
+ from openaivec.log import observe
10
+ from openaivec.proxy import AsyncBatchingMapProxy, BatchingMapProxy
11
+ from openaivec.util import backoff, backoff_async
12
12
 
13
13
  __all__ = [
14
14
  "BatchEmbeddings",
@@ -24,7 +24,8 @@ class BatchEmbeddings:
24
24
 
25
25
  Attributes:
26
26
  client (OpenAI): Configured OpenAI client.
27
- model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name (e.g., ``"text-embedding-3-small"``).
27
+ model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name
28
+ (e.g., ``"text-embedding-3-small"``).
28
29
  cache (BatchingMapProxy[str, NDArray[np.float32]]): Batching proxy for ordered, cached mapping.
29
30
  """
30
31
 
@@ -0,0 +1,108 @@
1
+ import threading
2
+ import time
3
+ from contextlib import contextmanager
4
+ from dataclasses import dataclass, field
5
+ from datetime import datetime, timezone
6
+ from typing import List
7
+
8
+
9
+ @dataclass(frozen=True)
10
+ class PerformanceMetric:
11
+ duration: float
12
+ batch_size: int
13
+ executed_at: datetime
14
+ exception: BaseException | None = None
15
+
16
+
17
+ @dataclass
18
+ class BatchSizeSuggester:
19
+ current_batch_size: int = 10
20
+ min_batch_size: int = 10
21
+ min_duration: float = 30.0
22
+ max_duration: float = 60.0
23
+ step_ratio: float = 0.1
24
+ sample_size: int = 10
25
+ _history: List[PerformanceMetric] = field(default_factory=list)
26
+ _lock: threading.RLock = field(default_factory=threading.RLock, repr=False)
27
+ _batch_size_changed_at: datetime | None = field(default=None, init=False)
28
+
29
+ def __post_init__(self) -> None:
30
+ if self.min_batch_size <= 0:
31
+ raise ValueError("min_batch_size must be > 0")
32
+ if self.current_batch_size < self.min_batch_size:
33
+ raise ValueError("current_batch_size must be >= min_batch_size")
34
+ if self.sample_size <= 0:
35
+ raise ValueError("sample_size must be > 0")
36
+ if self.step_ratio <= 0:
37
+ raise ValueError("step_ratio must be > 0")
38
+ if self.min_duration <= 0 or self.max_duration <= 0:
39
+ raise ValueError("min_duration and max_duration must be > 0")
40
+ if self.min_duration >= self.max_duration:
41
+ raise ValueError("min_duration must be < max_duration")
42
+
43
+ @contextmanager
44
+ def record(self, batch_size: int):
45
+ start_time = time.perf_counter()
46
+ executed_at = datetime.now(timezone.utc)
47
+ caught_exception: BaseException | None = None
48
+ try:
49
+ yield
50
+ except BaseException as e:
51
+ caught_exception = e
52
+ raise
53
+ finally:
54
+ duration = time.perf_counter() - start_time
55
+ with self._lock:
56
+ self._history.append(
57
+ PerformanceMetric(
58
+ duration=duration,
59
+ batch_size=batch_size,
60
+ executed_at=executed_at,
61
+ exception=caught_exception,
62
+ )
63
+ )
64
+
65
+ @property
66
+ def samples(self) -> List[PerformanceMetric]:
67
+ with self._lock:
68
+ selected: List[PerformanceMetric] = []
69
+ for metric in reversed(self._history):
70
+ if metric.exception is not None:
71
+ continue
72
+ if self._batch_size_changed_at and metric.executed_at < self._batch_size_changed_at:
73
+ continue
74
+ selected.append(metric)
75
+ if len(selected) >= self.sample_size:
76
+ break
77
+ return list(reversed(selected))
78
+
79
+ def clear_history(self):
80
+ with self._lock:
81
+ self._history.clear()
82
+
83
+ def suggest_batch_size(self) -> int:
84
+ selected = self.samples
85
+
86
+ if len(selected) < self.sample_size:
87
+ with self._lock:
88
+ return self.current_batch_size
89
+
90
+ average_duration = sum(m.duration for m in selected) / len(selected)
91
+
92
+ with self._lock:
93
+ current_size = self.current_batch_size
94
+
95
+ if average_duration < self.min_duration:
96
+ new_batch_size = int(current_size * (1 + self.step_ratio))
97
+ elif average_duration > self.max_duration:
98
+ new_batch_size = int(current_size * (1 - self.step_ratio))
99
+ else:
100
+ new_batch_size = current_size
101
+
102
+ new_batch_size = max(new_batch_size, self.min_batch_size)
103
+
104
+ if new_batch_size != self.current_batch_size:
105
+ self._batch_size_changed_at = datetime.now(timezone.utc)
106
+ self.current_batch_size = new_batch_size
107
+
108
+ return self.current_batch_size