Ai/rag/embedding.py

"""
سرویس تعبیه‌سازی متن — از Adapter Pattern برای سوئیچ بین providers استفاده می‌کند
"""
import logging
import time

from .api_provider import get_embedding_client
from .config import RAGConfig, load_rag_config
from .observability import classify_exception, log_event, observe_operation, record_metric

logger = logging.getLogger(__name__)

def embed_texts(
    texts: list[str],
    config: RAGConfig | None = None,
    model: str | None = None,
    dimensions: int | None = None,
) -> list[list[float]]:
    """
    تعبیه‌سازی لیست متن‌ها با Avalai.

    Args:
        texts: لیست رشته‌های ورودی
        config: تنظیمات RAG (پیش‌فرض: load_rag_config)
        model: نام مدل (override از config)
        dimensions: تعداد ابعاد (فقط برای مدل‌های پشتیبانی‌کننده)

    Returns:
        لیست وکتورها
    """
    if not texts:
        record_metric("rag.embedding.empty_input", operation="embed_texts")
        return []

    cfg = config or load_rag_config()
    client = get_embedding_client(cfg)
    model_name = model or cfg.embedding.model
    provider = cfg.embedding.provider or "unknown"
    batch_size = cfg.embedding.batch_size

    all_embeddings: list[list[float]] = []
    extra = {}
    if dimensions is not None:
        extra["dimensions"] = dimensions

    with observe_operation(source="rag.embedding", provider=provider, operation="embed_texts"):
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            started_at = time.monotonic()
            try:
                resp = client.embeddings.create(
                    model=model_name,
                    input=batch,
                    **extra,
                )
            except Exception as exc:
                failure = classify_exception(exc)
                log_event(
                    level=logging.ERROR,
                    message="embedding batch request failed",
                    source="rag.embedding",
                    provider=provider,
                    operation="embed_batch",
                    result_status="error",
                    duration_ms=(time.monotonic() - started_at) * 1000,
                    error_code=failure.error_code,
                    batch_size=len(batch),
                    model=model_name,
                )
                raise
            for item in sorted(resp.data, key=lambda x: x.index):
                all_embeddings.append(item.embedding)
            log_event(
                level=logging.INFO,
                message="embedding batch request completed",
                source="rag.embedding",
                provider=provider,
                operation="embed_batch",
                result_status="success",
                duration_ms=(time.monotonic() - started_at) * 1000,
                batch_size=len(batch),
                model=model_name,
            )

    return all_embeddings


def embed_single(text: str, config: RAGConfig | None = None, **kwargs) -> list[float]:
    """تعبیه‌سازی یک متن. خروجی مستقیماً یک وکتور است."""
    vecs = embed_texts([text], config=config, **kwargs)
    return vecs[0] if vecs else []