eval-framework 0.2.0__py3-none-any.whl

eval_framework/metrics/llm/llm_judge_mtbench_single.py

@@ -0,0 +1,188 @@
+import re
+
+from pydantic import BaseModel
+
+from eval_framework.logger import logger
+from eval_framework.metrics.base import MetricResult
+from eval_framework.metrics.llm.base import BaseLLMJudgeMetric
+from eval_framework.shared.types import BaseMetricContext, Completion, extract_context_metric
+from template_formatting.formatter import Message, Role
+
+SINGLE_JUDGE_PROMPTS = {
+    "single_assistant_single_turn": {
+        "prompt_template": '[System]\nYou are a helpful assistant.\n\n[Instruction]\nPlease act as an impartial judge and evaluate the quality of the response provided by an AI assistant to the user question displayed below. Your evaluation should consider factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of the response. Begin your evaluation by providing a short explanation. Be as objective as possible. After providing your explanation, you must rate the response on a scale of 1 to 10 by strictly following this format: "[[rating]]", for example: "Rating: [[5]]".\n\n[Question]\n{question}\n\n[The Start of Assistant\'s Answer]\n{answer}\n[The End of Assistant\'s Answer]'  # noqa: E501
+    },
+    "single_assistant_multi_turn": {
+        "prompt_template": 'Please act as an impartial judge and evaluate the quality of the response provided by an AI assistant to the user question displayed below. Your evaluation should consider factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of the response. You evaluation should focus on the assistant\'s answer to the second user question. Begin your evaluation by providing a short explanation. Be as objective as possible. After providing your explanation, you must rate the response on a scale of 1 to 10 by strictly following this format: "[[rating]]", for example: "Rating: [[5]]".\n\n<|The Start of Assistant A\'s Conversation with User|>\n\n### User:\n{question_1}\n\n### Assistant A:\n{answer_1}\n\n### User:\n{question_2}\n\n### Assistant A:\n{answer_2}\n\n<|The End of Assistant A\'s Conversation with User|>'  # noqa: E501
+    },
+    "single_assistant_single_turn_w_reference": {
+        "prompt_template": "[System]\nYou are a helpful assistant.\n\n[Instruction]\nPlease act as an impartial judge and evaluate the quality of the response provided by an AI assistant to the user question displayed below. Your evaluation should consider correctness and helpfulness. You will be given a reference answer and the assistant's answer. Begin your evaluation by comparing the assistant's answer with the reference answer. Identify and correct any mistakes. Be as objective as possible. After providing your explanation, you must rate the response on a scale of 1 to 10 by strictly following this format: \"[[rating]]\", for example: \"Rating: [[5]]\".\n\n[Question]\n{question}\n\n[The Start of Reference Answer]\n{ref_answer_1}\n[The End of Reference Answer]\n\n[The Start of Assistant's Answer]\n{answer}\n[The End of Assistant's Answer]"  # noqa: E501
+    },
+    "single_assistant_multi_turn_w_reference": {
+        "prompt_template": "Please act as an impartial judge and evaluate the quality of the response provided by an AI assistant to the user question. Your evaluation should consider correctness and helpfulness. You will be given a reference answer and the assistant's answer. You evaluation should focus on the assistant's answer to the second question. Begin your evaluation by comparing the assistant's answer with the reference answer. Identify and correct any mistakes. Be as objective as possible. After providing your explanation, you must rate the response on a scale of 1 to 10 by strictly following this format: \"[[rating]]\", for example: \"Rating: [[5]]\".\n\n<|The Start of Reference Answer|>\n\n### User:\n{question_1}\n\n### Reference answer:\n{ref_answer_1}\n\n### User:\n{question_2}\n\n### Reference answer:\n{ref_answer_2}\n\n<|The End of Reference Answer|>\n\n\n<|The Start of Assistant A's Conversation with User|>\n\n### User:\n{question_1}\n\n### Assistant A:\n{answer_1}\n\n### User:\n{question_2}\n\n### Assistant A:\n{answer_2}\n\n<|The End of Assistant A's Conversation with User|>"  # noqa: E501
+    },
+}
+
+SINGLE_JUDGE_PROMPTS_DE = {
+    "single_assistant_single_turn": {
+        "prompt_template": '[System]\nDu bist ein hilfreicher Assistent.\n\n[Anweisung]\nBitte agieren Sie als unparteiischer Richter und bewerten Sie die Qualität der Antwort, die von einem KI-Assistenten auf die unten angezeigte Nutzerfrage gegeben wurde. Ihre Bewertung sollte Faktoren wie Nützlichkeit, Relevanz, Genauigkeit, Tiefe, Kreativität und Detailliertheit der Antwort berücksichtigen. Beginnen Sie Ihre Bewertung mit einer kurzen Erklärung. Seien Sie so objektiv wie möglich. Nachdem Sie Ihre Erklärung gegeben haben, müssen Sie die Antwort auf einer Skala von 1 bis 10 bewerten und dabei streng dieses Format einhalten: "[[rating]]", zum Beispiel: "Bewertung: [[5]]".\n\n[Frage]\n{question}\n\n[Der Anfang der Assistentenantwort]\n{answer}\n[Das Ende der Assistentenantwort]'  # noqa: E501
+    },
+    "single_assistant_multi_turn": {
+        "prompt_template": 'Bitte agieren Sie als unparteiischer Richter und bewerten Sie die Qualität der Antwort, die von einem KI-Assistenten auf die unten angezeigte Nutzerfrage gegeben wurde. Ihre Bewertung sollte Faktoren wie Nützlichkeit, Relevanz, Genauigkeit, Tiefe, Kreativität und Detailliertheit der Antwort berücksichtigen. Ihre Bewertung sollte sich auf die Antwort des Assistenten auf die zweite Nutzerfrage konzentrieren. Beginnen Sie Ihre Bewertung mit einer kurzen Erklärung. Seien Sie so objektiv wie möglich. Nachdem Sie Ihre Erklärung gegeben haben, müssen Sie die Antwort auf einer Skala von 1 bis 10 bewerten, wobei Sie streng dieses Format einhalten: "[[rating]]", zum Beispiel: "Bewertung: [[5]]".\n\n<|Der Anfang von Assistent A\'s Unterhaltung mit dem Nutzer|>\n\n### Nutzer:\n{question_1}\n\n### Assistent A:\n{answer_1}\n\n### Nutzer:\n{question_2}\n\n### Assistent A:\n{answer_2}\n\n<|Das Ende von Assistent A\'s Unterhaltung mit dem Nutzer|>'  # noqa: E501
+    },
+    "single_assistant_single_turn_w_reference": {
+        "prompt_template": '[System]\nDu bist ein hilfreicher Assistent.\n\n[Anweisung]\nBitte agieren Sie als unparteiischer Richter und bewerten Sie die Qualität der Antwort, die von einem KI-Assistenten auf die unten angezeigte Nutzerfrage gegeben wurde. Ihre Bewertung sollte Korrektheit und Nützlichkeit berücksichtigen. Ihnen wird eine Referenzantwort und die Antwort des Assistenten gegeben. Beginnen Sie Ihre Bewertung, indem Sie die Antwort des Assistenten mit der Referenzantwort vergleichen. Identifizieren Sie und korrigieren Sie etwaige Fehler. Seien Sie so objektiv wie möglich. Nachdem Sie Ihre Erklärung gegeben haben, müssen Sie die Antwort auf einer Skala von 1 bis 10 bewerten und dabei streng dieses Format einhalten: "[[rating]]", zum Beispiel: "Bewertung: [[5]]".\n\n[Frage]\n{question}\n\n[Der Anfang der Referenzantwort]\n{ref_answer_1}\n[Das Ende der Referenzantwort]\n\n[Der Anfang der Assistentenantwort]\n{answer}\n[Das Ende der Assistentenantwort]'  # noqa: E501
+    },
+    "single_assistant_multi_turn_w_reference": {
+        "propmt_templte": 'Bitte agieren Sie als unparteiischer Richter und bewerten Sie die Qualität der Antwort, die von einem KI-Assistenten auf die Nutzerfrage gegeben wurde. Ihre Bewertung sollte Korrektheit und Nützlichkeit berücksichtigen. Ihnen wird eine Referenzantwort und die Antwort des Assistenten gegeben. Ihre Bewertung sollte sich auf die Antwort des Assistenten auf die zweite Frage konzentrieren. Beginnen Sie Ihre Bewertung, indem Sie die Antwort des Assistenten mit der Referenzantwort vergleichen. Identifizieren und korrigieren Sie etwaige Fehler. Seien Sie so objektiv wie möglich. Nachdem Sie Ihre Erklärung gegeben haben, müssen Sie die Antwort auf einer Skala von 1 bis 10 bewerten, wobei Sie streng dieses Format einhalten: "[[rating]]", zum Beispiel: "Bewertung: [[5]]".\n\n<|Der Anfang der Referenzantwort|>\n\n### Nutzer:\n{question_1}\n\n### Referenzantwort:\n{ref_answer_1}\n\n### Nutzer:\n{question_2}\n\n### Referenzantwort:\n{ref_answer_2}\n\n<|Das Ende der Referenzantwort|>\n\n\n<|Der Anfang von Assistent A\'s Unterhaltung mit dem Nutzer|>\n\n### Nutzer:\n{question_1}\n\n### Assistent A:\n{answer_1}\n\n### Nutzer:\n{question_2}\n\n### Assistent A:\n{answer_2}\n\n<|Das Ende von Assistent A\'s Unterhaltung mit dem Nutzer|>'  # noqa: E501
+    },
+}
+
+
+SINGLE_JUDGE_PROMPTS_FI = {
+    "single_assistant_single_turn": {
+        "prompt_template": '[Järjestelmä]\nOlet avulias avustaja.\n\n[Ohje]\nToimi puolueettomana tuomarina ja arvioi AI-avustajan antaman vastauksen laatua käyttäjän kysymykseen, joka näkyy alla. Arviosi tulisi ottaa huomioon tekijät kuten hyödyllisyys, asiaankuuluvuus, tarkkuus, syvällisyys, luovuus ja yksityiskohtien taso. Aloita arviointisi antamalla lyhyt selitys. Ole mahdollisimman objektiivinen. Selityksen jälkeen sinun on arvioitava vastaus asteikolla 1–10 noudattamalla tarkasti tätä muotoa: "[[arvosana]]", esimerkiksi: "Arvosana: [[5]]".\n\n[Kysymys]\n{question}\n\n[Avustajan vastauksen alku]\n{answer}\n[Avustajan vastauksen loppu]'  # noqa: E501
+    },
+    "single_assistant_multi_turn": {
+        "prompt_template": 'Toimi puolueettomana tuomarina ja arvioi AI-avustajan antaman vastauksen laatua käyttäjän kysymykseen, joka näkyy alla. Arviosi tulisi ottaa huomioon tekijät kuten hyödyllisyys, asiaankuuluvuus, tarkkuus, syvällisyys, luovuus ja yksityiskohtien taso. Arviosi tulisi keskittyä avustajan vastaukseen toiseen käyttäjän kysymykseen. Aloita arviointisi antamalla lyhyt selitys. Ole mahdollisimman objektiivinen. Selityksen jälkeen sinun on arvioitava vastaus asteikolla 1–10 noudattamalla tarkasti tätä muotoa: "[[arvosana]]", esimerkiksi: "Arvosana: [[5]]".\n\n<|Avustaja A:n keskustelun alku käyttäjän kanssa|>\n\n### Käyttäjä:\n{question_1}\n\n### Avustaja A:\n{answer_1}\n\n### Käyttäjä:\n{question_2}\n\n### Avustaja A:\n{answer_2}\n\n<|Avustaja A:n keskustelun loppu käyttäjän kanssa|>'  # noqa: E501
+    },
+    "single_assistant_single_turn_w_reference": {
+        "prompt_template": '[Järjestelmä]\nOlet avulias avustaja.\n\n[Ohje]\nToimi puolueettomana tuomarina ja arvioi AI-avustajan antaman vastauksen laatua käyttäjän kysymykseen, joka näkyy alla. Arviosi tulisi ottaa huomioon oikeellisuus ja hyödyllisyys. Sinulle annetaan viitevastaus ja avustajan vastaus. Aloita arviointisi vertaamalla avustajan vastausta viitevastaukseen. Tunnista ja korjaa mahdolliset virheet. Ole mahdollisimman objektiivinen. Selityksen jälkeen sinun on arvioitava vastaus asteikolla 1–10 noudattamalla tarkasti tätä muotoa: "[[arvosana]]", esimerkiksi: "Arvosana: [[5]]".\n\n[Kysymys]\n{question}\n\n[Viitevastauksen alku]\n{ref_answer_1}\n[Viitevastauksen loppu]\n\n[Avustajan vastauksen alku]\n{answer}\n[Avustajan vastauksen loppu]'  # noqa: E501
+    },
+    "single_assistant_multi_turn_w_reference": {
+        "prompt_template": 'Toimi puolueettomana tuomarina ja arvioi AI-avustajan antaman vastauksen laatua käyttäjän kysymykseen. Arviosi tulisi ottaa huomioon oikeellisuus ja hyödyllisyys. Sinulle annetaan viitevastaus ja avustajan vastaus. Arviosi tulisi keskittyä avustajan vastaukseen toiseen kysymykseen. Aloita arviointisi vertaamalla avustajan vastausta viitevastaukseen. Tunnista ja korjaa mahdolliset virheet. Ole mahdollisimman objektiivinen. Selityksen jälkeen sinun on arvioitava vastaus asteikolla 1–10 noudattamalla tarkasti tätä muotoa: "[[arvosana]]", esimerkiksi: "Arvosana: [[5]]".\n\n<|Viitevastauksen alku|>\n\n### Käyttäjä:\n{question_1}\n\n### Viitevastaus:\n{ref_answer_1}\n\n### Käyttäjä:\n{question_2}\n\n### Viitevastaus:\n{ref_answer_2}\n\n<|Viitevastauksen loppu|>\n\n\n<|Avustaja A:n keskustelun alku käyttäjän kanssa|>\n\n### Käyttäjä:\n{question_1}\n\n### Avustaja A:\n{answer_1}\n\n### Käyttäjä:\n{question_2}\n\n### Avustaja A:\n{answer_2}\n\n<|Avustaja A:n keskustelun loppu käyttäjän kanssa|>'  # noqa: E501
+    },
+}
+
+NEED_REF_CATEGORIES = ["math", "reasoning", "coding", "arena-hard-200"]
+
+
+class PromptToJudge(BaseModel):
+    comparison_type: str
+    prompt_text: str
+
+
+class MTBenchJudgeSingleMetricContext(BaseMetricContext):
+    category: str
+    reference: list[str] | str | None
+
+
+def generate_single_judge_prompts(response: Completion) -> list[PromptToJudge]:
+    context = extract_context_metric(response, MTBenchJudgeSingleMetricContext)
+
+    assert response.messages is not None
+
+    if response.subject.startswith("de"):
+        prompt_templates = SINGLE_JUDGE_PROMPTS_DE
+    elif response.subject.startswith("fi"):
+        prompt_templates = SINGLE_JUDGE_PROMPTS_FI
+    else:
+        prompt_templates = SINGLE_JUDGE_PROMPTS
+    prompts_to_judge = []
+
+    assert context.category is not None, "Category must be provided in the context for MTBenchJudgeSingleMetricContext"
+
+    # No reference answer needed
+    if context.category not in NEED_REF_CATEGORIES:
+        # SINLGE TURN
+        if len(response.messages) <= 2:
+            # turn 1
+            question = response.last_user_instruction
+            answer = response.completion
+            # format prompt
+            single_turn_prompt = prompt_templates["single_assistant_single_turn"]["prompt_template"].format(
+                question=question,
+                answer=answer,
+            )
+            prompts_to_judge.append(PromptToJudge(comparison_type="single_judgement", prompt_text=single_turn_prompt))
+            # MULTI TURN
+        else:
+            # turn 1
+            question_1 = response.first_user_instruction
+            answer_1 = response.messages[1].content
+            # turn 2
+            question_2 = response.last_user_instruction
+            answer_2 = response.completion
+            # format prompt
+            multi_turn_prompt = prompt_templates["single_assistant_multi_turn"]["prompt_template"].format(
+                question_1=question_1, answer_1=answer_1, question_2=question_2, answer_2=answer_2
+            )
+            prompts_to_judge.append(PromptToJudge(comparison_type="single_judgement", prompt_text=multi_turn_prompt))
+    # Reference answer needed
+    elif context.reference:
+        # SINGLE TURN
+        if len(response.messages) <= 2 and len(context.reference) >= 1:
+            # turn 1
+            question = response.last_user_instruction
+            answer = response.completion
+            ref_answer = context.reference[0]
+            # format prompt
+            single_turn_prompt = prompt_templates["single_assistant_single_turn_w_reference"]["prompt_template"].format(
+                question=question,
+                answer=answer,
+                ref_answer_1=ref_answer,
+            )
+            prompts_to_judge.append(PromptToJudge(comparison_type="single_judgement", prompt_text=single_turn_prompt))
+        # MULTI TURN
+        elif len(context.reference) >= 2:
+            # turn 1
+            question_1 = response.first_user_instruction
+            answer_1 = response.messages[1].content
+            ref_answer_1 = context.reference[0]
+            # turn 2
+            question_2 = response.last_user_instruction
+            answer_2 = response.completion
+            ref_answer_2 = context.reference[1]
+            # format prompt
+            multi_turn_prompt = prompt_templates["single_assistant_multi_turn_w_reference"]["prompt_template"].format(
+                question_1=question_1,
+                answer_1=answer_1,
+                ref_answer_1=ref_answer_1,
+                question_2=question_2,
+                answer_2=answer_2,
+                ref_answer_2=ref_answer_2,
+            )
+            prompts_to_judge.append(PromptToJudge(comparison_type="single_judgement", prompt_text=multi_turn_prompt))
+    else:
+        logger.info(
+            f"Warning: No reference answer found for this sample (category: "
+            f"{context.category}), even though it is needed."
+        )
+
+    return prompts_to_judge
+
+
+class MTBenchJudgeSingle(BaseLLMJudgeMetric):
+    NAME = "single_judgement"
+
+    def calculate(self, response: Completion) -> list[MetricResult]:
+        try:
+            prompts_to_judge: list[PromptToJudge] = generate_single_judge_prompts(response)
+
+            all_metrics = []
+            for prompt_to_judge in prompts_to_judge:
+                messages = [Message(role=Role.USER, content=prompt_to_judge.prompt_text)]
+                output = self._llm_judge.generate_from_messages([messages])
+                parsed_output = self._output_to_rating(output[0].completion)
+
+                all_metrics.append(
+                    MetricResult(
+                        metric_name=prompt_to_judge.comparison_type,
+                        value=parsed_output,
+                        higher_is_better=True,
+                    )
+                )
+
+            return all_metrics
+
+        except KeyError as e:
+            logger.info(f"LLM judge did not produce an expected output, sample will be ignored in aggregations. {e}")
+            return []
+
+    @staticmethod
+    def _output_to_rating(output: str) -> float:
+        match = re.search(r"\[\[(\d+)\]\]", output)
+
+        if match:
+            return float(match.group(1))
+        return 0

eval_framework/metrics/llm/llm_judge_refusal.py

@@ -0,0 +1,35 @@
+from eval_framework.llm.base import BaseLLM
+from eval_framework.metrics.base import MetricResult
+from eval_framework.metrics.llm.base import BaseLLMJudgeMetric
+from eval_framework.metrics.llm.graders.language import Language
+from eval_framework.metrics.llm.graders.refusal_grader import RefusalGrader
+from eval_framework.shared.types import Completion
+
+
+class LLMJudgeRefusal(BaseLLMJudgeMetric):
+    NAME = "refusal_classifier"
+
+    def __init__(self, llm_judge: BaseLLM):
+        super().__init__(llm_judge)
+        self._grader = RefusalGrader(llm_judge)
+
+    def calculate(self, response: Completion) -> list[MetricResult]:
+        if response.error is not None:
+            return [MetricResult(metric_name=self.NAME, value=None, higher_is_better=False, error=response.error)]
+
+        language = response.get_completion_language() or response.get_instruction_language() or "en"
+        grading = self._grader.grade(
+            completion=response.sanitized_completion,
+            language=Language(language),
+        )
+
+        return [
+            MetricResult(
+                metric_name=self.NAME,
+                value=float(grading.is_refusal) if grading.is_refusal is not None else None,
+                higher_is_better=True,
+                llm_judge_prompt=grading.judge_prompt,
+                llm_judge_response=grading.judge_response,
+                error=response.error,
+            )
+        ]

eval_framework/metrics/llm/llm_judge_sql.py

@@ -0,0 +1,394 @@
+import atexit
+import logging
+import random
+import re
+import signal
+import sqlite3
+import threading
+from enum import Enum
+from time import sleep
+from typing import Any
+from uuid import uuid4
+
+import docker
+import mysql.connector
+import mysql.connector.abstracts
+import psycopg2  # type: ignore
+from pydantic import BaseModel
+
+from eval_framework.llm.base import BaseLLM
+from eval_framework.metrics.base import MetricResult
+from eval_framework.metrics.llm.base import BaseLLMJudgeMetric
+from eval_framework.metrics.llm.graders.language import Language
+from eval_framework.metrics.llm.graders.sql_quality_grader import SqlQualityGrader
+from eval_framework.shared.types import Completion, LanguageMetricContext, extract_context_metric
+from eval_framework.tasks.utils import get_docker_address
+
+logger = logging.getLogger(__name__)
+
+
+class SqlDialects(Enum):
+    sqlite = "sqlite"
+    postgres = "postgresql"
+    mysql = "mysql"
+    standard_sql = "standard_sql"
+
+
+class SqlOutputComparison(BaseModel):
+    matches_results_count: bool
+    matches_column_count: bool
+    results_equal: bool
+
+
+class SqlValidationResult(BaseModel):
+    success: bool
+    schema_error: str | None = None
+    query_error: str | None = None
+    results: list[Any] = []
+
+
+class LLMJudgeSqlMetricContext(LanguageMetricContext):
+    dialect: str
+    db_schema: str
+
+
+_DOCKER_LAUNCH_LOCK = threading.Lock()
+_MYSQL_PORT = 0
+_POSTGRES_PORT = 0
+
+
+class LLMJudgeSql(BaseLLMJudgeMetric):
+    NAME = "SQL Quality"
+
+    def __init__(self, llm_judge: BaseLLM):
+        super().__init__(llm_judge)
+        self._grader = SqlQualityGrader(llm_judge)
+
+        self.postgres_password = "mysecretpassword"
+        self.postgres_user = "postgres"
+
+        self.mysql_password = "mysecretpassword"
+        self.mysql_user = "root"
+        self.mysql_db_name = "mysql"
+
+        with _DOCKER_LAUNCH_LOCK:
+            if _MYSQL_PORT != 0 and _POSTGRES_PORT != 0:
+                return
+            self.client = docker.from_env()
+            atexit.register(self._shutdown_dbs)
+            signal.signal(signal.SIGTERM, lambda *_: self._shutdown_dbs())
+            self._start_postgres_db()
+            self._start_mysql_db()
+            self._wait_for_db_containers()
+
+    def calculate(self, response: Completion) -> list[MetricResult]:
+        if response.error is not None:
+            return [
+                MetricResult(metric_name=f"{self.NAME}/{k}", value=None, higher_is_better=True, error=response.error)
+                for k in [
+                    "successfully_runs",
+                    "is_just_sql",
+                    "matches_results_count",
+                    "matches_column_count",
+                    "results_equal",
+                    "llm_quality_score",
+                ]
+            ]
+
+        context = extract_context_metric(response, LLMJudgeSqlMetricContext)
+
+        assert isinstance(response.ground_truth, str)
+
+        schema_id = str(uuid4()).replace("-", "_")
+
+        expected_result = self.validate_query(
+            SqlDialects(context.dialect),
+            context.db_schema,
+            response.ground_truth,
+            f"golden_{schema_id}",
+        )
+        completion_stripped = response.completion.strip().strip("```sql").strip("```")
+        completion_query = extract_query_from_completions(completion_stripped)
+        if completion_query:
+            result = self.validate_query(
+                SqlDialects(context.dialect),
+                context.db_schema,
+                completion_query,
+                f"completion_{schema_id}",
+            )
+        else:
+            result = None
+
+        results = [
+            MetricResult(
+                metric_name=f"{self.NAME}/successfully_runs",
+                value=float(result is not None and result.success),
+                higher_is_better=True,
+                error=response.error,
+            ),
+            MetricResult(
+                metric_name=f"{self.NAME}/is_just_sql",
+                value=float(completion_query == completion_stripped),
+                higher_is_better=True,
+                error=response.error,
+            ),
+        ]
+
+        if result is not None and result.success:
+            output_comparison = SqlOutputComparison(
+                matches_results_count=len(expected_result.results) == len(result.results),
+                matches_column_count=count_result_columns(expected_result.results)
+                == count_result_columns(result.results),
+                results_equal=expected_result.results == result.results,
+            )
+            results.extend(
+                [
+                    MetricResult(
+                        metric_name=f"{self.NAME}/matches_results_count",
+                        value=float(output_comparison.matches_results_count),
+                        higher_is_better=True,
+                        error=response.error,
+                    ),
+                    MetricResult(
+                        metric_name=f"{self.NAME}/matches_column_count",
+                        value=float(output_comparison.matches_column_count),
+                        higher_is_better=True,
+                        error=response.error,
+                    ),
+                    MetricResult(
+                        metric_name=f"{self.NAME}/results_equal",
+                        value=float(output_comparison.results_equal),
+                        higher_is_better=True,
+                        error=response.error,
+                    ),
+                ]
+            )
+
+        grading = self._grader.grade(
+            prompt=response.user_instruction,
+            completion=completion_stripped,
+            result=result.results if result and result.success else None,
+            language=Language(response.get_instruction_language()),
+        )
+        results.append(
+            MetricResult(
+                metric_name=f"{self.NAME}/llm_quality_score",
+                # [0, 1] normalization required for visualizer
+                value=(float(grading.query_quality) - 1) / 4 if grading.query_quality is not None else None,
+                higher_is_better=True,
+                llm_judge_prompt=grading.judge_prompt,
+                llm_judge_response=grading.judge_response,
+                error=response.error,
+            )
+        )
+
+        return results
+
+    def _start_postgres_db(self) -> None:
+        global _POSTGRES_PORT
+        for _ in range(10):  # find a free port
+            try:
+                _POSTGRES_PORT = random.randint(1000, 65535)
+                self.postgres_docker = self.client.containers.run(
+                    "docker.io/postgres",
+                    environment={"POSTGRES_PASSWORD": self.postgres_password},
+                    ports={5432: _POSTGRES_PORT},
+                    tty=True,
+                    auto_remove=True,
+                    detach=True,
+                    network_mode="bridge",
+                )
+                break
+            except docker.errors.APIError as e:
+                if "port is already allocated" not in str(e):
+                    raise e
+                continue
+
+    def _start_mysql_db(self) -> None:
+        global _MYSQL_PORT
+        for _ in range(10):  # find a free port
+            try:
+                _MYSQL_PORT = random.randint(1000, 65535)
+                self.mysql_docker = self.client.containers.run(
+                    "docker.io/mysql:latest",
+                    environment={"MYSQL_ROOT_PASSWORD": self.mysql_password, "MYSQL_DATABASE": self.mysql_db_name},
+                    ports={3306: _MYSQL_PORT},
+                    tty=True,
+                    auto_remove=True,
+                    detach=True,
+                    network_mode="bridge",
+                )
+                break
+            except docker.errors.APIError as e:
+                if "port is already allocated" not in str(e):
+                    raise e
+                continue
+
+    def _wait_for_db_containers(self) -> None:
+        for _ in range(600):
+            try:
+                con = self.connect_to_postgres()
+                con.close()
+                con = self.connect_to_mysql()
+                con.close()
+                return
+            except Exception:
+                logger.info("Could not connect to DBs yet...")
+                sleep(1)
+        raise Exception("DBs not available.")
+
+    def _shutdown_dbs(self) -> None:
+        if hasattr(self, "postgres_docker"):
+            self.postgres_docker.kill()
+        if hasattr(self, "mysql_docker"):
+            self.mysql_docker.kill()
+
+    def validate_query(
+        self,
+        dialect: SqlDialects,
+        create_db_statements: str,
+        sql_query: str,
+        db_schema: str,
+    ) -> SqlValidationResult:
+        match dialect:
+            case SqlDialects.sqlite | SqlDialects.standard_sql:
+                return self.validate_query_sqlite(create_db_statements, sql_query, f"{dialect.value}_{db_schema}")
+            case SqlDialects.postgres:
+                return self.validate_query_postgres(create_db_statements, sql_query, f"{dialect.value}_{db_schema}")
+            case SqlDialects.mysql:
+                return self.validate_query_mysql(create_db_statements, sql_query, f"{dialect.value}_{db_schema}")
+            case _:
+                raise NotImplementedError(f"Query validation not implemented for {dialect.value}.")
+
+    def validate_query_sqlite(self, create_db_statements: str, sql_query: str, db_schema: str) -> SqlValidationResult:
+        con = sqlite3.connect(":memory:")
+        cur = con.cursor()
+        try:
+            statements = separate_statements(create_db_statements)
+            for statement in statements:
+                cur.execute(statement)
+                con.commit()
+        except Exception as e:
+            logger.info(f"Create statements are not compatible with SQLite. Reason: {e}")
+            return SqlValidationResult(success=False, schema_error=str(e))
+        try:
+            queries = separate_statements(sql_query)
+            for query in queries:
+                cur.execute(query)
+                con.commit()
+            results = cur.fetchall()
+        except Exception as e:
+            logger.info(f"SQL query is not compatible with SQLite. Reason: {e}")
+            return SqlValidationResult(success=False, query_error=str(e))
+
+        con.close()
+        return SqlValidationResult(success=True, results=results)
+
+    def connect_to_postgres(self) -> psycopg2.extensions.connection:
+        conn_params = {
+            "dbname": "postgres",
+            "user": self.postgres_user,
+            "password": self.postgres_password,
+            "host": get_docker_address(),
+            "port": _POSTGRES_PORT,
+        }
+        return psycopg2.connect(**conn_params)
+
+    def validate_query_postgres(self, create_db_statements: str, sql_query: str, db_schema: str) -> SqlValidationResult:
+        con = self.connect_to_postgres()
+        cur = con.cursor()
+        cur.execute(f"CREATE SCHEMA {db_schema};")
+        con.commit()
+        cur.execute(f"ALTER USER {self.postgres_user} set SEARCH_PATH = {db_schema};")
+        con.commit()
+        try:
+            statements = separate_statements(create_db_statements)
+            for statement in statements:
+                cur.execute(statement)
+            con.commit()
+        except Exception as e:
+            logger.info(f"Create statements are not compatible with PostgreSQL. Reason: {e}")
+            return SqlValidationResult(success=False, schema_error=str(e))
+        try:
+            queries = separate_statements(sql_query)
+            for query in queries:
+                cur.execute(query)
+                con.commit()
+            results = cur.fetchall()
+        except Exception as e:
+            logger.info(f"SQL query is not compatible with PostgreSQL. Reason: {e}")
+            return SqlValidationResult(success=False, query_error=str(e))
+
+        con.commit()
+
+        con.close()
+        return SqlValidationResult(success=True, results=results)
+
+    def connect_to_mysql(
+        self,
+    ) -> mysql.connector.pooling.PooledMySQLConnection | mysql.connector.abstracts.MySQLConnectionAbstract:
+        conn_params = {
+            "database": self.mysql_db_name,
+            "user": self.mysql_user,
+            "password": self.mysql_password,
+            "host": get_docker_address(),
+            "port": _MYSQL_PORT,
+        }
+        return mysql.connector.connect(**conn_params)
+
+    def validate_query_mysql(self, create_db_statements: str, sql_query: str, db_schema: str) -> SqlValidationResult:
+        con = self.connect_to_mysql()
+        cur = con.cursor(buffered=True)
+        cur.execute(f"CREATE SCHEMA {db_schema};")
+        con.commit()
+        cur.execute(f"USE {db_schema};")
+        try:
+            statements = separate_statements(create_db_statements)
+            for statement in statements:
+                cur.execute(statement)
+                con.commit()
+        except Exception as e:
+            logger.info(f"Create statements are not compatible with MySQL. Reason: {e}")
+            con.close()
+            return SqlValidationResult(success=False, schema_error=str(e))
+        try:
+            queries = separate_statements(sql_query)
+            for query in queries:
+                cur.execute(query)
+                con.commit()
+            results = cur.fetchall()
+        except Exception as e:
+            logger.info(f"SQL query is not compatible with MySQL. Reason: {e}")
+            con.close()
+            return SqlValidationResult(success=False, query_error=str(e))
+
+        cur.close()
+        con.close()
+        return SqlValidationResult(success=True, results=results)
+
+
+def separate_statements(statements: str) -> list[str]:
+    return statements.split(";")[:-1]
+
+
+def is_create_table_statement(statement: str) -> bool:
+    return "CREATE TABLE" in statement
+
+
+def count_result_columns(result: list[Any]) -> int:
+    if len(result) == 0:
+        return 0
+    return len(result[0])
+
+
+def extract_query_from_completions(completion: str) -> str | None:
+    # Match SQL blocks starting with SELECT or WITH at line start
+    # (allowing punctuation/whitespace), ending at first semicolon
+    pattern = re.compile(r"(?:^|\n)[^a-zA-Z0-9_]*((?:select|with)\b.*?;)", re.IGNORECASE | re.DOTALL)
+
+    matches = pattern.findall(completion)
+
+    # Return the query only if exactly one match is found
+    if len(matches) == 1:
+        return matches[0].strip()
+    return None

eval_framework/metrics/llm/llm_judge_world_knowledge.py

@@ -0,0 +1,37 @@
+from eval_framework.llm.base import BaseLLM
+from eval_framework.metrics.base import MetricResult
+from eval_framework.metrics.llm.base import BaseLLMJudgeMetric
+from eval_framework.metrics.llm.graders.language import Language
+from eval_framework.metrics.llm.graders.summary_world_knowledge_grader import SummarizationWorldKnowledgeGrader
+from eval_framework.shared.types import Completion
+
+
+class LLMJudgeWorldKnowledge(BaseLLMJudgeMetric):
+    NAME = "World Knowledge"
+
+    def __init__(self, llm_judge: BaseLLM):
+        super().__init__(llm_judge)
+        self._grader = SummarizationWorldKnowledgeGrader(llm_judge)
+
+    def calculate(self, response: Completion) -> list[MetricResult]:
+        if response.error is not None:
+            return [MetricResult(metric_name=self.NAME, value=None, higher_is_better=False, error=response.error)]
+
+        language = Language(response.get_instruction_language())
+
+        grading = self._grader.grade(
+            reference_input=response.user_instruction,
+            completion=response.sanitized_completion,
+            language=language,
+        )
+
+        return [
+            MetricResult(
+                metric_name=self.NAME,
+                value=float(grading.contains_world_knowledge) if grading.contains_world_knowledge is not None else None,
+                higher_is_better=False,
+                llm_judge_prompt=grading.judge_prompt,
+                llm_judge_response=grading.judge_response,
+                error=response.error,
+            )
+        ]