pyinferencemanager 1.0.0__tar.gz → 1.3.0__tar.gz

{pyinferencemanager-1.0.0 → pyinferencemanager-1.3.0}/Cargo.lock

@@ -26,6 +26,16 @@ version = "1.0.104"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
 
+[[package]]
+name = "assert-json-diff"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "async-trait"
 version = "0.1.91"
@@ -233,6 +243,24 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "deadpool"
+version = "0.12.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b"
+dependencies = [
+ "deadpool-runtime",
+ "lazy_static",
+ "num_cpus",
+ "tokio",
+]
+
+[[package]]
+name = "deadpool-runtime"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b"
+
 [[package]]
 name = "deranged"
 version = "0.5.8"
@@ -576,6 +604,12 @@ version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
 
+[[package]]
+name = "httpdate"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
+
 [[package]]
 name = "hyper"
 version = "1.11.0"
@@ -590,6 +624,7 @@ dependencies = [
  "http",
  "http-body",
  "httparse",
+ "httpdate",
  "itoa",
  "pin-project-lite",
  "smallvec",
@@ -1304,7 +1339,7 @@ checksum = "106dd99e98437432fed6519dedecfade6a06a73bb7b2a1e019fdd2bee5778d94"
 
 [[package]]
 name = "pyinferencemanager"
-version = "1.0.0"
+version = "1.3.0"
 dependencies = [
  "chrono",
  "pyinferencemanager-core",
@@ -1315,7 +1350,7 @@ dependencies = [
 
 [[package]]
 name = "pyinferencemanager-core"
-version = "1.0.0"
+version = "1.3.0"
 dependencies = [
  "anyhow",
  "async-trait",
@@ -1330,6 +1365,7 @@ dependencies = [
  "rusqlite",
  "serde",
  "serde_json",
+ "serial_test",
  "sha2",
  "sqlite-vec",
  "tempfile",
@@ -1339,6 +1375,7 @@ dependencies = [
  "tracing-opentelemetry",
  "tracing-subscriber",
  "uuid",
+ "wiremock",
 ]
 
 [[package]]
@@ -1856,6 +1893,31 @@ dependencies = [
  "zmij",
 ]
 
+[[package]]
+name = "serial_test"
+version = "3.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "699f4197115b8a7e7ff19c9a315a4bd6fffec26cc4626ef45ecaea389e081c6d"
+dependencies = [
+ "futures-executor",
+ "futures-util",
+ "log",
+ "once_cell",
+ "parking_lot",
+ "serial_test_derive",
+]
+
+[[package]]
+name = "serial_test_derive"
+version = "3.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94e153fc76e1c6a068703d6d29c508a0b15c061c4b7e43da59cc097bc342673c"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
 [[package]]
 name = "sha2"
 version = "0.10.9"
@@ -2726,6 +2788,29 @@ version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
 
+[[package]]
+name = "wiremock"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031"
+dependencies = [
+ "assert-json-diff",
+ "base64",
+ "deadpool",
+ "futures",
+ "http",
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "log",
+ "once_cell",
+ "regex",
+ "serde",
+ "serde_json",
+ "tokio",
+ "url",
+]
+
 [[package]]
 name = "writeable"
 version = "0.6.3"

{pyinferencemanager-1.0.0 → pyinferencemanager-1.3.0}/Cargo.toml

@@ -3,7 +3,7 @@ members = ["crates/pyinferencemanager-core", "crates/pyinferencemanager-py"]
 resolver = "2"
 
 [workspace.package]
-version = "1.0.0"
+version = "1.3.0"
 edition = "2021"
 rust-version = "1.82"
 authors = ["Georgi Mammen Mullassery <mullassery@gmail.com>"]

pyinferencemanager-1.3.0/LICENSE

@@ -0,0 +1,22 @@
+Proprietary License — Free to use with explicit attribution
+
+Copyright © 2026 Georgi Mammen Mullassery
+
+Permission is granted to any person obtaining a copy of this software to use,
+copy, modify, and distribute this software and its documentation, provided that:
+
+1. The above copyright notice and this permission notice appear in all copies
+   or substantial portions of the software.
+
+2. Explicit attribution to the original author (Georgi Mammen Mullassery) is
+   provided in any project documentation and significant code comments where
+   this software is used.
+
+3. The software is provided "AS IS" without warranty of any kind, express or
+   implied, including but not limited to the warranties of merchantability,
+   fitness for a particular purpose, or non-infringement.
+
+4. In no event shall the author be liable for any claim, damages, or other
+   liability arising from the use of this software.
+
+For more information, visit: https://github.com/Mullassery

pyinferencemanager-1.3.0/PKG-INFO

@@ -0,0 +1,162 @@
+Metadata-Version: 2.4
+Name: pyinferencemanager
+Version: 1.3.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Rust
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Classifier: Typing :: Typed
+Requires-Dist: pytest>=7.0 ; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.21 ; extra == 'dev'
+Requires-Dist: maturin>=1.0 ; extra == 'dev'
+Requires-Dist: ruff>=0.5.0 ; extra == 'dev'
+Provides-Extra: dev
+License-File: LICENSE
+Summary: Intelligent AI workload orchestrator. Multi-provider routing, cost optimization, semantic caching, dynamic routing with budget enforcement. Production observability included.
+Keywords: ai,orchestrator,llm,ollama,anthropic,openai,local-first,cost-optimization,caching,multi-provider,routing,inference,workload-orchestration,hardware-aware,semantic-caching,distributed-tracing,observability,budget-enforcement
+Author-email: Georgi Mammen Mullassery <mullassery@gmail.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Changelog, https://github.com/Mullassery/pyinferencemanager/releases
+Project-URL: Documentation, https://github.com/Mullassery/pyinferencemanager/blob/main/README.md
+Project-URL: Homepage, https://github.com/Mullassery/pyinferencemanager
+Project-URL: Issues, https://github.com/Mullassery/pyinferencemanager/issues
+Project-URL: Repository, https://github.com/Mullassery/pyinferencemanager.git
+
+# PyInferenceManager
+
+**A multi-provider LLM inference orchestrator with cost/latency-aware routing, real retry + budget guardrails, and local-first execution.**
+
+Routes requests to Anthropic Claude, OpenAI GPT, Google Gemini, or a local Ollama model based on task complexity, privacy, and observed provider health — with real retry/backoff on failures and a real spend cap you configure.
+
+[![PyPI](https://img.shields.io/pypi/v/pyinferencemanager)](https://pypi.org/project/pyinferencemanager)
+[![Python 3.10+](https://img.shields.io/badge/Python-3.10%2B-blue)](https://www.python.org)
+[![Tests](https://img.shields.io/github/actions/workflow/status/Mullassery/PyInferenceManager/tests.yml?label=tests)](https://github.com/Mullassery/PyInferenceManager/actions)
+[![License: Proprietary](https://img.shields.io/badge/License-Proprietary-blue.svg)](./LICENSE)
+
+---
+
+## What this actually is
+
+The core (Rust, ~11k lines, PyO3 bindings) is a workload orchestrator: it builds a small execution DAG for a task, profiles your local hardware, picks an engine (local Ollama model vs. a cloud provider) based on complexity/privacy/cost, executes with real retry and budget enforcement, and semantically caches results.
+
+The Python API surface is a single `Orchestrator` class. There is no `Manager.chat()`, no streaming API, and no 11-provider marketplace — see [Providers](#providers) for the honest list.
+
+## Install
+
+```bash
+pip install pyinferencemanager
+```
+
+Requires Python 3.10+.
+
+Set whichever provider keys you plan to use (none are required for local-only Ollama use):
+
+```bash
+export ANTHROPIC_API_KEY=sk-ant-...
+export OPENAI_API_KEY=sk-...
+export GEMINI_API_KEY=...   # or GOOGLE_API_KEY
+```
+
+See [`.env.example`](.env.example).
+
+## Quick start
+
+```python
+from pyinferencemanager import Orchestrator
+
+orchestrator = Orchestrator(mode="local_first")  # or "cloud_first"
+
+result = orchestrator.run(task="question_answering", message="What is the capital of France?")
+print(result.output)
+print(f"Engines used: {result.engines_used}")
+print(f"Cost: ${result.total_cost_usd:.4f} | Latency: {result.total_latency_ms}ms | Tokens: {result.total_tokens}")
+```
+
+`mode="local_first"` runs on your local Ollama model when it's adequate for the task's complexity and escalates to a cloud provider otherwise. `mode="cloud_first"` prefers a cloud provider, falling back to local only for very low-complexity tasks. `privacy="high"` on any call always forces local execution regardless of mode.
+
+If the selected backend is unreachable or unauthenticated, `run()` doesn't raise — it returns a result whose `output` says so (e.g. `"[cloud inference unavailable: ...]"`), so a single bad provider doesn't crash your pipeline.
+
+## Core API
+
+```python
+orchestrator = Orchestrator(mode="local_first")
+
+# Execute a task — real inference against Ollama or a cloud provider.
+result = orchestrator.run(task="...", file=None, message=None, privacy="low")
+# result.output, .total_tokens, .total_cost_usd, .total_latency_ms, .engines_used, .cache_hits
+
+# Estimate cost/latency without executing anything.
+plan = orchestrator.plan("Summarize this document")
+# plan.stages, .estimated_cost_usd, .estimated_latency_ms, .local_first
+
+# Real-time provider health, from actually completed cloud calls.
+orchestrator.provider_ranking()      # [(provider_key, health_score), ...]
+orchestrator.provider_performance()  # {provider_key: {success_rate, avg_latency_ms, ...}}
+
+# Cost guardrails, enforced on every real cloud call.
+orchestrator.configure_budget(max_cost_usd=10.0, max_requests=1000,
+                               alert_threshold_percent=80.0, enforce_hard_limit=True)
+orchestrator.budget_status()
+
+# Retry/backoff policy for retryable errors (HTTP 429/408/5xx).
+orchestrator.configure_retry(max_attempts=3, backoff="exponential",
+                              initial_ms=100, max_ms=5000)
+
+# Synthetic load test exercising the same budget + dynamic-routing logic
+# run() uses, at a volume impractical against live APIs. Latencies/costs
+# here are simulated, not real network calls.
+orchestrator.run_load_test(num_requests=200, budget_usd=5.0)
+
+# Local hardware profile (memory tier, Apple Silicon/Metal, Ollama models).
+orchestrator.profile_hardware()
+orchestrator.available_backends()
+```
+
+See [`examples/`](examples/) for runnable scripts covering each of these.
+
+## Providers
+
+Real HTTP calls, with retry/backoff and cost tracking:
+
+| Provider | Env var | Notes |
+|---|---|---|
+| Anthropic Claude | `ANTHROPIC_API_KEY` | `claude-haiku-4-5`, `claude-opus-4-1` |
+| OpenAI | `OPENAI_API_KEY` | `gpt-4o-mini` |
+| Google Gemini | `GEMINI_API_KEY` / `GOOGLE_API_KEY` | `gemini-1.5-flash` |
+| Ollama (local) | — | whatever models you've pulled locally |
+| vLLM (local) | — | OpenAI-compatible endpoint, default `localhost:8000` |
+
+Additionally, `tensorrt_llm`, `mlc_llm`, and `colibri` exist as **cost-estimator-only stubs** (`BackendKind`/`RuntimeBackend` trait implementations with real cost/latency estimation logic, but no live inference) — they're placeholders for self-hosted inference servers, not currently wired to make real calls. `orchestrator.available_backends()` lists all of these honestly, including the stubs.
+
+Cloud execution dispatches through `BackendRegistry`/`RuntimeBackend` (not a hand-matched provider enum): `ProviderExecutor::execute` maps a `CloudProvider` to its `BackendKind`, registers the one backend it needs (reading the API key from the env vars above), and calls it through the same `RuntimeBackend::infer` trait object every backend implements — so adding a new cloud provider is a new `BackendKind`/backend + one dispatch arm, not edits scattered across every call site that used to match the old enum.
+
+## MCP tools
+
+`pyinferencemanager._mcp_tools.PyInferenceManagerMCPHandler` exposes 13 MCP-style tools (`list_available_models`, `execute_inference`, `estimate_inference_cost`, `get_provider_status`, etc.), all backed by a real `Orchestrator` instance — no hardcoded responses. See [`examples/mcp_pyinferencemanager.py`](examples/mcp_pyinferencemanager.py).
+
+The network connector (`_mcp_connector.InferenceManager.start_mcp_connector()`) binds to `127.0.0.1` by default with scoped CORS and permissions; binding elsewhere requires passing `allow_remote=True` explicitly.
+
+## Testing
+
+- Rust core: `cargo test --workspace` (350+ tests, including HTTP-mocked request/response tests for every cloud client via [`wiremock`](https://docs.rs/wiremock)).
+- Python: `pytest tests/` — covers `Orchestrator.run/plan/provider_ranking`, the MCP tool handlers, budget/retry configuration, load testing, and MCP connector security defaults, all without requiring live API keys.
+
+## Known issues
+
+- `tensorrt_llm`, `mlc_llm`, and `colibri` backends are cost-estimator-only stubs (see [Providers](#providers)) — they do not make live inference calls yet. `vllm` was a stub too as of 1.2.0 but is now a real backend.
+- No open GitHub issues and no `TODO`/`FIXME` markers in the codebase at the time of this writing.
+
+## License
+
+Proprietary License — free to use with explicit attribution. See [LICENSE](LICENSE).
+

pyinferencemanager-1.3.0/README.md

@@ -0,0 +1,127 @@
+# PyInferenceManager
+
+**A multi-provider LLM inference orchestrator with cost/latency-aware routing, real retry + budget guardrails, and local-first execution.**
+
+Routes requests to Anthropic Claude, OpenAI GPT, Google Gemini, or a local Ollama model based on task complexity, privacy, and observed provider health — with real retry/backoff on failures and a real spend cap you configure.
+
+[![PyPI](https://img.shields.io/pypi/v/pyinferencemanager)](https://pypi.org/project/pyinferencemanager)
+[![Python 3.10+](https://img.shields.io/badge/Python-3.10%2B-blue)](https://www.python.org)
+[![Tests](https://img.shields.io/github/actions/workflow/status/Mullassery/PyInferenceManager/tests.yml?label=tests)](https://github.com/Mullassery/PyInferenceManager/actions)
+[![License: Proprietary](https://img.shields.io/badge/License-Proprietary-blue.svg)](./LICENSE)
+
+---
+
+## What this actually is
+
+The core (Rust, ~11k lines, PyO3 bindings) is a workload orchestrator: it builds a small execution DAG for a task, profiles your local hardware, picks an engine (local Ollama model vs. a cloud provider) based on complexity/privacy/cost, executes with real retry and budget enforcement, and semantically caches results.
+
+The Python API surface is a single `Orchestrator` class. There is no `Manager.chat()`, no streaming API, and no 11-provider marketplace — see [Providers](#providers) for the honest list.
+
+## Install
+
+```bash
+pip install pyinferencemanager
+```
+
+Requires Python 3.10+.
+
+Set whichever provider keys you plan to use (none are required for local-only Ollama use):
+
+```bash
+export ANTHROPIC_API_KEY=sk-ant-...
+export OPENAI_API_KEY=sk-...
+export GEMINI_API_KEY=...   # or GOOGLE_API_KEY
+```
+
+See [`.env.example`](.env.example).
+
+## Quick start
+
+```python
+from pyinferencemanager import Orchestrator
+
+orchestrator = Orchestrator(mode="local_first")  # or "cloud_first"
+
+result = orchestrator.run(task="question_answering", message="What is the capital of France?")
+print(result.output)
+print(f"Engines used: {result.engines_used}")
+print(f"Cost: ${result.total_cost_usd:.4f} | Latency: {result.total_latency_ms}ms | Tokens: {result.total_tokens}")
+```
+
+`mode="local_first"` runs on your local Ollama model when it's adequate for the task's complexity and escalates to a cloud provider otherwise. `mode="cloud_first"` prefers a cloud provider, falling back to local only for very low-complexity tasks. `privacy="high"` on any call always forces local execution regardless of mode.
+
+If the selected backend is unreachable or unauthenticated, `run()` doesn't raise — it returns a result whose `output` says so (e.g. `"[cloud inference unavailable: ...]"`), so a single bad provider doesn't crash your pipeline.
+
+## Core API
+
+```python
+orchestrator = Orchestrator(mode="local_first")
+
+# Execute a task — real inference against Ollama or a cloud provider.
+result = orchestrator.run(task="...", file=None, message=None, privacy="low")
+# result.output, .total_tokens, .total_cost_usd, .total_latency_ms, .engines_used, .cache_hits
+
+# Estimate cost/latency without executing anything.
+plan = orchestrator.plan("Summarize this document")
+# plan.stages, .estimated_cost_usd, .estimated_latency_ms, .local_first
+
+# Real-time provider health, from actually completed cloud calls.
+orchestrator.provider_ranking()      # [(provider_key, health_score), ...]
+orchestrator.provider_performance()  # {provider_key: {success_rate, avg_latency_ms, ...}}
+
+# Cost guardrails, enforced on every real cloud call.
+orchestrator.configure_budget(max_cost_usd=10.0, max_requests=1000,
+                               alert_threshold_percent=80.0, enforce_hard_limit=True)
+orchestrator.budget_status()
+
+# Retry/backoff policy for retryable errors (HTTP 429/408/5xx).
+orchestrator.configure_retry(max_attempts=3, backoff="exponential",
+                              initial_ms=100, max_ms=5000)
+
+# Synthetic load test exercising the same budget + dynamic-routing logic
+# run() uses, at a volume impractical against live APIs. Latencies/costs
+# here are simulated, not real network calls.
+orchestrator.run_load_test(num_requests=200, budget_usd=5.0)
+
+# Local hardware profile (memory tier, Apple Silicon/Metal, Ollama models).
+orchestrator.profile_hardware()
+orchestrator.available_backends()
+```
+
+See [`examples/`](examples/) for runnable scripts covering each of these.
+
+## Providers
+
+Real HTTP calls, with retry/backoff and cost tracking:
+
+| Provider | Env var | Notes |
+|---|---|---|
+| Anthropic Claude | `ANTHROPIC_API_KEY` | `claude-haiku-4-5`, `claude-opus-4-1` |
+| OpenAI | `OPENAI_API_KEY` | `gpt-4o-mini` |
+| Google Gemini | `GEMINI_API_KEY` / `GOOGLE_API_KEY` | `gemini-1.5-flash` |
+| Ollama (local) | — | whatever models you've pulled locally |
+| vLLM (local) | — | OpenAI-compatible endpoint, default `localhost:8000` |
+
+Additionally, `tensorrt_llm`, `mlc_llm`, and `colibri` exist as **cost-estimator-only stubs** (`BackendKind`/`RuntimeBackend` trait implementations with real cost/latency estimation logic, but no live inference) — they're placeholders for self-hosted inference servers, not currently wired to make real calls. `orchestrator.available_backends()` lists all of these honestly, including the stubs.
+
+Cloud execution dispatches through `BackendRegistry`/`RuntimeBackend` (not a hand-matched provider enum): `ProviderExecutor::execute` maps a `CloudProvider` to its `BackendKind`, registers the one backend it needs (reading the API key from the env vars above), and calls it through the same `RuntimeBackend::infer` trait object every backend implements — so adding a new cloud provider is a new `BackendKind`/backend + one dispatch arm, not edits scattered across every call site that used to match the old enum.
+
+## MCP tools
+
+`pyinferencemanager._mcp_tools.PyInferenceManagerMCPHandler` exposes 13 MCP-style tools (`list_available_models`, `execute_inference`, `estimate_inference_cost`, `get_provider_status`, etc.), all backed by a real `Orchestrator` instance — no hardcoded responses. See [`examples/mcp_pyinferencemanager.py`](examples/mcp_pyinferencemanager.py).
+
+The network connector (`_mcp_connector.InferenceManager.start_mcp_connector()`) binds to `127.0.0.1` by default with scoped CORS and permissions; binding elsewhere requires passing `allow_remote=True` explicitly.
+
+## Testing
+
+- Rust core: `cargo test --workspace` (350+ tests, including HTTP-mocked request/response tests for every cloud client via [`wiremock`](https://docs.rs/wiremock)).
+- Python: `pytest tests/` — covers `Orchestrator.run/plan/provider_ranking`, the MCP tool handlers, budget/retry configuration, load testing, and MCP connector security defaults, all without requiring live API keys.
+
+## Known issues
+
+- `tensorrt_llm`, `mlc_llm`, and `colibri` backends are cost-estimator-only stubs (see [Providers](#providers)) — they do not make live inference calls yet. `vllm` was a stub too as of 1.2.0 but is now a real backend.
+- No open GitHub issues and no `TODO`/`FIXME` markers in the codebase at the time of this writing.
+
+## License
+
+Proprietary License — free to use with explicit attribution. See [LICENSE](LICENSE).

{pyinferencemanager-1.0.0 → pyinferencemanager-1.3.0}/crates/pyinferencemanager-core/Cargo.toml

@@ -40,3 +40,5 @@ python = ["pyo3"]
 [dev-dependencies]
 tokio = { version = "1.53", features = ["full", "test-util"] }
 tempfile = "3"
+wiremock = "0.6"
+serial_test = "3"

pyinferencemanager-1.3.0/crates/pyinferencemanager-core/src/backends/gemini_backend.rs

@@ -0,0 +1,86 @@
+use super::{BackendKind, InferenceRequest, InferenceResult, ModelProfile, RuntimeBackend};
+use crate::engines::GeminiClient;
+
+pub struct GeminiBackend {
+    api_key: String,
+}
+
+impl GeminiBackend {
+    pub fn new(api_key: String) -> Self {
+        GeminiBackend { api_key }
+    }
+}
+
+#[async_trait::async_trait]
+impl RuntimeBackend for GeminiBackend {
+    async fn infer(&self, request: InferenceRequest) -> crate::Result<InferenceResult> {
+        let start = std::time::Instant::now();
+        let client = GeminiClient::new(self.api_key.clone(), request.model.clone());
+
+        let response = client.complete(&request.prompt, request.max_tokens).await?;
+
+        let latency_ms = start.elapsed().as_millis() as u64;
+
+        Ok(InferenceResult {
+            output: response.text,
+            tokens_used: response.tokens_used,
+            latency_ms,
+        })
+    }
+
+    fn estimate_cost(&self, profile: &ModelProfile) -> f32 {
+        // Gemini Flash-tier pricing is roughly in the same band as
+        // OpenAI's smaller models.
+        let size_gb = profile.size_gb();
+        if size_gb < 10.0 {
+            0.0001
+        } else if size_gb < 30.0 {
+            0.00025
+        } else {
+            0.0005
+        }
+    }
+
+    fn estimate_latency(&self, profile: &ModelProfile) -> u64 {
+        let size_gb = profile.size_gb();
+        let base = (size_gb * 8.0) as u64;
+        base.clamp(80, 4000)
+    }
+
+    fn kind(&self) -> BackendKind {
+        BackendKind::Gemini
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_gemini_backend_kind() {
+        let backend = GeminiBackend::new("test-key".to_string());
+        assert_eq!(backend.kind(), BackendKind::Gemini);
+    }
+
+    #[test]
+    fn test_gemini_estimate_cost() {
+        let backend = GeminiBackend::new("test-key".to_string());
+        let small_profile = ModelProfile::new(4 * 1_073_741_824, "gemini".to_string());
+        let large_profile = ModelProfile::new(50 * 1_073_741_824, "gemini".to_string());
+
+        let small_cost = backend.estimate_cost(&small_profile);
+        let large_cost = backend.estimate_cost(&large_profile);
+
+        assert!(small_cost < large_cost);
+    }
+
+    #[test]
+    fn test_gemini_estimate_latency() {
+        let backend = GeminiBackend::new("test-key".to_string());
+        let profile = ModelProfile::new(8 * 1_073_741_824, "gemini".to_string());
+        let latency = backend.estimate_latency(&profile);
+
+        assert!(latency >= 80);
+        assert!(latency <= 4000);
+    }
+}

{pyinferencemanager-1.0.0 → pyinferencemanager-1.3.0}/crates/pyinferencemanager-core/src/backends/mod.rs

@@ -3,20 +3,25 @@ use std::collections::HashMap;
 
 pub mod anthropic_backend;
 pub mod colibri_backend;
+pub mod gemini_backend;
 pub mod ollama_backend;
 pub mod openai_backend;
 pub mod stub_backend;
+pub mod vllm_backend;
 
 pub use anthropic_backend::AnthropicBackend;
 pub use colibri_backend::ColibriBackend;
+pub use gemini_backend::GeminiBackend;
 pub use ollama_backend::OllamaBackend;
 pub use openai_backend::OpenAiBackend;
 pub use stub_backend::StubBackend;
+pub use vllm_backend::VLlmBackend;
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub enum BackendKind {
     Anthropic,
     OpenAi,
+    Gemini,
     Ollama,
     VLlm,
     TensorRtLlm,
@@ -29,6 +34,7 @@ impl BackendKind {
         match self {
             BackendKind::Anthropic => "anthropic",
             BackendKind::OpenAi => "openai",
+            BackendKind::Gemini => "gemini",
             BackendKind::Ollama => "ollama",
             BackendKind::VLlm => "vllm",
             BackendKind::TensorRtLlm => "tensorrt_llm",
@@ -99,8 +105,9 @@ pub trait RuntimeBackend: Send + Sync {
 pub enum AnyBackend {
     Anthropic(AnthropicBackend),
     OpenAi(OpenAiBackend),
+    Gemini(GeminiBackend),
     Ollama(OllamaBackend),
-    VLlm(StubBackend),
+    VLlm(VLlmBackend),
     TensorRtLlm(StubBackend),
     MlcLlm(StubBackend),
     Colibri(ColibriBackend),
@@ -111,6 +118,7 @@ impl AnyBackend {
         match self {
             AnyBackend::Anthropic(b) => b.infer(request).await,
             AnyBackend::OpenAi(b) => b.infer(request).await,
+            AnyBackend::Gemini(b) => b.infer(request).await,
             AnyBackend::Ollama(b) => b.infer(request).await,
             AnyBackend::VLlm(b) => b.infer(request).await,
             AnyBackend::TensorRtLlm(b) => b.infer(request).await,
@@ -123,6 +131,7 @@ impl AnyBackend {
         match self {
             AnyBackend::Anthropic(b) => b.estimate_cost(profile),
             AnyBackend::OpenAi(b) => b.estimate_cost(profile),
+            AnyBackend::Gemini(b) => b.estimate_cost(profile),
             AnyBackend::Ollama(b) => b.estimate_cost(profile),
             AnyBackend::VLlm(b) => b.estimate_cost(profile),
             AnyBackend::TensorRtLlm(b) => b.estimate_cost(profile),
@@ -135,6 +144,7 @@ impl AnyBackend {
         match self {
             AnyBackend::Anthropic(b) => b.estimate_latency(profile),
             AnyBackend::OpenAi(b) => b.estimate_latency(profile),
+            AnyBackend::Gemini(b) => b.estimate_latency(profile),
             AnyBackend::Ollama(b) => b.estimate_latency(profile),
             AnyBackend::VLlm(b) => b.estimate_latency(profile),
             AnyBackend::TensorRtLlm(b) => b.estimate_latency(profile),
@@ -147,6 +157,7 @@ impl AnyBackend {
         match self {
             AnyBackend::Anthropic(b) => b.kind(),
             AnyBackend::OpenAi(b) => b.kind(),
+            AnyBackend::Gemini(b) => b.kind(),
             AnyBackend::Ollama(b) => b.kind(),
             AnyBackend::VLlm(b) => b.kind(),
             AnyBackend::TensorRtLlm(b) => b.kind(),