gen_ai_api

get_embedding_vectors(texts)

Generates embedding vectors for a list of texts using the embedding-specific config.

Parameters:

Name	Type	Description	Default
texts	List[str]	List of texts to generate embeddings for.	required

Returns:

Type	Description
Dict[str, Any]	A dictionary containing: - "embeddings": A numpy array of embedding vectors. - "usage": A dictionary with token usage and cost details.

Source code in src/ariadne/utils/gen_ai_api.py

def get_embedding_vectors(texts: List[str]) -> Dict[str, Any]:
    """
    Generates embedding vectors for a list of texts using the embedding-specific config.

    Args:
        texts: List of texts to generate embeddings for.

    Returns:
        A dictionary containing:
            - "embeddings": A numpy array of embedding vectors.
            - "usage": A dictionary with token usage and cost details.
    """

    client, model, provider = _AIClientFactory.get_client(task_type="embedding")

    batch_size = 100
    verbose = len(texts) > batch_size
    batch_results = []
    total_tokens = 0
    for i in range(0, len(texts), batch_size):
        # if i != 9400:
        #     continue
        batch = texts[i : i + batch_size]
        if verbose:
            print(f"Getting embedding vectors for batch {i + 1} - {i + len(batch)} ({len(texts)} total)")
        response = client.embeddings.create(input=batch, model=model)
        data = sorted(response.data, key=lambda x: x.index)
        np_vectors = np.array([item.embedding for item in data])
        batch_results.append(np_vectors)
        usage = response.usage
        total_tokens = total_tokens + usage.prompt_tokens
    total_cost = _calculate_cost(model, total_tokens, 0, provider)

    return {
        "embeddings": np.concatenate(batch_results, axis=0),
        "usage": {
            "input_tokens": total_tokens,
            "output_tokens": 0,
            "reasoning_tokens": 0,
            "total_cost_usd": total_cost,
            "model_used": model,
        },
    }

get_llm_response(prompt, system_prompt=None, show_reasoning=False, json_schema=None, json_schema_name='ariadne_response', force_json_object=False)

Generates text response using the LLM-specific config.

Parameters:

Name	Type	Description	Default
prompt	str	The user prompt to send to the LLM.	required
system_prompt	Optional[str]	Optional system prompt to guide the LLM's behavior.	None
show_reasoning	bool	If False, strips any text before '' from local reasoning models.	False
json_schema	Optional[Dict[str, Any]]	Optional JSON schema to request strict structured output.	None
json_schema_name	str	Name attached to the schema for providers that require one.	'ariadne_response'
force_json_object	bool	If True, requests a JSON object without strict schema.	False

Returns: A dictionary containing: - "content": The generated text response from the LLM. - "parsed_json": Parsed JSON response when structured output was requested and valid. - "usage": A dictionary with token usage and cost details.

Source code in src/ariadne/utils/gen_ai_api.py

def get_llm_response(
    prompt: str,
    system_prompt: Optional[str] = None,
    show_reasoning: bool = False,
    json_schema: Optional[Dict[str, Any]] = None,
    json_schema_name: str = "ariadne_response",
    force_json_object: bool = False,
) -> Dict[str, Any]:
    """
    Generates text response using the LLM-specific config.

    Args:
        prompt: The user prompt to send to the LLM.
        system_prompt: Optional system prompt to guide the LLM's behavior.
        show_reasoning: If False, strips any text before '</think>' from local reasoning models.
        json_schema: Optional JSON schema to request strict structured output.
        json_schema_name: Name attached to the schema for providers that require one.
        force_json_object: If True, requests a JSON object without strict schema.
    Returns:
        A dictionary containing:
            - "content": The generated text response from the LLM.
            - "parsed_json": Parsed JSON response when structured output was requested and valid.
            - "usage": A dictionary with token usage and cost details.
    """

    client, model, provider = _AIClientFactory.get_client(task_type="llm")

    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": prompt})

    create_kwargs: Dict[str, Any] = {"model": model, "messages": messages}
    if model in _TEMPERATURE_OK_MODELS:
        create_kwargs["temperature"] = 0.0

    structured_requested = json_schema is not None or force_json_object
    if json_schema is not None:
        create_kwargs["response_format"] = {
            "type": "json_schema",
            "json_schema": {
                "name": json_schema_name,
                "strict": True,
                "schema": json_schema,
            },
        }
    elif force_json_object:
        create_kwargs["response_format"] = {"type": "json_object"}

    try:
        response = client.chat.completions.create(**create_kwargs)
    except Exception as e:
        msg = str(e).lower()
        if structured_requested and (
            "response_format" in msg
            or "json_schema" in msg
            or "unsupported" in msg
            or "not supported" in msg
            or "invalid_request" in msg
        ):
            fallback_kwargs = dict(create_kwargs)
            fallback_kwargs.pop("response_format", None)
            response = client.chat.completions.create(**fallback_kwargs)
        elif any(k in msg for k in ("content filter", "content_filter", "filtered", "policy", "moderation")):
            import warnings

            warnings.warn("Content filter triggered; returning None as response content.")
            return {
                "content": None,
                "parsed_json": None,
                "usage": {
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "reasoning_tokens": 0,
                    "total_cost_usd": 0.0,
                    "model_used": model,
                },
            }
        else:
            raise

    usage = response.usage
    reasoning_tokens = 0
    if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
        reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0)

    total_cost = _calculate_cost(model, usage.prompt_tokens, usage.completion_tokens, provider)

    content = response.choices[0].message.content or ""
    if not show_reasoning and isinstance(content, str):
        content = content.split("</think>", 1)[-1]

    parsed_json = None
    if structured_requested and content:
        try:
            parsed_json = json.loads(content)
        except json.JSONDecodeError:
            parsed_json = None

    return {
        "content": content,
        "parsed_json": parsed_json,
        "usage": {
            "input_tokens": usage.prompt_tokens,
            "output_tokens": usage.completion_tokens,
            "reasoning_tokens": reasoning_tokens,
            "total_cost_usd": total_cost,
            "model_used": model,
        },
    }