ff-aitoolkit 0.2.0__py3-none-any.whl

aitoolkit/__init__.py

@@ -0,0 +1,66 @@
+"""aitoolkit — centralized AI clients for self-hosted OpenAI-compatible services.
+
+Core capabilities (LLM, embeddings, STT, TTS) are always available. RAG and the
+LangChain bridge live behind extras and are imported from their own subpackages
+(``aitoolkit.rag``, ``aitoolkit.integrations.langchain``) so that importing the
+top-level package never forces an optional dependency.
+"""
+
+from __future__ import annotations
+
+from aitoolkit.config import AIToolkitSettings, configure, get_settings
+from aitoolkit.embeddings import EmbeddingsClient, get_embeddings_client
+from aitoolkit.exceptions import (
+    AIToolkitError,
+    ConfigurationError,
+    EmbeddingsError,
+    LLMError,
+    STTError,
+    TTSError,
+    VectorStoreError,
+)
+from aitoolkit.llm import LLMClient, get_llm_client
+from aitoolkit.stt import STTClient, get_stt_client
+from aitoolkit.tts import TTSClient, concat_wav, get_tts_client
+from aitoolkit.types import (
+    ChatMessage,
+    DialogueTurn,
+    RetrievedChunk,
+    TranscriptionResult,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "__version__",
+    # config
+    "AIToolkitSettings",
+    "configure",
+    "get_settings",
+    # llm
+    "LLMClient",
+    "get_llm_client",
+    # embeddings
+    "EmbeddingsClient",
+    "get_embeddings_client",
+    # stt
+    "STTClient",
+    "get_stt_client",
+    # tts
+    "TTSClient",
+    "get_tts_client",
+    "concat_wav",
+    # types
+    "ChatMessage",
+    "DialogueTurn",
+    "RetrievedChunk",
+    "TranscriptionResult",
+    # exceptions
+    "AIToolkitError",
+    "ConfigurationError",
+    "LLMError",
+    "EmbeddingsError",
+    "STTError",
+    "TTSError",
+    "VectorStoreError",
+]

aitoolkit/config.py

@@ -0,0 +1,107 @@
+"""Central configuration for aitoolkit.
+
+All settings are read from ``AITOOLKIT_*`` environment variables (or a ``.env``
+file) but every client also accepts explicit overrides, so the package can be
+used with zero environment configuration.
+
+Defaults intentionally point at ``localhost`` — they are NOT specific to any one
+deployment. Production endpoints are supplied by the consuming application via
+environment variables (see ``.env.example``).
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Optional
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+# A placeholder key for OpenAI-compatible servers that perform no app-layer auth
+# (our GPU services are firewalled, not key-gated). The openai SDK requires a
+# non-empty key, so we provide one.
+_NO_AUTH = "no-auth"
+
+
+class AIToolkitSettings(BaseSettings):
+    """Runtime configuration for every aitoolkit capability."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="AITOOLKIT_",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    # --- LLM (vLLM, OpenAI-compatible) ---
+    llm_base_url: str = Field(default="http://localhost:8000/v1")
+    llm_api_key: str = Field(default=_NO_AUTH)
+    # No default model id — set AITOOLKIT_LLM_MODEL to your served model.
+    llm_model: str = Field(default="")
+    llm_temperature: float = Field(default=0.2)
+    llm_timeout: float = Field(default=60.0)
+    llm_max_retries: int = Field(default=2)
+
+    # --- Embeddings (TEI, OpenAI-compatible) ---
+    embeddings_base_url: str = Field(default="http://localhost:8001/v1")
+    embeddings_api_key: str = Field(default=_NO_AUTH)
+    # No default model id — set AITOOLKIT_EMBEDDINGS_MODEL to your served model.
+    embeddings_model: str = Field(default="")
+    # TEI accepts modest batches; keep conservative and configurable.
+    embeddings_batch_size: int = Field(default=32)
+    embeddings_timeout: float = Field(default=60.0)
+
+    # --- Speech-to-Text (faster-whisper, OpenAI-compatible) ---
+    stt_base_url: str = Field(default="http://localhost:8003/v1")
+    stt_api_key: str = Field(default=_NO_AUTH)
+    stt_model: str = Field(default="whisper-1")
+    stt_language: Optional[str] = Field(default=None)
+    stt_timeout: float = Field(default=120.0)
+
+    # --- Text-to-Speech (custom /api/tts) ---
+    tts_base_url: str = Field(default="http://localhost:8002")
+    tts_default_voice: Optional[str] = Field(default=None)
+    tts_timeout: float = Field(default=120.0)
+
+    # --- Vector store (Qdrant) ---
+    qdrant_url: str = Field(default="http://localhost:6333")
+    qdrant_collection: str = Field(default="documents")
+    # Optional fixed vector size. When None, it is detected from the embedding model.
+    qdrant_vector_size: Optional[int] = Field(default=None)
+    # The qdrant-client refuses to talk to a server whose minor version differs by
+    # more than one, emitting a UserWarning. Self-hosted servers often lag the
+    # client; set False to silence the check when the API surface we use is stable.
+    qdrant_check_compatibility: bool = Field(default=True)
+
+    # --- Retriever cache (Redis, optional) ---
+    redis_url: Optional[str] = Field(default=None)
+    cache_ttl: int = Field(default=3600)
+
+
+_override: Optional[AIToolkitSettings] = None
+
+
+@lru_cache(maxsize=1)
+def _env_settings() -> AIToolkitSettings:
+    """Settings built from AITOOLKIT_* env vars (cached fallback)."""
+    return AIToolkitSettings()
+
+
+def get_settings() -> AIToolkitSettings:
+    """Return the process-wide settings.
+
+    Prefers an instance installed via :func:`configure`; otherwise builds one
+    from ``AITOOLKIT_*`` environment variables.
+    """
+    return _override if _override is not None else _env_settings()
+
+
+def configure(settings: AIToolkitSettings) -> None:
+    """Install an explicit settings object as the process-wide singleton.
+
+    Lets a consuming application own configuration directly instead of relying
+    on environment variables and this package's generic defaults. Call once at
+    startup, before any client is created.
+    """
+    global _override
+    _override = settings

aitoolkit/embeddings/__init__.py

@@ -0,0 +1,5 @@
+"""Embeddings capability — OpenAI-compatible embeddings (e.g. TEI)."""
+
+from aitoolkit.embeddings.client import EmbeddingsClient, get_embeddings_client
+
+__all__ = ["EmbeddingsClient", "get_embeddings_client"]

aitoolkit/embeddings/client.py

@@ -0,0 +1,133 @@
+"""Embeddings client backed by an OpenAI-compatible embeddings server (TEI).
+
+Async-first, with synchronous convenience wrappers so it can also satisfy the
+LangChain ``Embeddings`` interface. The embedding dimension is detected at
+runtime — never hardcoded — so swapping the underlying model requires no code
+change here (only a re-index in the vector store).
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import List, Optional
+
+from loguru import logger
+from openai import AsyncOpenAI, OpenAI
+
+from aitoolkit.config import AIToolkitSettings, get_settings
+from aitoolkit.exceptions import EmbeddingsError
+
+
+class EmbeddingsClient:
+    """Create embeddings via an OpenAI-compatible ``/v1/embeddings`` endpoint."""
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        timeout: Optional[float] = None,
+        settings: Optional[AIToolkitSettings] = None,
+    ) -> None:
+        settings = settings or get_settings()
+        self.model = model or settings.embeddings_model
+        self.batch_size = batch_size or settings.embeddings_batch_size
+        self._base_url = base_url or settings.embeddings_base_url
+        self._api_key = api_key or settings.embeddings_api_key
+        self._timeout = timeout if timeout is not None else settings.embeddings_timeout
+
+        self._aclient = AsyncOpenAI(
+            base_url=self._base_url, api_key=self._api_key, timeout=self._timeout
+        )
+        self._sclient: Optional[OpenAI] = None
+        self._dimension: Optional[int] = None
+        logger.info(
+            f"EmbeddingsClient ready (model={self.model}, base_url={self._base_url})"
+        )
+
+    @property
+    def sync_client(self) -> OpenAI:
+        if self._sclient is None:
+            self._sclient = OpenAI(
+                base_url=self._base_url, api_key=self._api_key, timeout=self._timeout
+            )
+        return self._sclient
+
+    @property
+    def dimension(self) -> Optional[int]:
+        """The embedding dimension, known after the first call."""
+        return self._dimension
+
+    # --------------------------------------------------------------- helpers
+    @staticmethod
+    def _batched(items: List[str], size: int):
+        for i in range(0, len(items), size):
+            yield items[i : i + size]
+
+    def _record_dim(self, vectors: List[List[float]]) -> None:
+        if vectors and self._dimension is None:
+            self._dimension = len(vectors[0])
+            logger.debug(f"Detected embedding dimension: {self._dimension}")
+
+    # ----------------------------------------------------------------- async
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed many documents, batching to respect server limits."""
+        if not texts:
+            return []
+        out: List[List[float]] = []
+        for batch in self._batched(texts, self.batch_size):
+            try:
+                resp = await self._aclient.embeddings.create(
+                    model=self.model, input=batch
+                )
+            except Exception as exc:  # noqa: BLE001
+                raise EmbeddingsError(f"embedding request failed: {exc}") from exc
+            # Preserve input order via the returned index.
+            ordered = sorted(resp.data, key=lambda d: d.index)
+            out.extend([d.embedding for d in ordered])
+        self._record_dim(out)
+        return out
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Embed a single query string."""
+        vectors = await self.aembed_documents([text])
+        if not vectors:
+            raise EmbeddingsError("empty embedding response for query")
+        return vectors[0]
+
+    # ------------------------------------------------------------------ sync
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Synchronous document embedding (LangChain ``Embeddings`` interface)."""
+        if not texts:
+            return []
+        out: List[List[float]] = []
+        for batch in self._batched(texts, self.batch_size):
+            try:
+                resp = self.sync_client.embeddings.create(
+                    model=self.model, input=batch
+                )
+            except Exception as exc:  # noqa: BLE001
+                raise EmbeddingsError(f"embedding request failed: {exc}") from exc
+            ordered = sorted(resp.data, key=lambda d: d.index)
+            out.extend([d.embedding for d in ordered])
+        self._record_dim(out)
+        return out
+
+    def embed_query(self, text: str) -> List[float]:
+        """Synchronous single-query embedding."""
+        vectors = self.embed_documents([text])
+        if not vectors:
+            raise EmbeddingsError("empty embedding response for query")
+        return vectors[0]
+
+    async def aclose(self) -> None:
+        await self._aclient.close()
+        if self._sclient is not None:
+            self._sclient.close()
+
+
+@lru_cache(maxsize=1)
+def get_embeddings_client() -> EmbeddingsClient:
+    """Return the process-wide embeddings client singleton."""
+    return EmbeddingsClient()

aitoolkit/exceptions.py

@@ -0,0 +1,35 @@
+"""Exception hierarchy for aitoolkit.
+
+A single base class lets callers catch every toolkit-originated error with one
+``except AIToolkitError``, while specific subclasses allow fine-grained handling.
+"""
+
+from __future__ import annotations
+
+
+class AIToolkitError(Exception):
+    """Base class for all aitoolkit errors."""
+
+
+class ConfigurationError(AIToolkitError):
+    """Raised when configuration is missing or invalid."""
+
+
+class LLMError(AIToolkitError):
+    """Raised when an LLM request fails or returns an unusable response."""
+
+
+class EmbeddingsError(AIToolkitError):
+    """Raised when an embeddings request fails."""
+
+
+class STTError(AIToolkitError):
+    """Raised when speech-to-text transcription fails."""
+
+
+class TTSError(AIToolkitError):
+    """Raised when text-to-speech synthesis fails."""
+
+
+class VectorStoreError(AIToolkitError):
+    """Raised when a vector-store operation fails."""

aitoolkit/integrations/__init__.py

@@ -0,0 +1 @@
+"""Optional third-party integrations (each gated behind its own extra)."""

aitoolkit/integrations/langchain.py

@@ -0,0 +1,69 @@
+"""LangChain bridge (optional extra ``aitoolkit[langchain]``).
+
+Provides:
+
+* :func:`to_chat_model` — a LangChain ``BaseChatModel`` pointed at the toolkit's
+  LLM endpoint, so existing LangGraph graphs keep working unchanged.
+* :class:`LangChainEmbeddings` — wraps :class:`EmbeddingsClient` as a LangChain
+  ``Embeddings`` so vector stores / chains can consume toolkit embeddings.
+
+Only this module imports LangChain; the toolkit core stays LangChain-free.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+try:  # pragma: no cover - import guard
+    from langchain_core.embeddings import Embeddings
+    from langchain_openai import ChatOpenAI
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "LangChain integration requires extra deps. "
+        "Install with: pip install 'aitoolkit[langchain]'"
+    ) from exc
+
+from aitoolkit.config import get_settings
+from aitoolkit.embeddings import EmbeddingsClient, get_embeddings_client
+
+
+def to_chat_model(
+    model: Optional[str] = None,
+    temperature: Optional[float] = None,
+    streaming: bool = False,
+    **kwargs,
+) -> ChatOpenAI:
+    """Return a LangChain ``ChatOpenAI`` bound to the toolkit's LLM endpoint."""
+    settings = get_settings()
+    return ChatOpenAI(
+        base_url=settings.llm_base_url,
+        api_key=settings.llm_api_key,
+        model=model or settings.llm_model,
+        temperature=temperature if temperature is not None else settings.llm_temperature,
+        streaming=streaming,
+        timeout=settings.llm_timeout,
+        max_retries=settings.llm_max_retries,
+        **kwargs,
+    )
+
+
+class LangChainEmbeddings(Embeddings):
+    """Adapt :class:`EmbeddingsClient` to the LangChain ``Embeddings`` interface."""
+
+    def __init__(self, client: Optional[EmbeddingsClient] = None) -> None:
+        self._client = client or get_embeddings_client()
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        return self._client.embed_documents(texts)
+
+    def embed_query(self, text: str) -> List[float]:
+        return self._client.embed_query(text)
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        return await self._client.aembed_documents(texts)
+
+    async def aembed_query(self, text: str) -> List[float]:
+        return await self._client.aembed_query(text)
+
+
+__all__ = ["to_chat_model", "LangChainEmbeddings"]

aitoolkit/llm/__init__.py

@@ -0,0 +1,5 @@
+"""LLM capability — OpenAI-compatible chat, streaming, and structured output."""
+
+from aitoolkit.llm.client import LLMClient, get_llm_client
+
+__all__ = ["LLMClient", "get_llm_client"]

aitoolkit/llm/client.py

@@ -0,0 +1,230 @@
+"""LLM client backed by any OpenAI-compatible server (e.g. vLLM).
+
+Provides three primitives, all provider-agnostic:
+
+* :meth:`LLMClient.chat`             — single completion, returns text
+* :meth:`LLMClient.stream`           — async token stream
+* :meth:`LLMClient.chat_structured`  — completion validated into a Pydantic model
+
+The client returns plain strings / Pydantic models and never leaks the underlying
+``openai`` SDK types to callers.
+"""
+
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from typing import AsyncIterator, List, Optional, Type, TypeVar
+
+from loguru import logger
+from openai import AsyncOpenAI, OpenAI
+from pydantic import BaseModel, ValidationError
+
+from aitoolkit.config import AIToolkitSettings, get_settings
+from aitoolkit.exceptions import LLMError
+from aitoolkit.types import ChatMessage, as_messages
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class LLMClient:
+    """A thin, stable wrapper over an OpenAI-compatible chat endpoint."""
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        temperature: Optional[float] = None,
+        timeout: Optional[float] = None,
+        max_retries: Optional[int] = None,
+        settings: Optional[AIToolkitSettings] = None,
+    ) -> None:
+        settings = settings or get_settings()
+        self.model = model or settings.llm_model
+        self.default_temperature = (
+            temperature if temperature is not None else settings.llm_temperature
+        )
+        self._base_url = base_url or settings.llm_base_url
+        self._api_key = api_key or settings.llm_api_key
+        self._timeout = timeout if timeout is not None else settings.llm_timeout
+        self._max_retries = (
+            max_retries if max_retries is not None else settings.llm_max_retries
+        )
+
+        self._aclient = AsyncOpenAI(
+            base_url=self._base_url,
+            api_key=self._api_key,
+            timeout=self._timeout,
+            max_retries=self._max_retries,
+        )
+        self._sclient: Optional[OpenAI] = None  # created lazily for sync calls
+        logger.info(f"LLMClient ready (model={self.model}, base_url={self._base_url})")
+
+    @property
+    def sync_client(self) -> OpenAI:
+        """Lazily-created synchronous client (for non-async call sites)."""
+        if self._sclient is None:
+            self._sclient = OpenAI(
+                base_url=self._base_url,
+                api_key=self._api_key,
+                timeout=self._timeout,
+                max_retries=self._max_retries,
+            )
+        return self._sclient
+
+    # ------------------------------------------------------------------ chat
+    async def chat(
+        self,
+        prompt: Optional[str] = None,
+        *,
+        system: Optional[str] = None,
+        messages: Optional[List[ChatMessage]] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> str:
+        """Return a single completion as text."""
+        msgs = as_messages(prompt, system=system, messages=messages)
+        try:
+            resp = await self._aclient.chat.completions.create(
+                model=self.model,
+                messages=msgs,  # type: ignore[arg-type]
+                temperature=self._temp(temperature),
+                max_tokens=max_tokens,
+                **kwargs,
+            )
+        except Exception as exc:  # noqa: BLE001 - surface as toolkit error
+            raise LLMError(f"chat completion failed: {exc}") from exc
+
+        return resp.choices[0].message.content or ""
+
+    def chat_sync(
+        self,
+        prompt: Optional[str] = None,
+        *,
+        system: Optional[str] = None,
+        messages: Optional[List[ChatMessage]] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> str:
+        """Synchronous counterpart of :meth:`chat`."""
+        msgs = as_messages(prompt, system=system, messages=messages)
+        try:
+            resp = self.sync_client.chat.completions.create(
+                model=self.model,
+                messages=msgs,  # type: ignore[arg-type]
+                temperature=self._temp(temperature),
+                max_tokens=max_tokens,
+                **kwargs,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise LLMError(f"chat completion failed: {exc}") from exc
+        return resp.choices[0].message.content or ""
+
+    # ---------------------------------------------------------------- stream
+    async def stream(
+        self,
+        prompt: Optional[str] = None,
+        *,
+        system: Optional[str] = None,
+        messages: Optional[List[ChatMessage]] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> AsyncIterator[str]:
+        """Yield completion text deltas as they arrive."""
+        msgs = as_messages(prompt, system=system, messages=messages)
+        try:
+            stream = await self._aclient.chat.completions.create(
+                model=self.model,
+                messages=msgs,  # type: ignore[arg-type]
+                temperature=self._temp(temperature),
+                max_tokens=max_tokens,
+                stream=True,
+                **kwargs,
+            )
+            async for chunk in stream:
+                if not chunk.choices:
+                    continue
+                delta = chunk.choices[0].delta.content
+                if delta:
+                    yield delta
+        except Exception as exc:  # noqa: BLE001
+            raise LLMError(f"streaming completion failed: {exc}") from exc
+
+    # ------------------------------------------------------------ structured
+    async def chat_structured(
+        self,
+        response_model: Type[T],
+        prompt: Optional[str] = None,
+        *,
+        system: Optional[str] = None,
+        messages: Optional[List[ChatMessage]] = None,
+        temperature: Optional[float] = None,
+        strict: bool = False,
+        **kwargs,
+    ) -> T:
+        """Return a completion validated into ``response_model``.
+
+        Uses the OpenAI ``response_format`` json_schema mechanism, which vLLM
+        implements via guided decoding. The raw JSON is validated with Pydantic,
+        so a malformed response raises :class:`LLMError` rather than passing
+        silently.
+        """
+        msgs = as_messages(prompt, system=system, messages=messages)
+        schema = response_model.model_json_schema()
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": response_model.__name__,
+                "schema": schema,
+                "strict": strict,
+            },
+        }
+        try:
+            resp = await self._aclient.chat.completions.create(
+                model=self.model,
+                messages=msgs,  # type: ignore[arg-type]
+                temperature=self._temp(temperature),
+                response_format=response_format,  # type: ignore[arg-type]
+                **kwargs,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise LLMError(f"structured completion request failed: {exc}") from exc
+
+        content = resp.choices[0].message.content or ""
+        try:
+            return response_model.model_validate_json(content)
+        except (ValidationError, json.JSONDecodeError) as exc:
+            raise LLMError(
+                f"structured output did not match {response_model.__name__}: {exc}\n"
+                f"raw content: {content[:500]}"
+            ) from exc
+
+    # ----------------------------------------------------------------- utils
+    def _temp(self, temperature: Optional[float]) -> float:
+        return temperature if temperature is not None else self.default_temperature
+
+    async def aclose(self) -> None:
+        await self._aclient.close()
+        if self._sclient is not None:
+            self._sclient.close()
+
+
+@lru_cache(maxsize=8)
+def _cached_client(model: Optional[str], temperature: Optional[float]) -> LLMClient:
+    return LLMClient(model=model, temperature=temperature)
+
+
+def get_llm_client(
+    model: Optional[str] = None,
+    temperature: Optional[float] = None,
+) -> LLMClient:
+    """Return a cached :class:`LLMClient` for the given model/temperature.
+
+    Caching mirrors the previous ``llm.py`` behaviour and avoids re-creating
+    HTTP clients on every request.
+    """
+    return _cached_client(model, temperature)

aitoolkit/rag/__init__.py

@@ -0,0 +1,25 @@
+"""RAG capability (optional extra ``aitoolkit[rag]``).
+
+Imports require ``qdrant-client``. Install with ``pip install 'aitoolkit[rag]'``.
+"""
+
+try:  # pragma: no cover - import guard
+    import qdrant_client  # noqa: F401
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "The RAG module requires qdrant-client. Install with: pip install 'aitoolkit[rag]'"
+    ) from exc
+
+from aitoolkit.rag.agent import RAGAgent, get_rag_agent
+from aitoolkit.rag.query_expansion import QueryExpander, get_query_expander
+from aitoolkit.rag.retriever import RAGRetriever
+from aitoolkit.rag.vector_store import UnifiedVectorStore
+
+__all__ = [
+    "RAGAgent",
+    "get_rag_agent",
+    "RAGRetriever",
+    "UnifiedVectorStore",
+    "QueryExpander",
+    "get_query_expander",
+]