openai

Classes for augmentation with OpenAI embeddings.

Classes:

• OpenAIAug –

  Class for augmentation with OpenAI embeddings.

OpenAIAug

OpenAIAug(
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
)

Bases: AugmentedBase

Class for augmentation with OpenAI embeddings.

Methods:

• get_embedding –

  Retrieve the embedding for a given text using OpenAI API.

• search –

  Search an encoded query into vector database.

Source code in src/rago/augmented/base.py

def __init__(
    self,
    model_name: Optional[str] = None,
    db: DBBase = FaissDB(),
    top_k: Optional[int] = None,
    api_key: str = '',
    cache: Optional[Cache] = None,
    logs: dict[str, Any] = DEFAULT_LOGS,
) -> None:
    """Initialize AugmentedBase."""
    if logs is DEFAULT_LOGS:
        logs = {}
    super().__init__(api_key=api_key, cache=cache, logs=logs)

    self.db = db

    self.top_k = top_k if top_k is not None else self.default_top_k
    self.model_name = (
        model_name if model_name is not None else self.default_model_name
    )
    self.model = None

    self._validate()
    self._setup()

get_embedding

get_embedding(content: list[str]) -> EmbeddingType

Retrieve the embedding for a given text using OpenAI API.

Source code in src/rago/augmented/openai.py

def get_embedding(self, content: list[str]) -> EmbeddingType:
    """Retrieve the embedding for a given text using OpenAI API."""
    cache_key = sha256(''.join(content).encode('utf-8')).hexdigest()
    cached = self._get_cache(cache_key)
    if cached is not None:
        return cast(EmbeddingType, cached)

    model = cast(openai.OpenAI, self.model)
    response = model.embeddings.create(
        input=content, model=self.model_name
    )
    result = np.array(
        [data.embedding for data in response.data], dtype=np.float32
    )

    self._save_cache(cache_key, result)

    return result

search

search(
    query: str, documents: list[str], top_k: int = 0
) -> list[str]

Search an encoded query into vector database.

Source code in src/rago/augmented/openai.py

def search(
    self, query: str, documents: list[str], top_k: int = 0
) -> list[str]:
    """Search an encoded query into vector database."""
    if not hasattr(self, 'db') or not self.db:
        raise Exception('Vector database (db) is not initialized.')

    # Encode the documents and query
    document_encoded = self.get_embedding(documents)
    query_encoded = self.get_embedding([query])
    top_k = top_k or self.top_k or self.default_top_k or 1

    self.db.embed(document_encoded)
    scores, indices = self.db.search(query_encoded, top_k=top_k)

    self.logs['indices'] = indices
    self.logs['scores'] = scores
    self.logs['search_params'] = {
        'query_encoded': query_encoded,
        'top_k': top_k,
    }

    retrieved_docs = [documents[i] for i in indices if i >= 0]

    return retrieved_docs