diff --git a/.env.example b/.env.example index 225f736c..982f9225 100644 --- a/.env.example +++ b/.env.example @@ -40,3 +40,6 @@ MINIMAX_API_KEY='xxxxx' # Optional: MiniMax API base URL (default: https://api.minimax.io/v1) # For mainland China users: https://api.minimaxi.com/v1 # MINIMAX_BASE_URL='https://api.minimax.io/v1' +AIMLAPI_API_KEY='xxxxx' +# Optional: aimlapi.com API base URL (default: https://api.aimlapi.com/v1) +# AIMLAPI_BASE_URL='https://api.aimlapi.com/v1' diff --git a/README.md b/README.md index 794d6c4c..fb2bfdfb 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,7 @@ OPENROUTER_API_KEY='optional' TOGETHER_API_KEY='optional' GOOGLE_API_KEY='optional' ANTHROPIC_API_KEY='optional' +AIMLAPI_API_KEY='optional' ``` @@ -242,6 +243,7 @@ provider_server_address = # Typically ignored or can be left blank when is_local | Provider | `provider_name` | Local? | Description | API Key Link (Examples) | |--------------|-----------------|--------|---------------------------------------------------|---------------------------------------------| +| aimlapi.com | `aimlapi` | No | One API for 350+ chat models from many vendors (e.g. `deepseek/deepseek-v4-flash`, `anthropic/claude-sonnet-4.6`). Model ids are listed at [api.aimlapi.com/v1/models](https://api.aimlapi.com/v1/models). | [aimlapi.com/app/keys](https://aimlapi.com/app/keys) | | OpenAI | `openai` | No | Use ChatGPT models via OpenAI's API. | [platform.openai.com/signup](https://platform.openai.com/signup) | | Google Gemini| `google` | No | Use Google Gemini models via Google AI Studio. | [aistudio.google.com/keys](https://aistudio.google.com/keys) | | Deepseek | `deepseek` | No | Use Deepseek models via their API. | [platform.deepseek.com](https://platform.deepseek.com) | diff --git a/sources/llm_provider.py b/sources/llm_provider.py index eb1c7560..b66e04a3 100644 --- a/sources/llm_provider.py +++ b/sources/llm_provider.py @@ -14,6 +14,27 @@ from sources.logger import Logger from sources.utility import pretty_print, animate_thinking +AIMLAPI_DEFAULT_BASE_URL = "https://api.aimlapi.com/v1" +AIMLAPI_ATTRIBUTION_HEADERS = { + "HTTP-Referer": "https://github.com/Fosowl/agenticSeek", + "X-Title": "AgenticSeek", + "X-AIMLAPI-Partner-ID": "part_l9PWfCWCyHgneiq7GLfDHUsf", + "X-AIMLAPI-Source": "agent/agenticseek", +} + +def aimlapi_attribution_headers(base_url: str) -> dict: + """ + Attribution headers identifying AgenticSeek to aimlapi.com. + + Keyed on the request origin, not on the provider name: if AIMLAPI_BASE_URL + points somewhere else (a proxy, a self-hosted gateway) we send nothing, so + these can never ride a request to a third party. Returns a fresh dict every + call so the module level constant is never mutated. + """ + if urlparse(base_url).hostname != "api.aimlapi.com": + return {} + return dict(AIMLAPI_ATTRIBUTION_HEADERS) + class Provider: def __init__(self, provider_name, model, server_address="127.0.0.1:5000", is_local=False): self.provider_name = provider_name.lower() @@ -38,12 +59,13 @@ def __init__(self, provider_name, model, server_address="127.0.0.1:5000", is_loc "anthropic": self.anthropic_fn, "minimax": self.minimax_fn, "litellm": self.litellm_fn, + "aimlapi": self.aimlapi_fn, "test": self.test_fn } self.logger = Logger("provider.log") self.api_key = None self.internal_url, self.in_docker = self.get_internal_url() - self.unsafe_providers = ["openai", "deepseek", "dsk_deepseek", "together", "google", "openrouter", "anthropic", "minimax"] + self.unsafe_providers = ["openai", "deepseek", "dsk_deepseek", "together", "google", "openrouter", "anthropic", "minimax", "aimlapi"] if self.provider_name not in self.available_providers: raise ValueError(f"Unknown provider: {provider_name}") if self.provider_name in self.unsafe_providers and self.is_local == False: @@ -541,6 +563,43 @@ def litellm_fn(self, history, verbose=False): except Exception as e: raise Exception(f"LiteLLM API error: {str(e)}") from e + def aimlapi_fn(self, history, verbose=False): + """ + Use aimlapi.com (AI/ML API) to generate text through its + OpenAI-compatible endpoint. + + Set AIMLAPI_API_KEY in your .env. provider_model is any chat model id + listed by https://api.aimlapi.com/v1/models (filter on + type == "openai/chat-completions"), e.g. deepseek/deepseek-v4-flash. + """ + if self.is_local: + raise Exception("aimlapi.com is not available for local use. Change config.ini") + load_dotenv() + base_url = os.getenv("AIMLAPI_BASE_URL", AIMLAPI_DEFAULT_BASE_URL) + + client = OpenAI( + api_key=self.api_key, + base_url=base_url, + default_headers=aimlapi_attribution_headers(base_url), + ) + # Optional parameters are omitted, never sent as None: aimlapi.com + # rejects an explicit null for temperature/top_p/tools/... with a 400 on + # some of its models, and the SDK serialises an unset optional as null. + params = { + "model": self.model, + "messages": history, + } + try: + response = client.chat.completions.create(**params) + if response is None: + raise Exception("aimlapi.com response is empty.") + thought = response.choices[0].message.content + if verbose: + print(thought) + return thought + except Exception as e: + raise Exception(f"aimlapi.com API error: {str(e)}") from e + def test_fn(self, history, verbose=True): """ This function is used to conduct tests. diff --git a/tests/test_aimlapi_provider.py b/tests/test_aimlapi_provider.py new file mode 100644 index 00000000..724391c6 --- /dev/null +++ b/tests/test_aimlapi_provider.py @@ -0,0 +1,171 @@ +import os +import re +import sys +import unittest +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from sources.llm_provider import ( + AIMLAPI_ATTRIBUTION_HEADERS, + AIMLAPI_DEFAULT_BASE_URL, + Provider, + aimlapi_attribution_headers, +) + +FAKE_ENV = {"AIMLAPI_API_KEY": "sk-test-123"} + + +def make_provider(is_local=False): + return Provider("aimlapi", "deepseek/deepseek-v4-flash", is_local=is_local) + + +def mocked_openai(mock_openai, content="42"): + """Wire an OpenAI mock so chat.completions.create returns `content`.""" + client = MagicMock() + client.chat.completions.create.return_value = MagicMock( + choices=[MagicMock(message=MagicMock(content=content))] + ) + mock_openai.return_value = client + return client + + +@patch.dict(os.environ, FAKE_ENV) +class TestAimlapiProvider(unittest.TestCase): + """Test cases for the aimlapi.com provider integration.""" + + def test_aimlapi_provider_registered(self): + """aimlapi is registered in available_providers.""" + self.assertIn("aimlapi", make_provider().available_providers) + + def test_aimlapi_is_unsafe_provider(self): + """aimlapi is a cloud API, so it must trigger the cloud data warning.""" + self.assertIn("aimlapi", make_provider().unsafe_providers) + + def test_aimlapi_reads_aimlapi_api_key(self): + """The key is read from AIMLAPI_API_KEY.""" + self.assertEqual(make_provider().api_key, "sk-test-123") + + def test_aimlapi_local_not_supported(self): + """aimlapi_fn refuses is_local=True.""" + provider = make_provider(is_local=True) + with self.assertRaises(Exception) as ctx: + provider.aimlapi_fn([{"role": "user", "content": "hi"}]) + self.assertIn("not available for local use", str(ctx.exception)) + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_returns_content(self, mock_openai): + """aimlapi_fn returns the assistant message content.""" + mocked_openai(mock_openai) + result = make_provider().aimlapi_fn([{"role": "user", "content": "What is 6*7?"}]) + self.assertEqual(result, "42") + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_uses_default_base_url(self, mock_openai): + """The OpenAI-compatible endpoint is https://api.aimlapi.com/v1.""" + mocked_openai(mock_openai) + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + self.assertEqual(mock_openai.call_args[1]["base_url"], AIMLAPI_DEFAULT_BASE_URL) + + @patch.dict(os.environ, {"AIMLAPI_BASE_URL": "https://proxy.example.com/v1"}) + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_base_url_is_overridable(self, mock_openai): + """AIMLAPI_BASE_URL overrides the endpoint.""" + mocked_openai(mock_openai) + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + self.assertEqual(mock_openai.call_args[1]["base_url"], "https://proxy.example.com/v1") + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_passes_model_and_messages(self, mock_openai): + """The configured model and the history reach the API unchanged.""" + client = mocked_openai(mock_openai) + history = [{"role": "user", "content": "hi"}] + make_provider().aimlapi_fn(history) + call_kwargs = client.chat.completions.create.call_args[1] + self.assertEqual(call_kwargs["model"], "deepseek/deepseek-v4-flash") + self.assertEqual(call_kwargs["messages"], history) + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_omits_unset_parameters(self, mock_openai): + """ + Unset optional parameters are omitted, never sent as None. + + aimlapi.com answers 400 for an explicit null on temperature, top_p, + seed, tools, tool_choice, response_format, stream and max_tokens on some + of its models, so passing None would fail live while mocks stay green. + """ + client = mocked_openai(mock_openai) + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + call_kwargs = client.chat.completions.create.call_args[1] + self.assertEqual( + [k for k, v in call_kwargs.items() if v is None], + [], + "no request parameter may be sent as None", + ) + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_sends_attribution_headers(self, mock_openai): + """All four attribution headers are attached to the client.""" + mocked_openai(mock_openai) + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + headers = mock_openai.call_args[1]["default_headers"] + self.assertEqual(headers["HTTP-Referer"], "https://github.com/Fosowl/agenticSeek") + self.assertEqual(headers["X-Title"], "AgenticSeek") + self.assertEqual(headers["X-AIMLAPI-Source"], "agent/agenticseek") + self.assertEqual(headers["X-AIMLAPI-Partner-ID"], "part_l9PWfCWCyHgneiq7GLfDHUsf") + + def test_partner_id_matches_gateway_pattern(self): + """ + A malformed partner id is dropped silently by the gateway, so its shape + is asserted here rather than discovered in production. + """ + self.assertRegex( + AIMLAPI_ATTRIBUTION_HEADERS["X-AIMLAPI-Partner-ID"], + re.compile(r"^part_[A-Za-z0-9]{1,64}$"), + ) + + def test_source_matches_gateway_pattern(self): + """X-AIMLAPI-Source is / with channel in web|agent|mcp.""" + self.assertRegex( + AIMLAPI_ATTRIBUTION_HEADERS["X-AIMLAPI-Source"], + re.compile(r"^(web|agent|mcp)/[a-z0-9-]{1,32}$"), + ) + + def test_attribution_headers_are_scoped_to_our_origin(self): + """Attribution never rides a request to another host.""" + self.assertEqual(aimlapi_attribution_headers(AIMLAPI_DEFAULT_BASE_URL), + AIMLAPI_ATTRIBUTION_HEADERS) + self.assertEqual(aimlapi_attribution_headers("https://proxy.example.com/v1"), {}) + + def test_attribution_headers_constant_is_not_mutated(self): + """Each call gets a fresh dict; the module level constant is immutable.""" + first = aimlapi_attribution_headers(AIMLAPI_DEFAULT_BASE_URL) + first["X-Title"] = "tampered" + self.assertIsNot(first, AIMLAPI_ATTRIBUTION_HEADERS) + self.assertEqual(AIMLAPI_ATTRIBUTION_HEADERS["X-Title"], "AgenticSeek") + self.assertEqual(aimlapi_attribution_headers(AIMLAPI_DEFAULT_BASE_URL)["X-Title"], + "AgenticSeek") + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_raises_on_empty_response(self, mock_openai): + """An empty response is reported instead of raising an IndexError.""" + client = MagicMock() + client.chat.completions.create.return_value = None + mock_openai.return_value = client + with self.assertRaises(Exception) as ctx: + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + self.assertIn("empty", str(ctx.exception).lower()) + + @patch('sources.llm_provider.OpenAI') + def test_aimlapi_fn_raises_on_api_error(self, mock_openai): + """API errors are wrapped with the provider name.""" + client = MagicMock() + client.chat.completions.create.side_effect = Exception("rate limit exceeded") + mock_openai.return_value = client + with self.assertRaises(Exception) as ctx: + make_provider().aimlapi_fn([{"role": "user", "content": "hi"}]) + self.assertIn("aimlapi.com API error", str(ctx.exception)) + + +if __name__ == '__main__': + unittest.main()