diff --git a/CHANGELOG.md b/CHANGELOG.md index a057ee8a6e..c630fc705e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ All notable changes to Parlant will be documented here. ### Added +- Add aimlapi.com as an NLP service (`parlant.adapters.nlp.aimlapi_service`), available as `NLPServices.aimlapi()` and `parlant-server --aimlapi`. Provides schematic generation, streaming and native embeddings over aimlapi.com's OpenAI-compatible API using the existing `openai` dependency; configured via `AIMLAPI_API_KEY`, `AIMLAPI_MODEL` and `AIMLAPI_EMBEDDER_MODEL` - Add `HealthReporter` (`parlant.core.health_reporter`) — a generic, per-process health-reporting service registered in the container. Subsystems call `report(kind, attributes)`; registered `HealthView` objects interpret reports per kind and contribute to the `/healthz` snapshot. Each kind has a configurable retention policy (`window` + `max_count`). Views declare `Criticality.CRITICAL` or `INFORMATIONAL`; only critical views feed the worst-of overall status rollup - Add `NLPHealthView` reporting NLP request health sliced by schema (success rate, p50/p95 latency, recent error breakdown) with configurable thresholds for `degraded`/`unhealthy` classification - Instrument `BaseSchematicGenerator.generate()` and `BaseEmbedder.embed()` to emit `nlp.request` / `nlp.embed` health reports on success and failure, providing dashboard visibility into LLM and embedding behavior across all adapters diff --git a/docs/adapters/nlp/aimlapi.md b/docs/adapters/nlp/aimlapi.md new file mode 100644 index 0000000000..13062139ca --- /dev/null +++ b/docs/adapters/nlp/aimlapi.md @@ -0,0 +1,120 @@ +# aimlapi.com Service Documentation + +The aimlapi.com service gives Parlant access to 350+ chat models and 15 embedding +models — OpenAI, Anthropic, Google, DeepSeek, Qwen, Llama and others — through a single +OpenAI-compatible API, with one key and one bill. + +Unlike most aggregators, aimlapi.com serves embeddings natively, so Parlant's vector +store does not need a separate provider or a local fallback embedder. + +## Prerequisites + +1. **Account**: sign up at [aimlapi.com](https://aimlapi.com) +2. **API key**: create one in the dashboard + +No extra package is needed — the adapter uses the `openai` client that Parlant already +depends on. + +## Quick Start + +```bash +export AIMLAPI_API_KEY="your-api-key-here" +``` + +```python +import parlant.sdk as p +from parlant.sdk import NLPServices + +async with p.Server(nlp_service=NLPServices.aimlapi) as server: + agent = await server.create_agent( + name="AI Assistant", + description="A helpful assistant powered by aimlapi.com.", + ) + # 🎉 Ready to use at http://localhost:8800 +``` + +Or from the CLI: + +```bash +parlant-server --aimlapi +``` + +## Environment Variables + +### Required + +| Variable | Description | +|----------|-------------| +| `AIMLAPI_API_KEY` | Your aimlapi.com API key | + +### Optional + +| Variable | Description | Default | +|----------|-------------|---------| +| `AIMLAPI_MODEL` | Chat model id | `openai/gpt-4.1` | +| `AIMLAPI_MAX_TOKENS` | Context window for a model the adapter does not know | `131072` | +| `AIMLAPI_EMBEDDER_MODEL` | Embedding model id | `openai/text-embedding-3-large` | +| `AIMLAPI_EMBEDDER_DIMENSIONS` | Override embedding dimensions | Looked up per model, else `1536` | +| `AIMLAPI_HTTP_REFERER` | Your app's URL, for analytics | Parlant's repository | +| `AIMLAPI_SITE_NAME` | Your app's name, for analytics | `Parlant` | + +## Models + +Parlant needs a model that reliably honours `response_format: {"type": "json_object"}`, +because guideline matching, tool calling and message generation all run through +schematic (JSON-schema) generation. These ids are pre-configured with their real +context windows: + +| Model | Context | Notes | +|-------|---------|-------| +| `openai/gpt-4.1` | 1,047,576 | Default | +| `openai/gpt-4.1-mini` | 1,000,000 | Cheaper, still structured-output capable | +| `anthropic/claude-sonnet-4.5` | 200,000 | Long-context reasoning | +| `google/gemini-2.5-flash` | 1,000,000 | Fast and inexpensive | + +Any other chat model id works too — set `AIMLAPI_MODEL` and, if the model's context +window differs from the 128K default, `AIMLAPI_MAX_TOKENS`. + +The catalog is public and needs no key: + +```bash +# `include=all` adds capabilities, modalities, pricing and providers, +# none of which appear in the default response. +curl 'https://api.aimlapi.com/v1/models?include=all' +``` + +Chat models are the entries whose `type` is `openai/chat-completions`; prefer a model +whose `capabilities` include `structured_output`. + +### Embedding models + +| Model | Dimensions | +|-------|------------| +| `openai/text-embedding-3-large` | 3072 (default) | +| `openai/text-embedding-3-small` | 1536 | +| `openai/text-embedding-ada-002` | 1536 | +| `alibaba/text-embedding-v4` | 1024 | + +⚠️ Changing the embedder model or its dimensions after data has been indexed requires +clearing the vector store, or you will hit dimension-mismatch errors. + +## Notes for Contributors + +**Unset optional parameters must be omitted, not sent as `null`.** aimlapi.com answers +`400` when `temperature`, `top_p`, `seed`, `tools`, `tool_choice`, `response_format`, +`stream`, `stream_options`, `parallel_tool_calls`, `max_tokens` or +`max_completion_tokens` arrive as an explicit JSON `null`, even though OpenAI accepts +`null` for all of them. The adapter routes every request through +`omit_unset_arguments()` for exactly this reason, and +`tests/adapters/nlp/test_aimlapi_service.py` guards it. The 400 body names the field in +`error.details[].path` / `.reason`; the top-level `message` is generic. + +## Troubleshooting + +**`AIMLAPI_API_KEY is not set`** — export the variable before starting the server. + +**`400 Bad Request` on every call** — check that no request field is being sent as +`null`; see the note above. + +**Dimension mismatch after switching embedder models** — clear the cached embeddings +in your `parlant-data` directory. diff --git a/src/parlant/adapters/nlp/aimlapi_service.py b/src/parlant/adapters/nlp/aimlapi_service.py new file mode 100644 index 0000000000..6e4f43392f --- /dev/null +++ b/src/parlant/adapters/nlp/aimlapi_service.py @@ -0,0 +1,695 @@ +# Copyright 2026 Emcie Co Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +import os +import re +import time +from types import MappingProxyType +from typing import Any, AsyncIterator, Callable, Mapping +from urllib.parse import urlparse + +import jsonfinder # type: ignore +import tiktoken +from openai import ( + APIConnectionError, + APIResponseValidationError, + APITimeoutError, + AsyncClient, + ConflictError, + InternalServerError, + RateLimitError, +) +from pydantic import ValidationError +from typing_extensions import override + +from parlant.adapters.nlp.common import normalize_json_output, record_llm_metrics +from parlant.core.engines.alpha.prompt_builder import PromptBuilder +from parlant.core.health import HealthReporter +from parlant.core.loggers import Logger +from parlant.core.meter import Meter +from parlant.core.nlp.embedding import BaseEmbedder, Embedder, EmbeddingResult +from parlant.core.nlp.generation import ( + T, + BaseSchematicGenerator, + BaseStreamingTextGenerator, + SchematicGenerationResult, + StreamingTextGenerator, +) +from parlant.core.nlp.generation_info import GenerationInfo, UsageInfo +from parlant.core.nlp.moderation import ModerationService, NoModeration +from parlant.core.nlp.policies import policy, retry +from parlant.core.nlp.service import ( + EmbedderHints, + NLPService, + SchematicGeneratorHints, + StreamingTextGeneratorHints, +) +from parlant.core.nlp.tokenization import EstimatingTokenizer +from parlant.core.tracer import Tracer + +AIMLAPI_BASE_URL = "https://api.aimlapi.com/v1" +AIMLAPI_HOST = "api.aimlapi.com" +AIMLAPI_DEFAULT_MODEL = "openai/gpt-4.1" +AIMLAPI_DEFAULT_EMBEDDER_MODEL = "openai/text-embedding-3-large" + +# aimlapi.com validates a documented set of optional chat-completion fields as +# "present but wrong type" when they are sent as an explicit JSON null, and answers +# 400 (`{"error": {"details": [{"path": "tools", "reason": "Expected array, received +# null"}]}}`). OpenAI itself accepts null for all of them, so a client that forwards +# an unset optional as None — the natural thing to write — fails on every request. +# +# The set below is what the provider currently rejects. It is kept as documentation; +# `omit_unset_arguments` drops *every* None, which is both simpler and strictly safer. +NULL_REJECTING_REQUEST_FIELDS = frozenset( + { + "temperature", + "top_p", + "seed", + "tools", + "tool_choice", + "response_format", + "stream", + "stream_options", + "parallel_tool_calls", + "max_tokens", + "max_completion_tokens", + } +) + +_ATTRIBUTION_HEADERS: Mapping[str, str] = MappingProxyType( + { + # HTTP-Referer and X-Title identify the *calling* application, not the provider. + "HTTP-Referer": "https://github.com/emcie-co/parlant", + "X-Title": "Parlant", + "X-AIMLAPI-Partner-ID": "part_UOT3mCwOdpOUQKX2gIvYrCmv", + "X-AIMLAPI-Source": "agent/parlant", + } +) + + +def omit_unset_arguments(arguments: Mapping[str, Any]) -> dict[str, Any]: + """Drops keys whose value is None instead of serializing them as JSON null. + + See NULL_REJECTING_REQUEST_FIELDS: aimlapi.com answers 400 for several optional + fields when they arrive as null, so "unset" has to mean "absent from the payload". + """ + return {k: v for k, v in arguments.items() if v is not None} + + +def build_attribution_headers(base_url: str = AIMLAPI_BASE_URL) -> dict[str, str]: + """Builds a fresh header dict for a request to aimlapi.com. + + Returns an empty dict for any other host, so that attribution cannot ride along + to a different provider or to a proxy that merely fronts the same API. + """ + if urlparse(base_url).hostname != AIMLAPI_HOST: + return {} + + headers = dict(_ATTRIBUTION_HEADERS) + + if referer := os.environ.get("AIMLAPI_HTTP_REFERER"): + headers["HTTP-Referer"] = referer + + if site_name := os.environ.get("AIMLAPI_SITE_NAME"): + headers["X-Title"] = site_name + + return headers + + +def _create_client() -> AsyncClient: + return AsyncClient( + base_url=AIMLAPI_BASE_URL, + api_key=os.environ["AIMLAPI_API_KEY"], + default_headers=build_attribution_headers(), + ) + + +class AIMLAPIEstimatingTokenizer(EstimatingTokenizer): + def __init__(self, model_name: str) -> None: + self.model_name = model_name + # aimlapi.com routes many vendors; gpt-4o's encoding is a good general estimate. + self.encoding = tiktoken.encoding_for_model("gpt-4o-2024-08-06") + + @override + async def estimate_token_count(self, prompt: str) -> int: + tokens = self.encoding.encode(prompt) + return len(tokens) + + +class AIMLAPISchematicGenerator(BaseSchematicGenerator[T]): + supported_aimlapi_params = ["temperature", "top_p", "logit_bias", "max_tokens"] + supported_hints = supported_aimlapi_params + ["strict"] + + def __init__( + self, + model_name: str, + logger: Logger, + tracer: Tracer, + meter: Meter, + health_reporter: HealthReporter, + ) -> None: + super().__init__( + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + model_name=model_name, + ) + + self._client = _create_client() + self._tokenizer = AIMLAPIEstimatingTokenizer(model_name=self.model_name) + + @property + @override + def id(self) -> str: + return f"aimlapi/{self.model_name}" + + @property + @override + def tokenizer(self) -> AIMLAPIEstimatingTokenizer: + return self._tokenizer + + @property + @override + def max_tokens(self) -> int: + return 128 * 1024 + + def build_request_arguments(self, hints: Mapping[str, Any]) -> dict[str, Any]: + return omit_unset_arguments( + {k: v for k, v in hints.items() if k in self.supported_aimlapi_params} + ) + + @policy( + [ + retry( + exceptions=( + APIConnectionError, + APITimeoutError, + ConflictError, + RateLimitError, + APIResponseValidationError, + ), + ), + retry(InternalServerError, max_exceptions=2, wait_times=(1.0, 5.0)), + ] + ) + @override + async def do_generate( + self, + prompt: str | PromptBuilder, + hints: Mapping[str, Any] = {}, + ) -> SchematicGenerationResult[T]: + with self.logger.scope(f"AI/ML API LLM Request ({self.schema.__name__})"): + return await self._do_generate(prompt, hints) + + async def _do_generate( + self, + prompt: str | PromptBuilder, + hints: Mapping[str, Any] = {}, + ) -> SchematicGenerationResult[T]: + if isinstance(prompt, PromptBuilder): + prompt = prompt.build() + + aimlapi_api_arguments = self.build_request_arguments(hints) + + t_start = time.time() + response = await self._client.chat.completions.create( + messages=[{"role": "user", "content": prompt}], + model=self.model_name, + response_format={"type": "json_object"}, + **omit_unset_arguments({"max_tokens": 8192, **aimlapi_api_arguments}), + ) + t_end = time.time() + + if response.usage: + self.logger.trace(response.usage.model_dump_json(indent=2)) + + raw_content = response.choices[0].message.content or "{}" + + try: + json_content = json.loads(normalize_json_output(raw_content)) + except json.JSONDecodeError: + self.logger.warning(f"Invalid JSON returned by {self.model_name}:\n{raw_content})") + json_content = jsonfinder.only_json(raw_content)[2] + self.logger.warning("Found JSON content within model response; continuing...") + + try: + content = self.schema.model_validate(json_content) + + assert response.usage + + cached_input_tokens = ( + getattr(response.usage, "prompt_cache_hit_tokens", 0) + or getattr( + getattr(response.usage, "prompt_tokens_details", None), "cached_tokens", 0 + ) + or 0 + ) + + await record_llm_metrics( + self.meter, + self.model_name, + schema_name=self.schema.__name__, + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + cached_input_tokens=cached_input_tokens, + ) + + return SchematicGenerationResult( + content=content, + info=GenerationInfo( + schema_name=self.schema.__name__, + model=self.id, + duration=(t_end - t_start), + usage=UsageInfo( + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + extra={"cached_input_tokens": cached_input_tokens}, + ), + ), + ) + except ValidationError: + self.logger.error( + f"JSON content returned by {self.model_name} does not match expected schema:\n{raw_content}" + ) + raise + + +class AIMLAPI_GPT_4_1(AIMLAPISchematicGenerator[T]): + def __init__( + self, logger: Logger, tracer: Tracer, meter: Meter, health_reporter: HealthReporter + ) -> None: + super().__init__( + model_name="openai/gpt-4.1", + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + ) + + @property + @override + def max_tokens(self) -> int: + return 1_047_576 + + +class AIMLAPI_GPT_4_1_Mini(AIMLAPISchematicGenerator[T]): + def __init__( + self, logger: Logger, tracer: Tracer, meter: Meter, health_reporter: HealthReporter + ) -> None: + super().__init__( + model_name="openai/gpt-4.1-mini", + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + ) + + @property + @override + def max_tokens(self) -> int: + return 1_000_000 + + +class AIMLAPI_ClaudeSonnet45(AIMLAPISchematicGenerator[T]): + def __init__( + self, logger: Logger, tracer: Tracer, meter: Meter, health_reporter: HealthReporter + ) -> None: + super().__init__( + model_name="anthropic/claude-sonnet-4.5", + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + ) + + @property + @override + def max_tokens(self) -> int: + return 200 * 1024 + + +class AIMLAPI_Gemini25Flash(AIMLAPISchematicGenerator[T]): + def __init__( + self, logger: Logger, tracer: Tracer, meter: Meter, health_reporter: HealthReporter + ) -> None: + super().__init__( + model_name="google/gemini-2.5-flash", + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + ) + + @property + @override + def max_tokens(self) -> int: + return 1_000_000 + + +# Pattern to detect word boundaries for chunking; matches after any whitespace character. +_WORD_BOUNDARY_PATTERN = re.compile(r"(?<=\s)") + +# Number of words to buffer before yielding a chunk. +_WORDS_PER_CHUNK = 3 + + +class AIMLAPIStreamingTextGenerator(BaseStreamingTextGenerator): + """Streaming text generator over aimlapi.com's OpenAI-compatible streaming API. + + Buffers tokens into word-sized chunks for smoother frontend rendering. + """ + + supported_aimlapi_params = ["temperature", "top_p", "max_tokens"] + + def __init__( + self, + model_name: str, + logger: Logger, + tracer: Tracer, + meter: Meter, + health_reporter: HealthReporter, + ) -> None: + super().__init__( + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + model_name=model_name, + ) + + self._client = _create_client() + self._tokenizer = AIMLAPIEstimatingTokenizer(model_name=self.model_name) + + @property + @override + def id(self) -> str: + return f"aimlapi-streaming/{self.model_name}" + + @property + @override + def tokenizer(self) -> AIMLAPIEstimatingTokenizer: + return self._tokenizer + + def build_request_arguments(self, hints: Mapping[str, Any]) -> dict[str, Any]: + return omit_unset_arguments( + {k: v for k, v in hints.items() if k in self.supported_aimlapi_params} + ) + + @override + async def do_generate( + self, + prompt: str | PromptBuilder, + hints: Mapping[str, Any] = {}, + ) -> tuple[AsyncIterator[str | None], Callable[[], UsageInfo]]: + if isinstance(prompt, PromptBuilder): + prompt = prompt.build() + + aimlapi_api_arguments = self.build_request_arguments(hints) + + stream = await self._client.chat.completions.create( + messages=[{"role": "user", "content": prompt}], + model=self.model_name, + stream=True, + stream_options={"include_usage": True}, + **aimlapi_api_arguments, + ) + + usage_info: UsageInfo | None = None + + async def chunk_generator() -> AsyncIterator[str | None]: + nonlocal usage_info + + buffer = "" + + async for chunk in stream: + if chunk.usage is not None: + self.logger.trace(chunk.usage.model_dump_json(indent=2)) + + cached_tokens = getattr(chunk.usage, "prompt_cache_hit_tokens", 0) or 0 + + usage_info = UsageInfo( + input_tokens=chunk.usage.prompt_tokens, + output_tokens=chunk.usage.completion_tokens, + extra={"cached_input_tokens": cached_tokens}, + ) + + if chunk.choices and chunk.choices[0].delta.content: + buffer += chunk.choices[0].delta.content + + boundaries = list(_WORD_BOUNDARY_PATTERN.finditer(buffer)) + if len(boundaries) >= _WORDS_PER_CHUNK: + last_boundary = boundaries[_WORDS_PER_CHUNK - 1] + chunk_text = buffer[: last_boundary.end()] + buffer = buffer[last_boundary.end() :] + yield chunk_text + + if buffer: + yield buffer + + if usage_info is not None: + await record_llm_metrics( + self.meter, + self.model_name, + schema_name="streaming", + input_tokens=usage_info.input_tokens, + output_tokens=usage_info.output_tokens, + cached_input_tokens=usage_info.extra.get("cached_input_tokens", 0) + if usage_info.extra + else 0, + ) + + yield None + + def get_usage() -> UsageInfo: + if usage_info is None: + return UsageInfo(input_tokens=0, output_tokens=0) + return usage_info + + return chunk_generator(), get_usage + + +class AIMLAPIEmbedder(BaseEmbedder): + supported_arguments = ["dimensions"] + + _KNOWN_DIMENSIONS: Mapping[str, int] = MappingProxyType( + { + "openai/text-embedding-3-large": 3072, + "openai/text-embedding-3-small": 1536, + "openai/text-embedding-ada-002": 1536, + "alibaba/text-embedding-v4": 1024, + } + ) + + def __init__( + self, + model_name: str, + logger: Logger, + tracer: Tracer, + meter: Meter, + health_reporter: HealthReporter, + ) -> None: + super().__init__(logger, tracer, meter, model_name, health_reporter) + + self._client = _create_client() + self._tokenizer = AIMLAPIEstimatingTokenizer(model_name=self.model_name) + + @property + @override + def id(self) -> str: + return f"aimlapi/{self.model_name}" + + @property + @override + def tokenizer(self) -> AIMLAPIEstimatingTokenizer: + return self._tokenizer + + @property + @override + def max_tokens(self) -> int: + return 8192 + + @property + @override + def dimensions(self) -> int: + if dimensions := os.environ.get("AIMLAPI_EMBEDDER_DIMENSIONS"): + return int(dimensions) + + return self._KNOWN_DIMENSIONS.get(self.model_name, 1536) + + @policy( + [ + retry( + exceptions=( + APIConnectionError, + APITimeoutError, + ConflictError, + RateLimitError, + APIResponseValidationError, + ), + ), + retry(InternalServerError, max_exceptions=2, wait_times=(1.0, 5.0)), + ] + ) + @override + async def do_embed( + self, + texts: list[str], + hints: Mapping[str, Any] = {}, + ) -> EmbeddingResult: + filtered_hints = omit_unset_arguments( + {k: v for k, v in hints.items() if k in self.supported_arguments} + ) + + response = await self._client.embeddings.create( + model=self.model_name, + input=texts, + **filtered_hints, + ) + + vectors = [data_point.embedding for data_point in response.data] + + return EmbeddingResult(vectors=vectors) + + +class AIMLAPIService(NLPService): + @staticmethod + def verify_environment() -> str | None: + """Returns an error message if the environment is not set up correctly.""" + + if not os.environ.get("AIMLAPI_API_KEY"): + return """\ +You're using the aimlapi.com NLP service, but AIMLAPI_API_KEY is not set. +Please set AIMLAPI_API_KEY in your environment before running Parlant. +""" + + return None + + def __init__( + self, + logger: Logger, + tracer: Tracer, + meter: Meter, + health_reporter: HealthReporter, + ) -> None: + self._logger = logger + self._tracer = tracer + self._meter = meter + self._health_reporter = health_reporter + + self.model_name = os.environ.get("AIMLAPI_MODEL", AIMLAPI_DEFAULT_MODEL) + self.embedder_model_name = os.environ.get( + "AIMLAPI_EMBEDDER_MODEL", AIMLAPI_DEFAULT_EMBEDDER_MODEL + ) + + self._logger.info(f"Initialized AIMLAPIService with model: {self.model_name}") + self._logger.info(f"aimlapi.com embedder model name: {self.embedder_model_name}") + + embedder_model = self.embedder_model_name + + class DynamicAIMLAPIEmbedder(AIMLAPIEmbedder): + def __init__( + self, + logger: Logger, + tracer: Tracer, + meter: Meter, + health_reporter: HealthReporter, + ) -> None: + super().__init__( + model_name=embedder_model, + logger=logger, + tracer=tracer, + meter=meter, + health_reporter=health_reporter, + ) + + self._embedder_class = DynamicAIMLAPIEmbedder + + @property + @override + def supports_streaming(self) -> bool: + return True + + @override + async def get_streaming_text_generator( + self, hints: StreamingTextGeneratorHints = {} + ) -> StreamingTextGenerator: + return AIMLAPIStreamingTextGenerator( + model_name=self.model_name, + logger=self._logger, + tracer=self._tracer, + meter=self._meter, + health_reporter=self._health_reporter, + ) + + def _get_specialized_generator_class( + self, + model_name: str, + schema_type: type[T], + ) -> Callable[[Logger, Tracer, Meter, HealthReporter], AIMLAPISchematicGenerator[T]] | None: + """Returns the specialized generator class for known models, or None for custom models.""" + model_to_class: dict[ + str, Callable[[Logger, Tracer, Meter, HealthReporter], AIMLAPISchematicGenerator[T]] + ] = { + "openai/gpt-4.1": AIMLAPI_GPT_4_1[schema_type], # type: ignore + "openai/gpt-4.1-mini": AIMLAPI_GPT_4_1_Mini[schema_type], # type: ignore + "anthropic/claude-sonnet-4.5": AIMLAPI_ClaudeSonnet45[schema_type], # type: ignore + "google/gemini-2.5-flash": AIMLAPI_Gemini25Flash[schema_type], # type: ignore + } + + return model_to_class.get(model_name) + + @override + async def get_schematic_generator( + self, t: type[T], hints: SchematicGeneratorHints = {} + ) -> AIMLAPISchematicGenerator[T]: + if specialized_class := self._get_specialized_generator_class( + self.model_name, schema_type=t + ): + self._logger.debug(f"Using specialized generator for model: {self.model_name}") + return specialized_class(self._logger, self._tracer, self._meter, self._health_reporter) + + self._logger.debug(f"Using custom generator for model: {self.model_name}") + + max_tokens = int(os.environ.get("AIMLAPI_MAX_TOKENS", 128 * 1024)) + + class DynamicAIMLAPISchematicGenerator(AIMLAPISchematicGenerator[T]): + @property + @override + def max_tokens(self) -> int: + return max_tokens + + return DynamicAIMLAPISchematicGenerator[t]( # type: ignore + model_name=self.model_name, + logger=self._logger, + tracer=self._tracer, + meter=self._meter, + health_reporter=self._health_reporter, + ) + + @override + async def get_embedder(self, hints: EmbedderHints = {}) -> Embedder: + return self._embedder_class( + logger=self._logger, + tracer=self._tracer, + meter=self._meter, + health_reporter=self._health_reporter, + ) + + @override + async def get_moderation_service(self) -> ModerationService: + return NoModeration() diff --git a/src/parlant/bin/server.py b/src/parlant/bin/server.py index 76ff6899e0..ee7198311f 100755 --- a/src/parlant/bin/server.py +++ b/src/parlant/bin/server.py @@ -276,6 +276,7 @@ def __init__(self, message: str) -> None: NLPServiceName = Literal[ + "aimlapi", "anthropic", "aws", "azure", @@ -410,6 +411,16 @@ def load_novita(container: Container) -> NLPService: ) +def load_aimlapi(container: Container) -> NLPService: + return load_nlp_service( + container, + "aimlapi.com", + "aimlapi", + "AIMLAPIService", + "parlant.adapters.nlp.aimlapi_service", + ) + + def load_litellm(container: Container) -> NLPService: from parlant.adapters.nlp.litellm_service import LiteLLMService @@ -432,6 +443,7 @@ def load_litellm(container: Container) -> NLPService: NLP_SERVICE_INITIALIZERS: dict[NLPServiceName, Callable[[Container], NLPService]] = { + "aimlapi": load_aimlapi, "anthropic": load_anthropic, "aws": load_aws, "azure": load_azure, @@ -1174,6 +1186,12 @@ def transform_and_exec_help(command: str) -> None: default=DEFAULT_PORT, help="Server port", ) + @click.option( + "--aimlapi", + is_flag=True, + help="Run with aimlapi.com. The environment variable AIMLAPI_API_KEY must be set.", + default=False, + ) @click.option( "--openai", is_flag=True, @@ -1300,6 +1318,7 @@ def run( cerebras: bool, together: bool, novita: bool, + aimlapi: bool, litellm: bool, modelscope: bool, log_level: str, @@ -1323,6 +1342,7 @@ def run( cerebras, together, novita, + aimlapi, litellm, modelscope, ] @@ -1342,6 +1362,7 @@ def run( cerebras, together, novita, + aimlapi, litellm, modelscope, ) @@ -1377,6 +1398,9 @@ def run( elif novita: nlp_service = "novita" require_env_keys(["NOVITA_API_KEY"]) + elif aimlapi: + nlp_service = "aimlapi" + require_env_keys(["AIMLAPI_API_KEY"]) elif litellm: nlp_service = "litellm" require_env_keys(["LITELLM_PROVIDER_MODEL_NAME"]) diff --git a/src/parlant/sdk.py b/src/parlant/sdk.py index 60518b567d..3c99baf0df 100644 --- a/src/parlant/sdk.py +++ b/src/parlant/sdk.py @@ -370,6 +370,18 @@ def emcie(container: Container) -> NLPService: container[HealthReporter], ) + @staticmethod + def aimlapi(container: Container) -> NLPService: + """Creates an aimlapi.com NLPService instance using the provided container.""" + from parlant.adapters.nlp.aimlapi_service import AIMLAPIService + + if error := AIMLAPIService.verify_environment(): + raise NLPServiceConfigurationError(error) + + return AIMLAPIService( + container[Logger], container[Tracer], container[Meter], container[HealthReporter] + ) + @staticmethod def azure(container: Container) -> NLPService: """Creates an Azure NLPService instance using the provided container.""" diff --git a/tests/adapters/nlp/test_aimlapi_service.py b/tests/adapters/nlp/test_aimlapi_service.py new file mode 100644 index 0000000000..b54d4859aa --- /dev/null +++ b/tests/adapters/nlp/test_aimlapi_service.py @@ -0,0 +1,331 @@ +# Copyright 2026 Emcie Co Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import os +import re +from collections.abc import Generator +from typing import Any +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from lagom import Container +from openai.types.chat import ChatCompletion, ChatCompletionMessage +from openai.types.chat.chat_completion import Choice +from openai.types.completion_usage import CompletionUsage + +from parlant.adapters.nlp.aimlapi_service import ( + AIMLAPI_BASE_URL, + AIMLAPI_DEFAULT_EMBEDDER_MODEL, + AIMLAPI_DEFAULT_MODEL, + NULL_REJECTING_REQUEST_FIELDS, + AIMLAPIEmbedder, + AIMLAPIEstimatingTokenizer, + AIMLAPISchematicGenerator, + AIMLAPIService, + AIMLAPIStreamingTextGenerator, + AIMLAPI_GPT_4_1, + build_attribution_headers, + omit_unset_arguments, +) +from parlant.core.common import DefaultBaseModel +from parlant.core.health import HealthReporter +from parlant.core.loggers import Logger +from parlant.core.meter import Meter +from parlant.core.tracer import Tracer + +PARTNER_ID_PATTERN = re.compile(r"^part_[A-Za-z0-9]{1,64}$") + + +class SchemaData(DefaultBaseModel): + """Test schema for type checking.""" + + test_field: str = "test_value" + + +@pytest.fixture(autouse=True) +def set_api_keys() -> Generator[None, None, None]: + with patch.dict( + os.environ, + {"OPENAI_API_KEY": "test-openai-key", "AIMLAPI_API_KEY": "test-aimlapi-key"}, + clear=False, + ): + yield + + +def test_that_missing_aimlapi_api_key_returns_error_message() -> None: + with patch.dict(os.environ, {}, clear=True): + error = AIMLAPIService.verify_environment() + assert error is not None + assert "AIMLAPI_API_KEY is not set" in error + + +def test_that_present_api_key_returns_none() -> None: + with patch.dict(os.environ, {"AIMLAPI_API_KEY": "test-key"}, clear=True): + assert AIMLAPIService.verify_environment() is None + + +def test_that_aimlapi_service_initializes_with_default_models() -> None: + with patch.dict(os.environ, {"AIMLAPI_API_KEY": "test-key"}, clear=True): + service = AIMLAPIService(logger=Mock(), tracer=Mock(), meter=Mock(), health_reporter=Mock()) + assert service.model_name == AIMLAPI_DEFAULT_MODEL + assert service.embedder_model_name == AIMLAPI_DEFAULT_EMBEDDER_MODEL + + +def test_that_aimlapi_service_uses_environment_model() -> None: + with patch.dict( + os.environ, + {"AIMLAPI_API_KEY": "test-key", "AIMLAPI_MODEL": "anthropic/claude-sonnet-4.5"}, + clear=True, + ): + service = AIMLAPIService(logger=Mock(), tracer=Mock(), meter=Mock(), health_reporter=Mock()) + assert service.model_name == "anthropic/claude-sonnet-4.5" + + +def test_that_aimlapi_estimating_tokenizer_counts_tokens() -> None: + tokenizer = AIMLAPIEstimatingTokenizer(model_name=AIMLAPI_DEFAULT_MODEL) + assert asyncio.run(tokenizer.estimate_token_count("Hello world")) > 0 + + +def test_that_gpt_4_1_generator_reports_its_full_context_window(container: Container) -> None: + generator = AIMLAPI_GPT_4_1[SchemaData]( + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + assert generator.model_name == "openai/gpt-4.1" + assert generator.id == "aimlapi/openai/gpt-4.1" + assert generator.max_tokens == 1_047_576 + + +def test_that_the_partner_id_matches_the_attribution_contract() -> None: + # A malformed partner id is silently ignored by the gateway rather than rejected, + # so nothing at runtime would ever surface a typo here. + headers = build_attribution_headers() + assert PARTNER_ID_PATTERN.match(headers["X-AIMLAPI-Partner-ID"]) + assert headers["X-AIMLAPI-Source"] == "agent/parlant" + + +def test_that_attribution_headers_identify_parlant_and_not_the_provider() -> None: + headers = build_attribution_headers() + assert headers["HTTP-Referer"] == "https://github.com/emcie-co/parlant" + assert headers["X-Title"] == "Parlant" + + +def test_that_attribution_headers_are_not_sent_to_another_host() -> None: + assert build_attribution_headers("https://api.openai.com/v1") == {} + assert build_attribution_headers("https://aimlapi-proxy.example.com/v1") == {} + assert build_attribution_headers(AIMLAPI_BASE_URL) != {} + + +def test_that_attribution_headers_are_a_fresh_dict_per_call() -> None: + first = build_attribution_headers() + first["X-Title"] = "mutated" + assert build_attribution_headers()["X-Title"] == "Parlant" + + +def test_that_referer_and_title_can_be_overridden_by_the_host_application() -> None: + with patch.dict( + os.environ, + { + "AIMLAPI_API_KEY": "test-key", + "AIMLAPI_HTTP_REFERER": "https://myapp.example", + "AIMLAPI_SITE_NAME": "My App", + }, + clear=True, + ): + headers = build_attribution_headers() + assert headers["HTTP-Referer"] == "https://myapp.example" + assert headers["X-Title"] == "My App" + # Attribution itself is not overridable by the host. + assert headers["X-AIMLAPI-Partner-ID"] == "part_UOT3mCwOdpOUQKX2gIvYrCmv" + + +@patch("parlant.adapters.nlp.aimlapi_service.AsyncClient") +def test_that_the_client_is_created_against_aimlapi_with_attribution_headers( + mock_client_class: Mock, +) -> None: + _ = AIMLAPISchematicGenerator[SchemaData]( + model_name=AIMLAPI_DEFAULT_MODEL, + logger=Mock(), + tracer=Mock(), + meter=Mock(), + health_reporter=Mock(), + ) + + call_kwargs = mock_client_class.call_args[1] + assert call_kwargs["base_url"] == AIMLAPI_BASE_URL + assert call_kwargs["default_headers"]["X-AIMLAPI-Partner-ID"] == "part_UOT3mCwOdpOUQKX2gIvYrCmv" + assert call_kwargs["default_headers"]["X-AIMLAPI-Source"] == "agent/parlant" + + +def test_that_unset_arguments_are_omitted_rather_than_passed_as_null() -> None: + # aimlapi.com answers 400 for these fields when they arrive as an explicit null, + # while OpenAI accepts null — so this must be checked here and not just in review. + unset = {field: None for field in NULL_REJECTING_REQUEST_FIELDS} + assert omit_unset_arguments({**unset, "model": "openai/gpt-4.1"}) == {"model": "openai/gpt-4.1"} + + +def _completion(content: str) -> Mock: + response = Mock(spec=ChatCompletion) + response.choices = [ + Choice( + message=ChatCompletionMessage(role="assistant", content=content), + finish_reason="stop", + index=0, + ) + ] + response.usage = CompletionUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + return response + + +async def test_that_a_none_valued_hint_is_never_sent_to_the_api(container: Container) -> None: + with patch("parlant.adapters.nlp.aimlapi_service.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.chat.completions.create = AsyncMock( + return_value=_completion('{"test_field": "test_value"}') + ) + mock_client_class.return_value = mock_client + + generator = AIMLAPISchematicGenerator[SchemaData]( + model_name=AIMLAPI_DEFAULT_MODEL, + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + + await generator.do_generate("Generate something", hints={"temperature": None}) + + request_kwargs: dict[str, Any] = mock_client.chat.completions.create.call_args[1] + assert "temperature" not in request_kwargs + assert [k for k, v in request_kwargs.items() if v is None] == [] + + +async def test_that_a_set_hint_is_forwarded_to_the_api(container: Container) -> None: + with patch("parlant.adapters.nlp.aimlapi_service.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.chat.completions.create = AsyncMock( + return_value=_completion('{"test_field": "test_value"}') + ) + mock_client_class.return_value = mock_client + + generator = AIMLAPISchematicGenerator[SchemaData]( + model_name=AIMLAPI_DEFAULT_MODEL, + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + + await generator.do_generate("Generate something", hints={"temperature": 0.0}) + + request_kwargs: dict[str, Any] = mock_client.chat.completions.create.call_args[1] + assert request_kwargs["temperature"] == 0.0 + assert request_kwargs["response_format"] == {"type": "json_object"} + + +async def test_that_no_none_valued_argument_is_sent_when_streaming(container: Container) -> None: + with patch("parlant.adapters.nlp.aimlapi_service.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.chat.completions.create = AsyncMock(return_value=AsyncMock()) + mock_client_class.return_value = mock_client + + generator = AIMLAPIStreamingTextGenerator( + model_name=AIMLAPI_DEFAULT_MODEL, + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + + await generator.do_generate("Say hello", hints={"temperature": None, "max_tokens": 64}) + + request_kwargs: dict[str, Any] = mock_client.chat.completions.create.call_args[1] + assert "temperature" not in request_kwargs + assert request_kwargs["max_tokens"] == 64 + assert [k for k, v in request_kwargs.items() if v is None] == [] + + +async def test_that_the_generator_parses_a_successful_response(container: Container) -> None: + with patch("parlant.adapters.nlp.aimlapi_service.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.chat.completions.create = AsyncMock( + return_value=_completion('{"test_field": "test_value"}') + ) + mock_client_class.return_value = mock_client + + generator = AIMLAPISchematicGenerator[SchemaData]( + model_name=AIMLAPI_DEFAULT_MODEL, + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + + result = await generator.do_generate('Generate {"test_field": "test_value"}') + + assert result.content.test_field == "test_value" + assert result.info.usage.input_tokens == 10 + assert result.info.usage.output_tokens == 20 + + +def test_that_the_default_embedder_reports_its_known_dimensions(container: Container) -> None: + embedder = AIMLAPIEmbedder( + model_name=AIMLAPI_DEFAULT_EMBEDDER_MODEL, + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + assert embedder.dimensions == 3072 + assert embedder.id == "aimlapi/openai/text-embedding-3-large" + + +def test_that_the_service_returns_the_specialized_generator_for_the_default_model( + container: Container, +) -> None: + with patch.dict(os.environ, {"AIMLAPI_API_KEY": "test-key"}, clear=True): + service = AIMLAPIService( + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + generator = asyncio.run(service.get_schematic_generator(SchemaData)) + assert isinstance(generator, AIMLAPISchematicGenerator) + assert generator.model_name == AIMLAPI_DEFAULT_MODEL + assert generator.max_tokens == 1_047_576 + + +def test_that_an_unknown_model_falls_back_to_a_dynamic_generator(container: Container) -> None: + with patch.dict( + os.environ, + { + "AIMLAPI_API_KEY": "test-key", + "AIMLAPI_MODEL": "some-vendor/some-model", + "AIMLAPI_MAX_TOKENS": "4096", + }, + clear=True, + ): + service = AIMLAPIService( + logger=container[Logger], + tracer=container[Tracer], + meter=container[Meter], + health_reporter=container[HealthReporter], + ) + generator = asyncio.run(service.get_schematic_generator(SchemaData)) + assert generator.model_name == "some-vendor/some-model" + assert generator.max_tokens == 4096