swiftpy-framework 0.1.0__py3-none-any.whl

swiftpy/ai/rag/retriever.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import copy
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from swiftpy.ai.embeddings import EmbeddingsClient
+from swiftpy.ai.vectorstore.base import SearchResult, VectorRecord, VectorStore
+from swiftpy.core.container import Container, get_container
+
+
+def _estimate_text_tokens(text: str) -> int:
+    """Approximate token count cleanly (~4 chars per token)."""
+    if not text:
+        return 0
+    return max(1, (len(text) + 3) // 4)
+
+
+def _extract_document_text(doc: SearchResult | str | dict[str, Any], index: int) -> str:
+    """Extract plain text content from a document or raise TypeError."""
+    if isinstance(doc, SearchResult):
+        return doc.content
+    if isinstance(doc, str):
+        return doc
+    if isinstance(doc, dict):
+        raw = doc.get("content", "")
+        return "" if raw is None else str(raw)
+    raise TypeError(f"Unsupported document type at index {index}: {type(doc).__name__}")
+
+
+def chunk_text(
+    text: str,
+    chunk_size: int = 1000,
+    overlap: int = 200,
+) -> list[str]:
+    """
+    Split text into overlapping chunks bounded by `chunk_size` characters.
+
+    Prefers splitting on whitespace boundaries when available while strictly
+    enforcing ``len(chunk) <= chunk_size`` and monotonic forward progress.
+    """
+    if not isinstance(text, str):
+        raise TypeError(f"text must be a str, got {type(text).__name__}")
+    if isinstance(chunk_size, bool) or not isinstance(chunk_size, int) or chunk_size <= 0:
+        raise ValueError(f"chunk_size must be a positive integer, got {chunk_size!r}")
+    if isinstance(overlap, bool) or not isinstance(overlap, int) or overlap < 0:
+        raise ValueError(f"overlap must be a non-negative integer, got {overlap!r}")
+    if overlap >= chunk_size:
+        raise ValueError(f"overlap ({overlap}) must be smaller than chunk_size ({chunk_size})")
+
+    if not text.strip():
+        return []
+
+    chunks: list[str] = []
+    start = 0
+    length = len(text)
+
+    while start < length:
+        while start < length and text[start].isspace():
+            start += 1
+        if start >= length:
+            break
+
+        end = min(start + chunk_size, length)
+        if end < length and not text[end].isspace():
+            split_at = text.rfind(" ", start, end)
+            if split_at == -1:
+                split_at = text.rfind("\n", start, end)
+            if split_at > start:
+                end = split_at
+
+        chunk = text[start:end].strip()
+        if chunk:
+            chunks.append(chunk)
+
+        if end >= length:
+            break
+
+        next_start = end - overlap
+        if next_start <= start:
+            next_start = start + 1
+        start = next_start
+
+    return chunks
+
+
+def truncate_to_token_budget[DocT: SearchResult | str | dict[str, Any]](
+    documents: Sequence[DocT],
+    max_tokens: int,
+    token_estimator: Callable[[str], int] | None = None,
+) -> list[DocT]:
+    """
+    Retain an ordered prefix of documents whose cumulative token estimate stays
+    within `max_tokens`.
+    """
+    if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0:
+        raise ValueError(f"max_tokens must be a positive integer, got {max_tokens!r}")
+    if token_estimator is not None and not callable(token_estimator):
+        raise TypeError("token_estimator must be callable when provided")
+
+    kept: list[DocT] = []
+    used_tokens = 0
+    for idx, doc in enumerate(documents):
+        content = _extract_document_text(doc, idx)
+        est = token_estimator(content) if token_estimator is not None else _estimate_text_tokens(content)
+        if est < 0:
+            raise ValueError(f"Token estimate must be non-negative, got {est}")
+        if used_tokens + est > max_tokens:
+            break
+        kept.append(doc)
+        used_tokens += est
+    return kept
+
+
+class Retriever:
+    """
+    Asynchronous document retriever for vector similarity search.
+
+    Embeds incoming text queries asynchronously and queries the configured
+    vector store to return ranked search result chunks.
+    """
+
+    def __init__(
+        self,
+        vectorstore: VectorStore,
+        embeddings: EmbeddingsClient | None = None,
+        top_k: int = 5,
+        filter: dict[str, Any] | None = None,
+        min_score: float | None = None,
+        model: str | None = None,
+        provider: str | None = None,
+    ) -> None:
+        if top_k <= 0:
+            raise ValueError(f"top_k must be greater than 0, got {top_k}")
+
+        self._vectorstore = vectorstore
+        if embeddings is None:
+            try:
+                container = get_container()
+            except RuntimeError:
+                container = Container()
+
+            if container.has(EmbeddingsClient):
+                self._embeddings = container.resolve(EmbeddingsClient)
+            else:
+                self._embeddings = EmbeddingsClient()
+        else:
+            self._embeddings = embeddings
+
+        self._top_k = top_k
+        self._filter = filter
+        self._min_score = min_score
+        self._model = model
+        self._provider = provider
+
+    @property
+    def vectorstore(self) -> VectorStore:
+        """The underlying vector store instance."""
+        return self._vectorstore
+
+    @property
+    def embeddings(self) -> EmbeddingsClient:
+        """The underlying embeddings client instance."""
+        return self._embeddings
+
+    @property
+    def top_k(self) -> int:
+        """Default top_k results to retrieve."""
+        return self._top_k
+
+    @property
+    def filter(self) -> dict[str, Any] | None:
+        """Default metadata filter."""
+        return self._filter
+
+    @property
+    def min_score(self) -> float | None:
+        """Optional minimum score threshold."""
+        return self._min_score
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int | None = None,
+        filter: dict[str, Any] | None = None,
+        min_score: float | None = None,
+    ) -> list[SearchResult]:
+        """
+        Embed the query asynchronously, search the vector store, and return ranked chunks.
+        """
+        if not isinstance(query, str) or not query.strip():
+            raise ValueError("Query must be a non-empty string.")
+
+        limit = top_k if top_k is not None else self._top_k
+        if limit <= 0:
+            raise ValueError(f"top_k must be greater than 0, got {limit}")
+
+        active_filter = filter if filter is not None else self._filter
+        cutoff = min_score if min_score is not None else self._min_score
+
+        query_vector = await self._embeddings.embed(
+            query,
+            model=self._model,
+            provider=self._provider,
+        )
+
+        results = await self._vectorstore.search(
+            vector=query_vector,
+            limit=limit,
+            filter=active_filter,
+        )
+
+        if cutoff is not None:
+            results = [r for r in results if r.score >= cutoff]
+
+        return results
+
+    async def add_documents(
+        self,
+        documents: Sequence[str | dict[str, Any] | VectorRecord],
+        metadatas: Sequence[dict[str, Any]] | None = None,
+        ids: Sequence[str] | None = None,
+        chunk_size: int = 100,
+        concurrency: int = 5,
+    ) -> list[str]:
+        """
+        Convenience method to embed and ingest documents into the underlying vector store.
+
+        Accepts a sequence of strings (texts), dictionaries, or VectorRecords.
+        Returns the list of stored document identifiers.
+        """
+        if not documents:
+            return []
+
+        if ids is not None and len(ids) != len(documents):
+            raise ValueError(
+                f"Length of ids ({len(ids)}) does not match length of documents ({len(documents)})"
+            )
+
+        if metadatas is not None and len(metadatas) != len(documents):
+            raise ValueError(
+                f"Length of metadatas ({len(metadatas)}) does not match length of documents ({len(documents)})"
+            )
+
+        # 1. Validate & partition EVERY document up front (supporting mixed types)
+        records_by_idx: dict[int, VectorRecord] = {}
+        embed_indices: list[int] = []
+        texts_to_embed: list[str] = []
+        embed_ids: list[str] = []
+        embed_metas: list[dict[str, Any]] = []
+
+        for idx, doc in enumerate(documents):
+            custom_id = str(ids[idx]) if ids is not None else f"doc_{idx}"
+            custom_meta = copy.deepcopy(dict(metadatas[idx])) if metadatas is not None else {}
+
+            if isinstance(doc, VectorRecord):
+                merged_meta = (
+                    {**custom_meta, **copy.deepcopy(doc.metadata)}
+                    if custom_meta
+                    else copy.deepcopy(doc.metadata)
+                )
+                records_by_idx[idx] = VectorRecord(
+                    id=str(ids[idx]) if ids is not None else doc.id,
+                    vector=list(doc.vector),
+                    metadata=merged_meta,
+                    content=doc.content,
+                )
+            elif isinstance(doc, str):
+                embed_indices.append(idx)
+                texts_to_embed.append(doc)
+                embed_ids.append(custom_id)
+                embed_metas.append(custom_meta)
+            elif isinstance(doc, dict):
+                raw_id = doc.get("id")
+                item_id = str(raw_id) if raw_id is not None and str(raw_id) != "" else custom_id
+                raw_content = doc.get("content", "")
+                item_content = "" if raw_content is None else str(raw_content)
+                doc_meta = doc.get("metadata")
+                if doc_meta is not None and not isinstance(doc_meta, dict):
+                    raise TypeError(
+                        f"Document metadata at index {idx} must be a dict, got {type(doc_meta).__name__}"
+                    )
+                item_meta = (
+                    {**custom_meta, **copy.deepcopy(doc_meta)} if doc_meta is not None else custom_meta
+                )
+                embed_indices.append(idx)
+                texts_to_embed.append(item_content)
+                embed_ids.append(item_id)
+                embed_metas.append(item_meta)
+            else:
+                raise TypeError(
+                    f"Unsupported document type at index {idx}: {type(doc).__name__}"
+                )
+
+        # 2. Batch-embed only text/dict entries asynchronously
+        if texts_to_embed:
+            vectors = await self._embeddings.embed_batch(
+                texts=texts_to_embed,
+                model=self._model,
+                provider=self._provider,
+                chunk_size=chunk_size,
+                concurrency=concurrency,
+            )
+            for pos, orig_idx in enumerate(embed_indices):
+                records_by_idx[orig_idx] = VectorRecord(
+                    id=embed_ids[pos],
+                    vector=vectors[pos],
+                    metadata=embed_metas[pos],
+                    content=texts_to_embed[pos],
+                )
+
+        # 3. Preserve original document ordering and upsert atomically
+        records_to_upsert = [records_by_idx[i] for i in range(len(documents))]
+        await self._vectorstore.upsert_batch(records_to_upsert)
+        return [r.id for r in records_to_upsert]

swiftpy/ai/redis_store.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import json
+import time
+from collections import OrderedDict
+from collections.abc import Awaitable
+from typing import Protocol
+
+import redis.asyncio as aioredis
+from redis.exceptions import RedisError
+
+from swiftpy.ai.conversation import Conversation
+from swiftpy.ai.exceptions import AIConnectionError
+from swiftpy.core.config import Config
+from swiftpy.support.redis_config import resolve_redis_params
+
+
+class _ConversationRedisClient(Protocol):
+    def get(self, name: str, /) -> Awaitable[bytes | str | None]: ...
+    def set(
+        self, name: str, value: str, /, ex: int | None = None
+    ) -> Awaitable[bool | str | bytes | None]: ...
+    def delete(self, name: str, /) -> Awaitable[int]: ...
+    def exists(self, name: str, /) -> Awaitable[int]: ...
+    def aclose(self) -> Awaitable[None]: ...
+
+
+class ConversationStore:
+    """Abstract protocol for conversation persistence."""
+
+    async def save(
+        self,
+        session_id: str,
+        conversation: Conversation,
+        ttl: int | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    async def load(self, session_id: str) -> Conversation:
+        raise NotImplementedError
+
+    async def delete(self, session_id: str) -> bool:
+        raise NotImplementedError
+
+    async def exists(self, session_id: str) -> bool:
+        raise NotImplementedError
+
+
+class InMemoryConversationStore(ConversationStore):
+    """Bounded LRU + TTL-aware in-memory conversation store."""
+
+    def __init__(
+        self,
+        max_entries: int = 1000,
+        default_ttl: int | None = None,
+    ) -> None:
+        if max_entries <= 0:
+            raise ValueError(f"max_entries must be greater than 0, got {max_entries}")
+        self.max_entries = max_entries
+        self.default_ttl = default_ttl
+        self._data: OrderedDict[str, tuple[str, float | None]] = OrderedDict()
+
+    def _get_valid_payload(self, session_id: str) -> str | None:
+        entry = self._data.get(session_id)
+        if entry is None:
+            return None
+        raw, expires_at = entry
+        if expires_at is not None and time.monotonic() >= expires_at:
+            del self._data[session_id]
+            return None
+        self._data.move_to_end(session_id)
+        return raw
+
+    async def save(
+        self,
+        session_id: str,
+        conversation: Conversation,
+        ttl: int | None = None,
+    ) -> None:
+        effective_ttl = self.default_ttl if ttl is None else ttl
+        if effective_ttl is not None and effective_ttl < 0:
+            self._data.pop(session_id, None)
+            return
+
+        expires_at = (
+            time.monotonic() + effective_ttl
+            if effective_ttl is not None and effective_ttl > 0
+            else None
+        )
+        self._data[session_id] = (conversation.to_json(), expires_at)
+        self._data.move_to_end(session_id)
+        while len(self._data) > self.max_entries:
+            self._data.popitem(last=False)
+
+    async def load(self, session_id: str) -> Conversation:
+        raw = self._get_valid_payload(session_id)
+        if raw is None:
+            return Conversation()
+        try:
+            return Conversation.from_json(raw)
+        except (ValueError, TypeError) as exc:
+            raise ValueError(f"Corrupted in-memory conversation data: {exc}") from exc
+
+    async def delete(self, session_id: str) -> bool:
+        return self._data.pop(session_id, None) is not None
+
+    async def exists(self, session_id: str) -> bool:
+        return self._get_valid_payload(session_id) is not None
+
+
+class RedisConversationStore(ConversationStore):
+    """
+    Redis-backed conversation store serializing dialogue history keyed by session/user ID.
+    """
+
+    def __init__(
+        self,
+        redis_client: _ConversationRedisClient | None = None,
+        config: Config | None = None,
+        key_prefix: str = "ai:conversation",
+        default_ttl: int = 86400,
+    ) -> None:
+        self._client: _ConversationRedisClient | None = redis_client
+        self._config = config
+        self._key_prefix = key_prefix
+        self.default_ttl = default_ttl
+        self._owns_client = redis_client is None
+
+    def _format_key(self, session_id: str) -> str:
+        # Sanitize session_id to prevent injection or whitespace issues
+        clean_id = session_id.strip()
+        return f"{self._key_prefix}:{clean_id}"
+
+    async def get_client(self) -> _ConversationRedisClient:
+        """Lazily initialize Redis client if not explicitly provided."""
+        if self._client is None:
+            cfg = self._config or Config.get_instance()
+            params = resolve_redis_params(cfg, section="ai")
+            client: _ConversationRedisClient = aioredis.from_url(
+                params["url"], decode_responses=True
+            )
+            self._client = client
+            return client
+        return self._client
+
+    async def save(
+        self,
+        session_id: str,
+        conversation: Conversation,
+        ttl: int | None = None,
+    ) -> None:
+        """
+        Serialize and persist a conversation into Redis.
+
+        If ttl is None, self.default_ttl is used.
+        If ttl == 0, the key is saved with no expiration.
+        If ttl < 0, any existing key is removed immediately.
+        """
+        client = await self.get_client()
+        key = self._format_key(session_id)
+        effective_ttl = self.default_ttl if ttl is None else ttl
+        payload = conversation.to_json()
+
+        try:
+            if effective_ttl < 0:
+                await client.delete(key)
+            elif effective_ttl > 0:
+                await client.set(key, payload, ex=effective_ttl)
+            else:
+                await client.set(key, payload)
+        except RedisError as exc:
+            raise AIConnectionError(
+                f"Failed to save conversation to Redis: {exc}",
+                provider="redis",
+            ) from exc
+
+    async def load(self, session_id: str) -> Conversation:
+        """Retrieve and deserialize a conversation from Redis. Returns empty Conversation if missing."""
+        client = await self.get_client()
+        key = self._format_key(session_id)
+
+        try:
+            raw = await client.get(key)
+        except RedisError as exc:
+            raise AIConnectionError(
+                f"Failed to load conversation from Redis: {exc}",
+                provider="redis",
+            ) from exc
+
+        if raw is None or not str(raw).strip():
+            return Conversation()
+
+        try:
+            return Conversation.from_json(str(raw))
+        except (ValueError, TypeError, json.JSONDecodeError) as exc:
+            raise ValueError(f"Corrupted conversation data in Redis for key '{key}': {exc}") from exc
+
+    async def delete(self, session_id: str) -> bool:
+        """Remove a conversation from Redis."""
+        client = await self.get_client()
+        key = self._format_key(session_id)
+        try:
+            deleted_count = await client.delete(key)
+            return bool(deleted_count > 0)
+        except RedisError as exc:
+            raise AIConnectionError(
+                f"Failed to delete conversation from Redis: {exc}",
+                provider="redis",
+            ) from exc
+
+    async def exists(self, session_id: str) -> bool:
+        """Check whether a conversation exists in Redis."""
+        client = await self.get_client()
+        key = self._format_key(session_id)
+        try:
+            return bool(await client.exists(key))
+        except RedisError as exc:
+            raise AIConnectionError(
+                f"Failed to check conversation existence in Redis: {exc}",
+                provider="redis",
+            ) from exc
+
+    async def close(self) -> None:
+        """Close the Redis client connection if owned by this store."""
+        if self._owns_client and self._client is not None:
+            await self._client.aclose()
+            self._client = None

swiftpy/ai/streaming.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import codecs
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+from swiftpy.ai.exceptions import AIProviderError
+
+
+async def _iter_sse_lines(lines: AsyncIterator[str]) -> AsyncIterator[str]:
+    """Expand chunks containing embedded newlines into individual SSE lines."""
+    async for chunk in lines:
+        if not chunk:
+            yield ""
+            continue
+        for line in chunk.splitlines():
+            yield line
+
+
+def _parse_sse_event_payload(full_data: str, current_event: str) -> dict[str, Any]:
+    """Parse buffered SSE event data and raise AIProviderError on active provider errors."""
+    try:
+        parsed_json = json.loads(full_data)
+    except (json.JSONDecodeError, ValueError):
+        return {"event": current_event, "data": full_data}
+
+    if isinstance(parsed_json, dict):
+        err_obj = parsed_json.get("error")
+        if parsed_json.get("type") == "error" or (
+            err_obj is not None and err_obj is not False
+        ):
+            payload_err = err_obj if err_obj is not None else parsed_json
+            err_msg = (
+                payload_err.get("message")
+                if isinstance(payload_err, dict)
+                else str(payload_err)
+            )
+            raise AIProviderError(f"Stream error: {err_msg}")
+    return {"event": current_event, "data": parsed_json}
+
+
+async def parse_sse_events(
+    lines: AsyncIterator[str],
+) -> AsyncIterator[dict[str, Any]]:
+    """
+    Parse an asynchronous stream of Server-Sent Event (SSE) lines into event dictionaries.
+
+    Supports:
+        - CRLF (\\r\\n) and LF (\\n) line endings
+        - Multiple data: lines within a single event (joined by newline per SSE specification)
+        - Event type identifiers (event: <name>)
+        - SSE comments (: ping/keepalive)
+        - Provider error detection
+        - Clean [DONE] termination
+    """
+    current_event = "message"
+    data_buffer: list[str] = []
+
+    async for raw_line in _iter_sse_lines(lines):
+        line = raw_line.rstrip("\r\n")
+
+        if not line:
+            # Empty line indicates event boundary
+            if data_buffer:
+                full_data = "\n".join(data_buffer)
+                data_buffer.clear()
+                if full_data == "[DONE]":
+                    break
+                yield _parse_sse_event_payload(full_data, current_event)
+                current_event = "message"
+            continue
+
+        if line.startswith(":"):
+            # SSE comment line, ignore
+            continue
+
+        if line.startswith("event:"):
+            current_event = line[len("event:") :].strip()
+        elif line.startswith("data:"):
+            # Strip leading space after data: if present per spec
+            val = line[len("data:") :]
+            val = val.removeprefix(" ")
+            data_buffer.append(val)
+
+    # Flush remaining data if stream closed without trailing empty line
+    if data_buffer:
+        full_data = "\n".join(data_buffer)
+        if full_data != "[DONE]":
+            yield _parse_sse_event_payload(full_data, current_event)
+
+
+async def split_chunk_lines(
+    chunks: AsyncIterator[str | bytes],
+) -> AsyncIterator[str]:
+    """
+    Split arbitrary incoming network chunks into individual lines, buffering
+    partial lines and multi-byte UTF-8 sequences across chunk boundaries.
+    """
+    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
+    buffer = ""
+    async for chunk in chunks:
+        if isinstance(chunk, (bytes, bytearray)):
+            text = decoder.decode(chunk, final=False)
+        else:
+            text = chunk
+        buffer += text
+        while "\n" in buffer:
+            line, buffer = buffer.split("\n", 1)
+            yield line
+    tail = decoder.decode(b"", final=True)
+    if tail:
+        buffer += tail
+    if buffer:
+        yield buffer

swiftpy/ai/tools/__init__.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from swiftpy.ai.exceptions import AgentLoopError
+from swiftpy.ai.tools.base_tool import BaseTool, FunctionTool, tool
+from swiftpy.ai.tools.first_party import (
+    CalculatorTool,
+    DateTimeTool,
+    WebSearchTool,
+)
+from swiftpy.ai.tools.tool_loop import ToolLoop
+
+__all__ = [
+    "AgentLoopError",
+    "BaseTool",
+    "CalculatorTool",
+    "DateTimeTool",
+    "FunctionTool",
+    "ToolLoop",
+    "WebSearchTool",
+    "tool",
+]
+