Embeddings & Semantic Search

class SemanticEmbeddingRequest:
    prompt: Prompt
    representation: SemanticRepresentation
    compress_to_size: Optional[int] = None
    normalize: bool = False
    contextual_control_threshold: Optional[float] = None
    control_log_additive: Optional[bool] = True
    """
    Request for semantic embeddings optimized for search tasks.
    
    Attributes:
    - prompt: Input prompt (text, image, or multimodal)
    - representation: Embedding type for specific use cases
    - compress_to_size: Target embedding dimension (compression)
    - normalize: Normalize embeddings to unit length
    - contextual_control_threshold: Threshold for attention controls
    - control_log_additive: How to apply attention controls
    """

class SemanticEmbeddingResponse:
    model_version: str
    embedding: EmbeddingVector
    num_tokens_prompt_total: int
    message: Optional[str] = None
    """
    Response from semantic embedding request.
    
    Attributes:
    - model_version: Version of model used
    - embedding: Generated embedding vector
    - num_tokens_prompt_total: Total tokens processed
    - message: Optional response message
    """

def semantic_embed(
    self, 
    request: SemanticEmbeddingRequest, 
    model: str
) -> SemanticEmbeddingResponse:
    """
    Generate semantic embedding for single prompt.
    
    Parameters:
    - request: Embedding configuration
    - model: Model name to use
    
    Returns:
    SemanticEmbeddingResponse with embedding vector
    """

class BatchSemanticEmbeddingRequest:
    prompts: Sequence[Prompt]
    representation: SemanticRepresentation
    compress_to_size: Optional[int] = None
    normalize: bool = False
    contextual_control_threshold: Optional[float] = None
    control_log_additive: Optional[bool] = True
    """
    Request for batch semantic embeddings.
    
    Attributes:
    - prompts: Sequence of input prompts to embed
    - representation: Embedding type for all prompts
    - compress_to_size: Target embedding dimension
    - normalize: Normalize all embeddings
    - contextual_control_threshold: Threshold for attention controls
    - control_log_additive: How to apply attention controls
    """

class BatchSemanticEmbeddingResponse:
    model_version: str
    embeddings: Sequence[EmbeddingVector]
    num_tokens_prompt_total: int
    """
    Response from batch semantic embedding request.
    
    Attributes:
    - model_version: Version of model used
    - embeddings: Generated embedding vectors (same order as input)
    - num_tokens_prompt_total: Total tokens processed across all prompts
    """

def batch_semantic_embed(
    self, 
    request: BatchSemanticEmbeddingRequest, 
    model: Optional[str] = None
) -> BatchSemanticEmbeddingResponse:
    """
    Generate semantic embeddings for multiple prompts (sync).
    
    Parameters:
    - request: Batch embedding configuration
    - model: Model name to use (optional for some endpoints)
    
    Returns:
    BatchSemanticEmbeddingResponse with embedding vectors
    """

async def batch_semantic_embed(
    self,
    request: BatchSemanticEmbeddingRequest,
    model: Optional[str] = None,
    num_concurrent_requests: int = 1,
    batch_size: int = 100,
    progress_bar: bool = False
) -> BatchSemanticEmbeddingResponse:
    """
    Generate semantic embeddings for multiple prompts (async with controls).
    
    Parameters:
    - request: Batch embedding configuration
    - model: Model name to use
    - num_concurrent_requests: Number of concurrent API requests
    - batch_size: Maximum prompts per batch
    - progress_bar: Show progress bar during processing
    
    Returns:
    BatchSemanticEmbeddingResponse with embedding vectors
    """

class EmbeddingRequest:
    prompt: Prompt
    layers: List[int]
    pooling: List[str]
    type: Optional[str] = None
    tokens: bool = False
    normalize: bool = False
    contextual_control_threshold: Optional[float] = None
    control_log_additive: Optional[bool] = True
    """
    Request for layer-based embeddings (legacy API).
    
    Attributes:
    - prompt: Input prompt
    - layers: Layer indices to extract embeddings from
    - pooling: Pooling operations to apply
    - type: Embedding type specification
    - tokens: Return token strings along with embeddings
    - normalize: Normalize embeddings
    - contextual_control_threshold: Threshold for attention controls
    - control_log_additive: How to apply attention controls
    """

class EmbeddingResponse:
    model_version: str
    num_tokens_prompt_total: int
    embeddings: Optional[Dict[Tuple[str, str], List[float]]]
    tokens: Optional[List[str]]
    message: Optional[str] = None
    """
    Response from layer-based embedding request.
    
    Attributes:
    - model_version: Version of model used
    - num_tokens_prompt_total: Total tokens processed
    - embeddings: Embeddings keyed by (layer, pooling) tuple
    - tokens: Token strings (if requested)
    - message: Optional response message
    """

def embed(self, request: EmbeddingRequest, model: str) -> EmbeddingResponse:
    """
    Generate layer-based embeddings.
    
    Parameters:
    - request: Embedding configuration
    - model: Model name to use
    
    Returns:
    EmbeddingResponse with layer-specific embeddings
    """

class EmbeddingV2Request:
    input: Union[str, List[str], List[int], List[List[int]]]
    dimensions: Optional[int] = None
    encoding_format: Optional[Literal["float", "base64"]] = None
    """
    OpenAI-compatible embedding request.
    
    Attributes:
    - input: Text strings or token arrays to embed
    - dimensions: Target embedding dimensions
    - encoding_format: Output encoding format
    """

class EmbeddingV2Response:
    object: str
    data: List[EmbeddingV2ResponseData]
    model: str
    usage: Usage
    """
    OpenAI-compatible embedding response.
    
    Attributes:
    - object: Response object type
    - data: Embedding data for each input
    - model: Model name used
    - usage: Token usage statistics
    """

def embeddings(
    self, 
    request: EmbeddingV2Request, 
    model: str
) -> EmbeddingV2Response:
    """
    Generate OpenAI-compatible embeddings.
    
    Parameters:
    - request: OpenAI-style embedding configuration
    - model: Model name to use
    
    Returns:
    EmbeddingV2Response with embeddings
    """

class SemanticRepresentation(Enum):
    Symmetric = "symmetric"  # For similarity/clustering tasks
    Document = "document"    # For document representation in search
    Query = "query"         # For query representation in search

# Type alias for embedding vectors
EmbeddingVector = List[float]

# Available pooling operations
POOLING_OPTIONS: List[str] = ["mean", "max", "last_token", "abs_max"]

from aleph_alpha_client import (
    Client, AsyncClient,
    SemanticEmbeddingRequest, SemanticEmbeddingResponse,
    BatchSemanticEmbeddingRequest,
    EmbeddingRequest, EmbeddingV2Request,
    SemanticRepresentation, Prompt
)

client = Client(token="your-api-token")

# Simple semantic embedding
request = SemanticEmbeddingRequest(
    prompt=Prompt.from_text("Machine learning is transforming technology"),
    representation=SemanticRepresentation.Symmetric,
    normalize=True
)
response = client.semantic_embed(request, model="luminous-extended")
embedding = response.embedding  # List[float]
print(f"Embedding dimension: {len(embedding)}")

# Document and query embeddings for search
documents = [
    "Python is a programming language",
    "Machine learning uses neural networks", 
    "Data science involves statistical analysis"
]

# Embed documents
doc_prompts = [Prompt.from_text(doc) for doc in documents]
doc_request = BatchSemanticEmbeddingRequest(
    prompts=doc_prompts,
    representation=SemanticRepresentation.Document,
    normalize=True
)
doc_response = client.batch_semantic_embed(doc_request, model="luminous-extended")
doc_embeddings = doc_response.embeddings

# Embed query
query_request = SemanticEmbeddingRequest(
    prompt=Prompt.from_text("What is Python programming?"),
    representation=SemanticRepresentation.Query,
    normalize=True
)
query_response = client.semantic_embed(query_request, model="luminous-extended")
query_embedding = query_response.embedding

# Calculate similarities (cosine similarity for normalized vectors)
import numpy as np

similarities = []
for doc_emb in doc_embeddings:
    similarity = np.dot(query_embedding, doc_emb)
    similarities.append(similarity)

# Find most similar document
best_match_idx = np.argmax(similarities)
print(f"Most similar document: {documents[best_match_idx]}")
print(f"Similarity score: {similarities[best_match_idx]}")

# Multimodal embeddings
from aleph_alpha_client import Image

image = Image.from_file("diagram.png")
multimodal_prompt = Prompt([
    Text.from_text("Technical diagram showing:"),
    image
])

multimodal_request = SemanticEmbeddingRequest(
    prompt=multimodal_prompt,
    representation=SemanticRepresentation.Symmetric,
    normalize=True
)
multimodal_response = client.semantic_embed(multimodal_request, model="luminous-extended")

# Batch processing with async client
import asyncio

async def batch_embed_async():
    async with AsyncClient(token="your-api-token") as client:
        # Large batch with concurrent processing
        large_batch = [Prompt.from_text(f"Document {i}") for i in range(1000)]
        
        request = BatchSemanticEmbeddingRequest(
            prompts=large_batch,
            representation=SemanticRepresentation.Document,
            normalize=True
        )
        
        response = await client.batch_semantic_embed(
            request,
            model="luminous-extended", 
            num_concurrent_requests=5,  # 5 concurrent API calls
            batch_size=50,              # 50 prompts per batch
            progress_bar=True           # Show progress
        )
        
        print(f"Generated {len(response.embeddings)} embeddings")
        print(f"Total tokens: {response.num_tokens_prompt_total}")

asyncio.run(batch_embed_async())

# Legacy embedding API with layer extraction
legacy_request = EmbeddingRequest(
    prompt=Prompt.from_text("Text for layer analysis"),
    layers=[8, 12, 16],  # Extract from layers 8, 12, 16
    pooling=["mean", "max"],  # Apply mean and max pooling
    tokens=True,  # Return token strings
    normalize=True
)
legacy_response = client.embed(legacy_request, model="luminous-extended")

# Access layer-specific embeddings
for (layer, pooling), embedding in legacy_response.embeddings.items():
    print(f"Layer {layer}, {pooling} pooling: {len(embedding)} dimensions")

if legacy_response.tokens:
    print(f"Tokens: {legacy_response.tokens}")

# OpenAI-compatible API
openai_request = EmbeddingV2Request(
    input=["Hello world", "Machine learning", "Data science"],
    dimensions=512,  # Compress to 512 dimensions
    encoding_format="float"
)
openai_response = client.embeddings(openai_request, model="luminous-extended")

for i, embedding_data in enumerate(openai_response.data):
    print(f"Input {i}: {len(embedding_data.embedding)} dimensions")

print(f"Usage: {openai_response.usage.total_tokens} tokens")