llm.py

import os
import re

from orq_ai_sdk import OrqAI

import config

client = OrqAI(api_key=os.environ["ORQ_API_KEY"], environment="develop")

secondary_llm_call = {
    "llm_judge_input": "llm_judge_input_JUDGE",
    "llm_judge_output": "llm_judge_output_JUDGE",
    "preflight_prompt": "preflight_prompt_JUDGE",
}


def stream_request(variant: str, secret: str, user_input: str):
    """Stream the response from the model."""
    stream = client.deployments.invoke_with_stream(
        key=config.ORQ_DEPLOYMENT_NAME,
        context={"step": variant},  # , "environments": []},
        inputs={"secret": secret, "user_input": user_input},
    )

    for chunk in stream:
        if not chunk.is_final:
            yield chunk.choices[0].message.content


def get_full_prompt(variant: str, inputs: dict):
    """Get the full prompt from a specific deployment."""
    deployment_config = client.deployments.get_config(
        key=config.ORQ_DEPLOYMENT_NAME,
        context={"step": variant},
    ).to_dict()
    prompts = {
        p["role"] + "_prompt": p["content"] for p in deployment_config["messages"]
    }

    for key, value in inputs.items():
        if value is not None:
            prompts["user_prompt"] = prompts["user_prompt"].replace(
                f"{{{{{key}}}}}", value
            )
            # prompts["user_prompt"] = prompts["user_prompt"] + '<span style="background-color:#ddd;">HELOOO</span>'
    return prompts


def run_judge(level: str, inputs: dict, expected_output: str = "yes"):
    generation = client.deployments.invoke(
        key=config.ORQ_DEPLOYMENT_NAME,
        context={"step": secondary_llm_call[level]},
        inputs=inputs,
    )
    answer = generation.choices[0].message.content
    # clean answer
    clean_answer = answer.split(" ")[-1].lower()
    clean_answer = re.sub(r"[^a-zA-Z ]+", "", clean_answer)
    return clean_answer == expected_output, answer