omniai-kit 1.0.1__tar.gz → 1.1.0__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 (47) hide show
  1. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/PKG-INFO +148 -90
  2. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/README.md +147 -89
  3. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/__init__.py +57 -0
  4. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/_version.py +1 -1
  5. omniai_kit-1.1.0/omniai/client.py +536 -0
  6. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/resources/__init__.py +6 -0
  7. omniai_kit-1.1.0/omniai/resources/agents.py +716 -0
  8. omniai_kit-1.1.0/omniai/resources/chat.py +607 -0
  9. omniai_kit-1.1.0/omniai/resources/keys.py +222 -0
  10. omniai_kit-1.1.0/omniai/resources/livekit.py +116 -0
  11. omniai_kit-1.1.0/omniai/resources/models.py +185 -0
  12. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/resources/tasks.py +78 -26
  13. omniai_kit-1.1.0/omniai/types/__init__.py +60 -0
  14. omniai_kit-1.1.0/omniai/types/agents.py +91 -0
  15. omniai_kit-1.1.0/omniai/types/chat.py +139 -0
  16. omniai_kit-1.1.0/omniai/types/keys.py +41 -0
  17. omniai_kit-1.1.0/omniai/types/livekit.py +47 -0
  18. omniai_kit-1.1.0/omniai/types/models.py +51 -0
  19. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai_kit.egg-info/PKG-INFO +148 -90
  20. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai_kit.egg-info/SOURCES.txt +6 -1
  21. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/pyproject.toml +1 -1
  22. omniai_kit-1.1.0/tests/test_v1_features.py +391 -0
  23. omniai_kit-1.0.1/omniai/client.py +0 -316
  24. omniai_kit-1.0.1/omniai/resources/chat.py +0 -403
  25. omniai_kit-1.0.1/omniai/resources/livekit.py +0 -71
  26. omniai_kit-1.0.1/omniai/resources/models.py +0 -63
  27. omniai_kit-1.0.1/omniai/types/__init__.py +0 -27
  28. omniai_kit-1.0.1/omniai/types/chat.py +0 -48
  29. omniai_kit-1.0.1/omniai/types/livekit.py +0 -17
  30. omniai_kit-1.0.1/omniai/types/models.py +0 -30
  31. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/LICENSE +0 -0
  32. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/_base.py +0 -0
  33. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/_constants.py +0 -0
  34. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/exceptions.py +0 -0
  35. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/py.typed +0 -0
  36. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/resources/router.py +0 -0
  37. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai/types/tasks.py +0 -0
  38. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai_kit.egg-info/dependency_links.txt +0 -0
  39. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai_kit.egg-info/requires.txt +0 -0
  40. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/omniai_kit.egg-info/top_level.txt +0 -0
  41. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/setup.cfg +0 -0
  42. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_async_client.py +0 -0
  43. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_client.py +0 -0
  44. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_exceptions.py +0 -0
  45. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_providers.py +0 -0
  46. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_retry.py +0 -0
  47. {omniai_kit-1.0.1 → omniai_kit-1.1.0}/tests/test_router.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: omniai-kit
3
- Version: 1.0.1
3
+ Version: 1.1.0
4
4
  Summary: Official Python client SDK for OmniAI Gateway and Dynamic Multi-Provider LLM Router
5
5
  Author-email: OmniAI Team <dev@omniai.ai>, Jayed Bin Jahangir <jayedbinjahangir@gmail.com>
6
6
  Maintainer-email: Jayed Bin Jahangir <jayedbinjahangir@gmail.com>
@@ -56,19 +56,21 @@ Dynamic: license-file
56
56
  [![Type Checked: PEP 561](https://img.shields.io/badge/typing-PEP%20561-blue.svg)](https://peps.python.org/pep-0561/)
57
57
  [![Status: Production/Stable](https://img.shields.io/badge/status-production-success.svg)](https://omniai.go-agents.app)
58
58
 
59
- Official Python Client SDK for the **OmniAI Platform & AI Gateway** — providing zero-config multi-provider LLM routing, real-time token streaming, native Pydantic structured outputs, background task queues, and WebRTC multimodal live voice/video sessions.
59
+ Official Python Client SDK for the **OmniAI Platform & Universal AI Gateway** — providing zero-config multi-provider LLM routing, real-time SSE token streaming, native Pydantic structured outputs, extended reasoning token budgets (`thinking`), background task queues, Custom AI Agents with RAG knowledge bases, scoped API key management, and real-time WebRTC multimodal live voice/video sessions.
60
60
 
61
61
  ---
62
62
 
63
- ## ⚡ Features
63
+ ## ⚡ Key Capabilities
64
64
 
65
- - **Zero-Configuration Setup:** Automatically loads credentials (`OMNIAI_API_KEY`, `OMNI_API_KEY`, `X_API_KEY`) and defaults to `https://omniai.go-agents.app`.
66
- - **Modern Resource Architecture:** Clean OpenAI-style resource hierarchy (`client.chat.completions.create`, `client.models.list`, `client.tasks.submit`, `client.livekit.create_session`) plus ergonomic top-level shortcuts (`client.chat`, `client.stream`).
65
+ - **Zero-Configuration Setup:** Auto-discovers credentials (`OMNIAI_API_KEY`, `OMNI_API_KEY`, `X_API_KEY`) and defaults to `https://omniai.go-agents.app`.
66
+ - **Modern Resource Architecture:** Clean OpenAI-style resource hierarchy (`client.chat.completions.create`, `client.models.list`, `client.tasks.submit`, `client.livekit.create_session`, `client.agents`, `client.keys`) plus ergonomic top-level shortcuts (`client.chat`, `client.stream`, `client.execute_agent`).
67
+ - **Full Response & Billing Visibility:** Pass `full_response=True` to inspect completion ID, exact micro-dollar request cost (`0.000001`), cumulative API key spend, remaining budget, account tier, and granular prompt token details (`system_tokens`, `input_tokens`).
68
+ - **Extended Reasoning / Thinking Tokens:** Pass `thinking=1024` for reasoning models (e.g. Gemini 3.7 Flash, Claude 3.7 Sonnet, o3-mini) with automatic thinking block handling.
67
69
  - **Synchronous & Asynchronous:** Identical API surfaces for both synchronous (`OmniaiClient`) and async (`AsyncOmniaiClient`) applications.
68
70
  - **Strongly Typed Pydantic Schemas:** Pass Pydantic models directly to `response_model` for guaranteed schema validation and markdown JSON stripping.
69
- - **Resilient Real-Time Streaming:** Server-Sent Events (SSE) streaming engine with auto-reconnect and keep-alive handling.
71
+ - **Custom AI Agents & RAG:** Create autonomous agents, upload documents to Qdrant vector knowledge bases, attach dynamic tools, and execute agents with context.
72
+ - **Scoped API Key Management:** Programmatically provision API keys with spending limits (`total_limit`), expiration dates (`expires_at`), and model access control (`allowed_models`).
70
73
  - **Enterprise Fault Tolerance:** Built-in exponential backoff with full jitter, honoring upstream `Retry-After` headers on 429 and 5xx errors.
71
- - **Background Tasks:** Submit compute-heavy completions to the Celery background queue and poll reliably with `job.wait()`.
72
74
  - **LiveKit Multimodal Sessions:** Generate signed WebRTC access tokens for ultra-low latency Gemini Live audio and video agents.
73
75
 
74
76
  ---
@@ -91,20 +93,20 @@ pip install "omniai-kit[anthropic]" # Installs Anthropic SDK
91
93
 
92
94
  ## 🏛️ Supported Model Providers
93
95
 
94
- OmniAI natively routes and normalizes parameters across the following providers:
96
+ OmniAI natively routes, caches, and normalizes parameters across the following providers:
95
97
 
96
98
  | Provider | Schema | Example Models |
97
99
  | :--- | :--- | :--- |
100
+ | **Google Gemini** | `gemini:<model>` | `gemini:gemini-2.5-flash`, `gemini:gemini-3.7-flash`, `gemini-3.8-live` |
98
101
  | **OpenAI** | `openai:<model>` | `openai:gpt-4o`, `openai:o3-mini`, `openai:gpt-4o-mini` |
99
- | **Google Gemini** | `gemini:<model>` | `gemini:gemini-2.5-flash`, `gemini:gemini-3.7-flash` |
100
- | **Anthropic** | `anthropic:<model>` | `anthropic:claude-3-5-sonnet-20241022`, `anthropic:claude-3-opus` |
102
+ | **Anthropic** | `anthropic:<model>` | `anthropic:claude-3-7-sonnet`, `anthropic:claude-3-5-haiku` |
103
+ | **OpenRouter** | `openrouter:<model>` | `openrouter:inclusionai/ling-3.2-flash:free`, `openrouter:meta-llama/llama-3.3-70b-instruct` |
101
104
  | **vLLM (Self-Hosted)** | `vllm:<model>` | `vllm:meta-llama/Llama-3.1-8B-Instruct` |
102
- | **OpenRouter** | `openrouter:<model>` | `openrouter:anthropic/claude-3.5-sonnet` |
103
105
  | **Ollama (Local)** | `ollama:<model>` | `ollama:qwen2.5-coder:7b` |
104
106
 
105
107
  ---
106
108
 
107
- ## 🚀 Quickstart
109
+ ## 🚀 Usage Guide
108
110
 
109
111
  ### 1. Basic Chat Completion
110
112
 
@@ -132,7 +134,52 @@ response = client.chat.completions.create(model="gemini:gemini-2.5-flash", promp
132
134
 
133
135
  ---
134
136
 
135
- ### 2. Real-Time Token Streaming
137
+ ### 2. Full Response Envelope & Billing Inspection
138
+
139
+ Pass `full_response=True` to inspect exact token counts, cache hits, and request cost:
140
+
141
+ ```python
142
+ from omniai import OmniaiClient, ChatCompletionResponse
143
+
144
+ client = OmniaiClient()
145
+
146
+ resp: ChatCompletionResponse = client.chat(
147
+ model="openrouter:inclusionai/ling-3.2-flash:free",
148
+ prompt="Explain quantum entanglement briefly.",
149
+ full_response=True,
150
+ )
151
+
152
+ print(f"Result: {resp.result}")
153
+ print(f"ID: {resp.id}")
154
+ print(f"Cached: {resp.cached}")
155
+ print(f"Tokens: {resp.usage.total_tokens} (prompt={resp.usage.prompt_tokens}, completion={resp.usage.completion_tokens})")
156
+ print(f" System Tokens: {resp.usage.details.system_tokens}")
157
+ print(f" User Input Tokens: {resp.usage.details.input_tokens}")
158
+ print(f"Request Cost: ${resp.billing.request_cost:.6f} USD")
159
+ print(f"Key Cumulative Spend: ${resp.billing.api_key.spent:.6f} USD")
160
+ print(f"Account Balance Remaining: ${resp.billing.account.balance_remaining:.6f} USD")
161
+ ```
162
+
163
+ ---
164
+
165
+ ### 3. Extended Reasoning / Thinking Budget
166
+
167
+ Allocate reasoning tokens for models that support chain-of-thought thinking (e.g. Gemini 3.7 Flash, Claude 3.7 Sonnet):
168
+
169
+ ```python
170
+ response = client.chat(
171
+ model="gemini:gemini-3.7-flash",
172
+ prompt="How many r's are in strawberry?",
173
+ thinking=1024, # Allocates 1,024 reasoning tokens
174
+ )
175
+ print(response)
176
+ ```
177
+
178
+ ---
179
+
180
+ ### 4. Real-Time Token Streaming
181
+
182
+ Stream completions token-by-token with Server-Sent Events (SSE):
136
183
 
137
184
  ```python
138
185
  from omniai import OmniaiClient
@@ -150,7 +197,7 @@ print()
150
197
 
151
198
  ---
152
199
 
153
- ### 3. Strongly Typed Structured Outputs (Pydantic V2)
200
+ ### 5. Strongly Typed Structured Outputs (Pydantic V2)
154
201
 
155
202
  Pass any Pydantic class to `response_model`. The SDK extracts JSON schema, enforces structured output from the model, and automatically parses and validates the response into your Python model:
156
203
 
@@ -179,68 +226,69 @@ print(f"Target SLA: {analysis.availability_sla}%")
179
226
 
180
227
  ---
181
228
 
182
- ### 4. Asynchronous Client (`AsyncOmniaiClient`)
229
+ ### 6. Custom AI Agents & Vector Knowledge Bases
183
230
 
184
- Full async/await parity for high-concurrency microservices, FastAPI, and async workers:
231
+ Create autonomous custom agents, ingest reference documents into Qdrant vector memory, and execute queries:
185
232
 
186
233
  ```python
187
- import asyncio
188
- from omniai import AsyncOmniaiClient
234
+ from omniai import OmniaiClient
189
235
 
190
- async def main():
191
- async with AsyncOmniaiClient() as client:
192
- # Async chat
193
- res = await client.chat(
194
- model="gemini:gemini-2.5-flash",
195
- prompt="Summarize CQRS in 20 words.",
196
- )
197
- print("Response:", res)
236
+ client = OmniaiClient()
198
237
 
199
- # Async streaming
200
- print("Streaming: ", end="", flush=True)
201
- async for chunk in client.stream(
202
- model="gemini:gemini-2.5-flash",
203
- prompt="Count from 1 to 5.",
204
- ):
205
- print(chunk, end="", flush=True)
206
- print()
238
+ # 1. Create a custom agent
239
+ agent = client.agents.create(
240
+ name="Refund Policy Assistant",
241
+ prompt="You are a refund specialist. Answer user questions strictly based on the company's 30-day guarantee.",
242
+ model="gemini:gemini-3.8-flash",
243
+ )
207
244
 
208
- asyncio.run(main())
245
+ # 2. Add document to agent's vectorized knowledge base
246
+ client.agents.knowledge.add(
247
+ agent_id=agent.id,
248
+ title="Refund Policy v2",
249
+ content="Customers can request a 100% refund within 30 days of purchase for unopened items.",
250
+ )
251
+
252
+ # 3. Execute the agent
253
+ answer = client.execute_agent(
254
+ agent_id=agent.id,
255
+ message="What is the refund window for an unopened box?",
256
+ )
257
+ print("Agent:", answer)
209
258
  ```
210
259
 
211
260
  ---
212
261
 
213
- ### 5. Background Task Queue (Celery Integration)
262
+ ### 7. Scoped API Key Management
214
263
 
215
- Queue long-running jobs to the background Celery cluster with fault-tolerant polling that survives transient network dropouts:
264
+ Programmatically provision, monitor, and revoke API keys with spend quotas and allowed models:
216
265
 
217
266
  ```python
218
267
  from omniai import OmniaiClient
219
- from pydantic import BaseModel
220
-
221
- class ComprehensiveAudit(BaseModel):
222
- vulnerabilities: list[str]
223
- risk_score: int
224
268
 
225
269
  client = OmniaiClient()
226
270
 
227
- # 1. Submit task asynchronously
228
- job = client.submit_task(
229
- model="gemini:gemini-2.5-flash",
230
- prompt="Perform in-depth threat modeling for an open S3 bucket with public access.",
231
- response_model=ComprehensiveAudit,
271
+ # 1. Create API key with $25 lifetime spend limit and model restrictions
272
+ new_key = client.keys.create(
273
+ name="staging-bot",
274
+ total_limit=25.0, # $25.00 lifetime spending limit
275
+ allowed_models=["openai:gpt-4o", "gemini:gemini-3.7-flash"],
276
+ expires_at="2027-12-31T23:59:59Z",
232
277
  )
233
- print(f"Dispatched Celery Task ID: {job.task_id}")
278
+ print(f"Key created: {new_key.name} -> {new_key.key}")
279
+
280
+ # 2. List all active keys
281
+ keys = client.keys.list()
282
+ for k in keys:
283
+ print(f"Key: {k.name} | Spent: ${k.current_spend:.4f} / ${k.total_limit or 'Unlimited'}")
234
284
 
235
- # 2. Poll until complete with automatic error-tolerance
236
- result: ComprehensiveAudit = job.wait(timeout=120.0, interval=2.0)
237
- print("Task finished successfully!")
238
- print(f"Risk Score: {result.risk_score}")
285
+ # 3. Revoke key
286
+ client.keys.delete(new_key.id)
239
287
  ```
240
288
 
241
289
  ---
242
290
 
243
- ### 6. LiveKit Multimodal Real-Time Voice & Video
291
+ ### 8. LiveKit Multimodal Real-Time Voice & Video
244
292
 
245
293
  Create WebRTC sessions for sub-second bidirectional voice and vision agents:
246
294
 
@@ -250,7 +298,7 @@ from omniai import OmniaiClient
250
298
  client = OmniaiClient()
251
299
 
252
300
  session = client.create_live_session(
253
- model="gemini:gemini-2.5-flash",
301
+ model="gemini-3.8-live",
254
302
  voice="Aoede",
255
303
  video_input=True,
256
304
  )
@@ -262,19 +310,34 @@ print(f"Token: {session.participant_token}")
262
310
 
263
311
  ---
264
312
 
265
- ### 7. Model Discovery & Catalog
313
+ ### 9. Asynchronous Client (`AsyncOmniaiClient`)
266
314
 
267
- Query available models dynamically from the gateway:
315
+ Full async/await parity for high-concurrency microservices, FastAPI, and async workers:
268
316
 
269
317
  ```python
270
- from omniai import OmniaiClient
318
+ import asyncio
319
+ from omniai import AsyncOmniaiClient
271
320
 
272
- client = OmniaiClient()
321
+ async def main():
322
+ async with AsyncOmniaiClient() as client:
323
+ # Async chat with reasoning tokens
324
+ res = await client.chat(
325
+ model="gemini:gemini-3.7-flash",
326
+ prompt="Summarize CQRS in 20 words.",
327
+ thinking=1024,
328
+ )
329
+ print("Response:", res)
330
+
331
+ # Async streaming
332
+ print("Streaming: ", end="", flush=True)
333
+ async for chunk in client.stream(
334
+ model="gemini:gemini-2.5-flash",
335
+ prompt="Count from 1 to 5.",
336
+ ):
337
+ print(chunk, end="", flush=True)
338
+ print()
273
339
 
274
- models = client.models.list()
275
- print(f"Total available models: {len(models.data)}")
276
- for m in models.data[:5]:
277
- print(f" • {m.id} (provider: {m.provider}, release: {m.release_date})")
340
+ asyncio.run(main())
278
341
  ```
279
342
 
280
343
  ---
@@ -290,11 +353,6 @@ for m in models.data[:5]:
290
353
  | **`timeout`** | `OmniaiClient(timeout=30.0)` | — | `60.0` seconds |
291
354
  | **`max_retries`**| `OmniaiClient(max_retries=5)`| — | `3` attempts |
292
355
 
293
- Create a `.env` file in your project directory:
294
- ```env
295
- OMNIAI_API_KEY=sk-or-v1-your-api-key-here
296
- ```
297
-
298
356
  ---
299
357
 
300
358
  ## 🛡️ Exception Hierarchy
@@ -316,29 +374,30 @@ OmniaiError
316
374
 
317
375
  ---
318
376
 
319
- ## 📂 Repository Directory Pattern
377
+ ## 📂 Repository Directory Layout
320
378
 
321
379
  ```text
322
380
  omniai-sdk/
323
- ├── .github/
324
- │ └── workflows/
325
- │ └── pypi-publish.yml # PyPI publishing workflow
326
381
  ├── omniai/
327
382
  │ ├── resources/ # Modular resource namespaces
328
- │ │ ├── chat.py
329
- │ │ ├── livekit.py
330
- │ │ ├── models.py
331
- │ │ ├── router.py
332
- │ │ └── tasks.py
383
+ │ │ ├── chat.py # Chat completions & SSE streaming
384
+ │ │ ├── agents.py # Custom AI agents & RAG knowledge
385
+ │ │ ├── keys.py # Scoped API key management
386
+ │ │ ├── livekit.py # Multimodal WebRTC voice/video
387
+ │ │ ├── models.py # Foundation models catalog
388
+ │ │ ├── router.py # Multi-provider LLM router
389
+ │ │ └── tasks.py # Celery background tasks
333
390
  │ ├── types/ # Typed Pydantic schemas
334
- │ │ ├── chat.py
335
- │ │ ├── livekit.py
336
- │ │ ├── models.py
337
- │ │ └── tasks.py
338
- │ ├── __init__.py # Top-level exports
391
+ │ │ ├── chat.py # Response, usage, billing envelopes
392
+ │ │ ├── agents.py # Agent, tool, knowledge models
393
+ │ │ ├── keys.py # API key records
394
+ │ │ ├── livekit.py # Live session credentials
395
+ │ │ ├── models.py # Foundation model metadata
396
+ │ │ └── tasks.py # Task handles and job wrappers
397
+ │ ├── __init__.py # Top-level SDK exports
339
398
  │ ├── _base.py # Base HTTP transport & retry policy
340
399
  │ ├── _constants.py # Default constants & endpoints
341
- │ ├── _version.py # Package version (1.0.0)
400
+ │ ├── _version.py # Package version (1.1.0)
342
401
  │ ├── client.py # OmniaiClient & AsyncOmniaiClient
343
402
  │ ├── exceptions.py # Typed exception hierarchy
344
403
  │ └── py.typed # PEP 561 typing compliance
@@ -346,15 +405,14 @@ omniai-sdk/
346
405
  │ ├── quickstart.py
347
406
  │ ├── streaming.py
348
407
  │ ├── structured_output.py
349
- │ └── async_chat.py
408
+ │ ├── async_chat.py
409
+ │ ├── thinking_reasoning.py
410
+ │ ├── custom_agents.py
411
+ │ └── api_keys.py
350
412
  ├── tests/ # Unit and integration test suite
351
- ├── .env.example
352
- ├── .gitignore
353
- ├── LICENSE # MIT License
354
- ├── README.md # Documentation & Guides
413
+ ├── test.py # Root verification script
355
414
  ├── pyproject.toml # PEP 517/621 packaging metadata
356
- ├── requirements.txt # Core dependencies (httpx, pydantic)
357
- └── test.py # Root test script for quick validation
415
+ └── README.md # Documentation & Guides
358
416
  ```
359
417
 
360
418
  ---
@@ -367,7 +425,7 @@ Run the test suite with `pytest`:
367
425
  pytest tests
368
426
  ```
369
427
 
370
- Or run the quick root validation script:
428
+ Or run the root verification script:
371
429
 
372
430
  ```bash
373
431
  python test.py