Skip to content

Commit 26cf502

Browse files
committed
Add AIMLAPI provider integration and example
Introduces AIMLAPIProvider in marvin.providers for AI/ML API support, including provider registration in __init__.py. Adds an example script demonstrating agent usage with AIMLAPI in examples/provider_specific/aimlapi/run_agent.py.
1 parent 578c0a0 commit 26cf502

6 files changed

Lines changed: 104 additions & 2 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ Install `marvin`:
1919
uv pip install marvin
2020
```
2121

22-
Configure your LLM provider (Marvin uses OpenAI by default but natively supports [all Pydantic AI models](https://ai.pydantic.dev/models/)):
22+
Configure your LLM provider (Marvin uses OpenAI by default but also supports providers like [AI/ML API](https://aimlapi.com/) and other [Pydantic AI models](https://ai.pydantic.dev/models/)):
2323

2424
```bash
2525
export OPENAI_API_KEY=your-api-key

docs/guides/configure-llms.mdx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ Marvin supports any model provider that is compatible with Pydantic AI. Common p
7777
- Anthropic
7878
- Azure OpenAI
7979
- Google
80+
- [AI/ML API](https://aimlapi.com) - 300+ models with an OpenAI-compatible interface
8081

8182
Each provider may require its own API key and configuration. Refer to the provider's [documentation](https://ai.pydantic.dev/models/) for specific setup instructions.
8283

docs/installation.mdx

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ icon: download
77
## Requirements
88

99
- Python 3.10 or higher
10-
- An API key from an LLM provider (OpenAI by default)
10+
- An API key from an LLM provider (OpenAI by default, or [AI/ML API](https://aimlapi.com/))
1111

1212
## Install `marvin`
1313

@@ -40,6 +40,12 @@ By default, Marvin uses OpenAI's models. Set your API key as an environment vari
4040
export OPENAI_API_KEY="your-api-key"
4141
```
4242

43+
To use AI/ML API, set the `AIML_API_KEY` variable instead:
44+
45+
```bash
46+
export AIML_API_KEY="your-api-key"
47+
```
48+
4349
To use another provider, see the docs on [configuring LLMs](/guides/configure-llms).
4450

4551
## Development Installation
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
import os
2+
from pathlib import Path
3+
4+
from pydantic_ai.models.openai import OpenAIModel
5+
6+
import marvin
7+
from marvin.providers.aimlapi import AIMLAPIProvider
8+
9+
10+
def write_file(path: str, content: str) -> None:
11+
"""Write content to a file."""
12+
Path(path).write_text(content)
13+
14+
15+
writer = marvin.Agent(
16+
model=OpenAIModel(
17+
"gpt-4o-mini",
18+
provider=AIMLAPIProvider(api_key=os.getenv("AIML_API_KEY")),
19+
),
20+
name="AI/ML Writer",
21+
instructions="Write concise, engaging content for developers",
22+
tools=[write_file],
23+
)
24+
25+
result = marvin.run("how to use pydantic? write haiku to docs.md", agents=[writer])
26+
print(result)

src/marvin/providers/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from .aimlapi import AIMLAPIProvider
2+
3+
__all__ = ["AIMLAPIProvider"]

src/marvin/providers/aimlapi.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from __future__ import annotations
2+
3+
import os
4+
from typing import overload
5+
6+
from httpx import AsyncClient as AsyncHTTPClient
7+
from openai import AsyncOpenAI
8+
9+
from pydantic_ai.exceptions import UserError
10+
from pydantic_ai.models import cached_async_http_client
11+
from pydantic_ai.profiles import ModelProfile
12+
from pydantic_ai.profiles.openai import openai_model_profile
13+
from pydantic_ai.providers import Provider
14+
15+
16+
class AIMLAPIProvider(Provider[AsyncOpenAI]):
17+
"""Provider for the AI/ML API."""
18+
19+
@property
20+
def name(self) -> str: # pragma: no cover - simple property
21+
return "aimlapi"
22+
23+
@property
24+
def base_url(self) -> str: # pragma: no cover - simple property
25+
return "https://api.aimlapi.com/v1"
26+
27+
@property
28+
def client(self) -> AsyncOpenAI:
29+
return self._client
30+
31+
def model_profile(self, model_name: str) -> ModelProfile | None: # pragma: no cover - thin wrapper
32+
return openai_model_profile(model_name)
33+
34+
@overload
35+
def __init__(self) -> None: ...
36+
37+
@overload
38+
def __init__(self, *, api_key: str) -> None: ...
39+
40+
@overload
41+
def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
42+
43+
@overload
44+
def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
45+
46+
def __init__(
47+
self,
48+
*,
49+
api_key: str | None = None,
50+
openai_client: AsyncOpenAI | None = None,
51+
http_client: AsyncHTTPClient | None = None,
52+
) -> None:
53+
api_key = api_key or os.getenv("AIML_API_KEY")
54+
if not api_key and openai_client is None:
55+
raise UserError(
56+
"Set the `AIML_API_KEY` environment variable or pass it via `AIMLAPIProvider(api_key=...)` "
57+
"to use the AI/ML API provider."
58+
)
59+
60+
if openai_client is not None:
61+
self._client = openai_client
62+
elif http_client is not None:
63+
self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
64+
else:
65+
http_client = cached_async_http_client(provider="aimlapi")
66+
self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

0 commit comments

Comments
 (0)