swiftpy-framework 0.1.0__py3-none-any.whl

swiftpy/ai/providers/openai_provider.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+_OPENAI_HTTP_ERRORS: tuple[type[Exception], ...] = (httpx.TimeoutException, httpx.RequestError)
+
+from swiftpy.ai.exceptions import (
+    AIAuthenticationError,
+    AIConnectionError,
+    AIProviderError,
+    AIRateLimitError,
+)
+from swiftpy.ai.pool import AIHttpPool
+from swiftpy.ai.providers.base import (
+    LLMRequest,
+    LLMResponse,
+    ProviderInterface,
+    TokenUsage,
+    ToolCall,
+    extract_error_message,
+    parse_retry_after,
+    strip_internal_extra,
+)
+from swiftpy.ai.streaming import parse_sse_events
+
+
+class OpenAIProvider(ProviderInterface):
+    """
+    OpenAI provider adapter mapping unified requests to the /v1/chat/completions endpoint.
+    """
+
+    def __init__(
+        self,
+        config: dict[str, Any] | None = None,
+        http_pool: AIHttpPool | None = None,
+    ) -> None:
+        super().__init__(config=config, http_pool=http_pool)
+        self._api_key: str = str(self.config.get("api_key", ""))
+        self._default_model: str = str(
+            self.config.get("model") or self.config.get("default_model") or "gpt-4o"
+        )
+        self._base_url: str = str(
+            self.config.get("base_url", "https://api.openai.com/v1")
+        ).rstrip("/")
+        self._timeout: float = float(self.config.get("timeout", 60.0))
+
+    @property
+    def name(self) -> str:
+        return "openai"
+
+    def get_default_model(self) -> str:
+        return self._default_model
+
+    def _build_payload(self, request: LLMRequest, *, stream: bool = False) -> tuple[str, dict[str, Any]]:
+        model = request.model or self.get_default_model()
+        payload: dict[str, Any] = {
+            "model": model,
+            "messages": [m.to_dict() for m in request.messages],
+        }
+        if stream:
+            payload["stream"] = True
+        if request.temperature is not None:
+            payload["temperature"] = request.temperature
+        if request.max_tokens is not None:
+            payload["max_tokens"] = request.max_tokens
+        if request.top_p is not None:
+            payload["top_p"] = request.top_p
+        if request.stop is not None:
+            payload["stop"] = request.stop
+        if request.tools is not None:
+            payload["tools"] = request.tools
+        if request.extra:
+            payload.update(strip_internal_extra(request.extra))
+        return model, payload
+
+    def _raise_for_http_error(self, response: httpx.Response, body_text: str, data: Any) -> None:
+        if response.status_code in (401, 403):
+            err_msg = extract_error_message(data, body_text) or "Authentication failed"
+            raise AIAuthenticationError(
+                f"OpenAI authentication failed ({response.status_code}): {err_msg}",
+                provider=self.name,
+                status_code=response.status_code,
+                response_body=body_text,
+            )
+        if response.status_code == 429:
+            retry_after = parse_retry_after(response.headers.get("Retry-After"))
+            err_msg = extract_error_message(data, body_text) or "Rate limit reached"
+            raise AIRateLimitError(
+                f"OpenAI rate limit exceeded: {err_msg}",
+                provider=self.name,
+                status_code=response.status_code,
+                response_body=body_text,
+                retry_after=retry_after,
+            )
+        err_msg = extract_error_message(data, body_text) or f"HTTP {response.status_code}"
+        raise AIProviderError(
+            f"OpenAI API error ({response.status_code}): {err_msg}",
+            provider=self.name,
+            status_code=response.status_code,
+            response_body=body_text,
+        )
+
+    async def send(self, request: LLMRequest) -> LLMResponse:
+        url = f"{self._base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+        }
+        model, payload = self._build_payload(request, stream=False)
+        client = self._http_pool.get_client()
+
+        async def _attempt() -> tuple[httpx.Response, float, str, Any]:
+            start_time = time.perf_counter()
+            try:
+                response = await client.post(
+                    url,
+                    headers=headers,
+                    json=payload,
+                    timeout=self._timeout,
+                )
+            except _OPENAI_HTTP_ERRORS as exc:
+                sanitized_exc = self.sanitize_error(str(exc))
+                raise AIConnectionError(
+                    f"Failed to connect to OpenAI ({type(exc).__name__}): {sanitized_exc}",
+                    provider=self.name,
+                ) from exc
+
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            body_text = self.sanitize_error(response.text)
+            try:
+                data: Any = response.json()
+            except (json.JSONDecodeError, ValueError):
+                data = None
+
+            if response.status_code >= 400:
+                self._raise_for_http_error(response, body_text, data)
+            return response, latency_ms, body_text, data
+
+        response, latency_ms, body_text, data = await self._call_with_retries(_attempt)
+
+        if not isinstance(data, dict):
+            raise AIProviderError(
+                f"OpenAI returned non-JSON response ({response.status_code}): {body_text[:200]}",
+                provider=self.name,
+                status_code=response.status_code,
+                response_body=body_text,
+            )
+
+        choices = data.get("choices")
+        if not isinstance(choices, list) or len(choices) == 0:
+            raise AIProviderError(
+                "OpenAI returned a response with no choices",
+                provider=self.name,
+                status_code=response.status_code,
+                response_body=body_text,
+            )
+
+        choice = choices[0]
+        if not isinstance(choice, dict):
+            raise AIProviderError(
+                "OpenAI returned malformed choice object",
+                provider=self.name,
+                status_code=response.status_code,
+                response_body=body_text,
+            )
+
+        msg = choice.get("message") or {}
+        content = str(msg.get("content") or "") if isinstance(msg, dict) else ""
+        role = str(msg.get("role") or "assistant") if isinstance(msg, dict) else "assistant"
+        finish_reason = choice.get("finish_reason")
+
+        extracted_tool_calls: list[ToolCall] = []
+        raw_tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
+        if isinstance(raw_tool_calls, list):
+            for tc in raw_tool_calls:
+                if isinstance(tc, dict):
+                    tc_id = str(tc.get("id") or "")
+                    fn = tc.get("function") or {}
+                    fn_name = str(fn.get("name") or "") if isinstance(fn, dict) else ""
+                    args_raw = fn.get("arguments") if isinstance(fn, dict) else {}
+                    if isinstance(args_raw, str):
+                        try:
+                            args_parsed = json.loads(args_raw)
+                        except (json.JSONDecodeError, TypeError):
+                            args_parsed = {"raw": args_raw}
+                    elif isinstance(args_raw, dict):
+                        args_parsed = args_raw
+                    else:
+                        args_parsed = {}
+                    extracted_tool_calls.append(
+                        ToolCall(id=tc_id, name=fn_name, arguments=args_parsed)
+                    )
+
+        raw_usage = data.get("usage")
+        if isinstance(raw_usage, dict):
+            usage = TokenUsage(
+                prompt_tokens=int(raw_usage.get("prompt_tokens") or 0),
+                completion_tokens=int(raw_usage.get("completion_tokens") or 0),
+                total_tokens=int(raw_usage.get("total_tokens") or 0),
+            )
+        else:
+            usage = TokenUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
+
+        model_reported = str(data.get("model") or model)
+
+        return LLMResponse(
+            content=content,
+            role=role,
+            model=model_reported,
+            usage=usage,
+            raw=data,
+            latency_ms=round(latency_ms, 2),
+            finish_reason=str(finish_reason) if finish_reason else None,
+            tool_calls=extracted_tool_calls,
+        )
+
+    async def stream(self, request: LLMRequest) -> AsyncIterator[str]:
+        url = f"{self._base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream",
+        }
+        _, payload = self._build_payload(request, stream=True)
+        client = self._http_pool.get_client()
+
+        async def _open_stream() -> httpx.Response:
+            try:
+                req = client.build_request(
+                    "POST",
+                    url,
+                    headers=headers,
+                    json=payload,
+                    timeout=self._timeout,
+                )
+                resp = await client.send(req, stream=True)
+            except _OPENAI_HTTP_ERRORS as exc:
+                sanitized_exc = self.sanitize_error(str(exc))
+                raise AIConnectionError(
+                    f"Failed to connect to OpenAI stream ({type(exc).__name__}): {sanitized_exc}",
+                    provider=self.name,
+                ) from exc
+
+            if resp.status_code >= 400:
+                raw_bytes = await resp.aread()
+                body_text = self.sanitize_error(raw_bytes.decode("utf-8", errors="replace"))
+                await resp.aclose()
+                try:
+                    err_data: Any = json.loads(body_text)
+                except (json.JSONDecodeError, ValueError):
+                    err_data = None
+                self._raise_for_http_error(resp, body_text, err_data)
+            return resp
+
+        response = await self._call_with_retries(_open_stream)
+
+        try:
+            async for event in parse_sse_events(response.aiter_lines()):
+                data = event.get("data")
+                if isinstance(data, dict):
+                    choices = data.get("choices")
+                    if isinstance(choices, list) and choices:
+                        delta = choices[0].get("delta")
+                        if isinstance(delta, dict):
+                            content = delta.get("content")
+                            if content and isinstance(content, str):
+                                yield content
+        finally:
+            await response.aclose()
+

swiftpy/ai/rag/__init__.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from swiftpy.ai.rag.context_formatter import (
+    ContextFormatter,
+    NumberedListFormatter,
+    PlainConcatenationFormatter,
+    XMLBlocksFormatter,
+)
+from swiftpy.ai.rag.pipeline_builder import (
+    DEFAULT_PROMPT_TEMPLATE,
+    DEFAULT_SYSTEM_PROMPT,
+    RAGPipelineBuilder,
+)
+from swiftpy.ai.rag.reranker import (
+    BaseReranker,
+    CohereReranker,
+    CrossEncoderReranker,
+)
+from swiftpy.ai.rag.result import (
+    DEFAULT_MODEL_PRICING,
+    RAGResult,
+    calculate_token_cost,
+    register_model_pricing,
+)
+from swiftpy.ai.rag.retriever import Retriever, chunk_text, truncate_to_token_budget
+
+
+def rag() -> RAGPipelineBuilder:
+    """Convenience factory returning a fluent RAGPipelineBuilder."""
+    return RAGPipelineBuilder()
+
+
+__all__ = [
+    "DEFAULT_MODEL_PRICING",
+    "DEFAULT_PROMPT_TEMPLATE",
+    "DEFAULT_SYSTEM_PROMPT",
+    "BaseReranker",
+    "CohereReranker",
+    "ContextFormatter",
+    "CrossEncoderReranker",
+    "NumberedListFormatter",
+    "PlainConcatenationFormatter",
+    "RAGPipelineBuilder",
+    "RAGResult",
+    "Retriever",
+    "XMLBlocksFormatter",
+    "calculate_token_cost",
+    "chunk_text",
+    "rag",
+    "register_model_pricing",
+    "truncate_to_token_budget",
+]
+

swiftpy/ai/rag/context_formatter.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import html
+import re
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any
+
+from swiftpy.ai.vectorstore.base import SearchResult
+
+_CONTEXT_BOUNDARY_RE = re.compile(
+    r"-{3,}\s*(BEGIN|END)\s+UNTRUSTED\s+REFERENCE\s+CONTEXT\s*-{3,}",
+    re.IGNORECASE,
+)
+_NONCE_BOUNDARY_RE = re.compile(
+    r"<<<\s*(?:END\s+)?CONTEXT_[A-Za-z0-9_-]*>>>",
+    re.IGNORECASE,
+)
+_XML_TAG_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
+
+
+def escape_context_boundaries(text: str) -> str:
+    """Neutralize embedded context boundary delimiters in untrusted text."""
+    escaped = _CONTEXT_BOUNDARY_RE.sub(
+        lambda m: f"[escaped: {m.group(1).upper()} UNTRUSTED REFERENCE CONTEXT]",
+        text,
+    )
+    return _NONCE_BOUNDARY_RE.sub("[escaped: CONTEXT BOUNDARY]", escaped)
+
+
+def _extract_doc_content(
+    doc: SearchResult | str | dict[str, Any],
+    index: int,
+) -> tuple[str, str, dict[str, Any], float]:
+    """
+    Extract (content, id, metadata, score) from a document input.
+    """
+    if isinstance(doc, SearchResult):
+        return doc.content, doc.id, doc.metadata, doc.score
+    if isinstance(doc, str):
+        return doc, f"doc_{index}", {}, 0.0
+    if isinstance(doc, dict):
+        return (
+            str(doc.get("content", "")),
+            str(doc.get("id", f"doc_{index}")),
+            dict(doc.get("metadata") or {}),
+            float(doc.get("score", 0.0)),
+        )
+    raise TypeError(f"Unsupported document type at index {index}: {type(doc).__name__}")
+
+
+class ContextFormatter(ABC):
+    """
+    Abstract contract for context formatting strategies in RAG pipelines.
+    """
+
+    @abstractmethod
+    def format(
+        self,
+        documents: Sequence[SearchResult | str | dict[str, Any]],
+    ) -> str:
+        """
+        Format a sequence of retrieved chunks into a prompt-ready context string.
+        """
+
+    @classmethod
+    def create(cls, strategy: str | ContextFormatter) -> ContextFormatter:
+        """
+        Factory method resolving a formatter instance from a strategy name or instance.
+
+        Supported strategies:
+        - 'numbered', 'numbered_list' -> NumberedListFormatter
+        - 'xml', 'xml_blocks' -> XMLBlocksFormatter
+        - 'plain', 'concatenation', 'plain_concatenation' -> PlainConcatenationFormatter
+        """
+        if isinstance(strategy, ContextFormatter):
+            return strategy
+
+        norm = str(strategy).lower().strip().replace("-", "_")
+        if norm in ("numbered", "numbered_list", "number"):
+            return NumberedListFormatter()
+        if norm in ("xml", "xml_blocks", "xml_block"):
+            return XMLBlocksFormatter()
+        if norm in ("plain", "concatenation", "plain_concatenation", "concat"):
+            return PlainConcatenationFormatter()
+
+        raise ValueError(
+            f"Unknown context formatter strategy '{strategy}'. "
+            f"Supported strategies: 'numbered', 'xml', 'concatenation'."
+        )
+
+    @classmethod
+    def numbered(
+        cls,
+        separator: str = "\n\n",
+        include_metadata: bool = False,
+    ) -> NumberedListFormatter:
+        """Create a numbered list context formatter."""
+        return NumberedListFormatter(separator=separator, include_metadata=include_metadata)
+
+    @classmethod
+    def xml(
+        cls,
+        root_tag: str = "context",
+        document_tag: str = "document",
+        include_attributes: bool = True,
+        escape_content: bool = True,
+    ) -> XMLBlocksFormatter:
+        """Create an XML blocks context formatter."""
+        return XMLBlocksFormatter(
+            root_tag=root_tag,
+            document_tag=document_tag,
+            include_attributes=include_attributes,
+            escape_content=escape_content,
+        )
+
+    @classmethod
+    def concatenation(
+        cls,
+        separator: str = "\n\n",
+        prefix: str = "",
+        suffix: str = "",
+    ) -> PlainConcatenationFormatter:
+        """Create a plain concatenation context formatter."""
+        return PlainConcatenationFormatter(separator=separator, prefix=prefix, suffix=suffix)
+
+
+class NumberedListFormatter(ContextFormatter):
+    """
+    Formats retrieved documents as a numbered list:
+    [1] First document content
+    [2] Second document content
+    """
+
+    def __init__(
+        self,
+        separator: str = "\n\n",
+        include_metadata: bool = False,
+    ) -> None:
+        self.separator = separator
+        self.include_metadata = include_metadata
+
+    def format(
+        self,
+        documents: Sequence[SearchResult | str | dict[str, Any]],
+    ) -> str:
+        if not documents:
+            return ""
+
+        blocks: list[str] = []
+        for idx, doc in enumerate(documents, start=1):
+            content, doc_id, metadata, _ = _extract_doc_content(doc, idx)
+            safe_content = escape_context_boundaries(content)
+            header = f"[{idx}]"
+            if self.include_metadata:
+                meta_parts = [f"id: {escape_context_boundaries(str(doc_id))}"]
+                source = metadata.get("source") or metadata.get("title")
+                if source:
+                    meta_parts.append(
+                        f"source: {escape_context_boundaries(str(source))}"
+                    )
+                header += f" ({', '.join(meta_parts)})"
+
+            blocks.append(f"{header} {safe_content}")
+
+        return self.separator.join(blocks)
+
+
+class XMLBlocksFormatter(ContextFormatter):
+    """
+    Formats retrieved documents into structured XML blocks:
+    <context>
+      <document index="1" id="doc_1">
+        Content of document 1
+      </document>
+    </context>
+    """
+
+    def __init__(
+        self,
+        root_tag: str = "context",
+        document_tag: str = "document",
+        include_attributes: bool = True,
+        escape_content: bool = True,
+    ) -> None:
+        if not isinstance(root_tag, str) or not _XML_TAG_NAME_RE.match(root_tag):
+            raise ValueError(
+                f"Invalid XML root_tag {root_tag!r}; expected match for {_XML_TAG_NAME_RE.pattern}"
+            )
+        if not isinstance(document_tag, str) or not _XML_TAG_NAME_RE.match(document_tag):
+            raise ValueError(
+                f"Invalid XML document_tag {document_tag!r}; expected match for {_XML_TAG_NAME_RE.pattern}"
+            )
+        self.root_tag = root_tag
+        self.document_tag = document_tag
+        self.include_attributes = include_attributes
+        self.escape_content = escape_content
+
+    def format(
+        self,
+        documents: Sequence[SearchResult | str | dict[str, Any]],
+    ) -> str:
+        if not documents:
+            return f"<{self.root_tag}></{self.root_tag}>"
+
+        lines: list[str] = [f"<{self.root_tag}>"]
+        for idx, doc in enumerate(documents, start=1):
+            content, doc_id, metadata, _score = _extract_doc_content(doc, idx)
+            attrs = ""
+            if self.include_attributes:
+                safe_id = html.escape(str(doc_id), quote=True)
+                attrs = f' index="{idx}" id="{safe_id}"'
+                source = metadata.get("source") or metadata.get("title")
+                if source:
+                    safe_source = html.escape(str(source), quote=True)
+                    attrs += f' source="{safe_source}"'
+
+            lines.append(f"  <{self.document_tag}{attrs}>")
+            # Indent each line of content (preserving trailing newlines) and escape untrusted markup
+            safe_content = (
+                content if self.escape_content else escape_context_boundaries(content)
+            )
+            for content_line in safe_content.split("\n"):
+                formatted_line = (
+                    html.escape(escape_context_boundaries(content_line))
+                    if self.escape_content
+                    else content_line
+                )
+                lines.append(f"    {formatted_line}")
+            lines.append(f"  </{self.document_tag}>")
+
+        lines.append(f"</{self.root_tag}>")
+        return "\n".join(lines)
+
+
+class PlainConcatenationFormatter(ContextFormatter):
+    """
+    Formats retrieved documents by joining their raw text with a separator.
+    """
+
+    def __init__(
+        self,
+        separator: str = "\n\n",
+        prefix: str = "",
+        suffix: str = "",
+    ) -> None:
+        self.separator = separator
+        self.prefix = prefix
+        self.suffix = suffix
+
+    def format(
+        self,
+        documents: Sequence[SearchResult | str | dict[str, Any]],
+    ) -> str:
+        if not documents:
+            return ""
+
+        chunks = [
+            escape_context_boundaries(_extract_doc_content(doc, idx)[0])
+            for idx, doc in enumerate(documents, start=1)
+        ]
+        joined = self.separator.join(chunks)
+        return f"{self.prefix}{joined}{self.suffix}"