{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Judging and Routing — Optimizing Resource Usage by Evaluating Problem Complexity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the original Lab 2, we explored the **Orchestrator–Worker pattern**, where a planner sent the same question to multiple agents, and a judge assessed their responses to evaluate agent intelligence.\n",
    "\n",
    "In this notebook, we extend that design by adding multiple judges and a routing component to optimize model usage based on task complexity. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports and Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "from anthropic import Anthropic\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')\n",
    "if openai_api_key and google_api_key and deepseek_api_key:\n",
    "    print(\"All keys were loaded successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ollama pull llama3.2\n",
    "!ollama pull mistral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The notebook uses instances of GPT, Gemini and DeepSeek APIs, along with two local models served via Ollama: ```llama3.2```  and ```mistral```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_specs = {\n",
    "    \"gpt-4o-mini\" : None,\n",
    "    \"gemini-2.0-flash\": {\n",
    "        \"api_key\" : google_api_key,\n",
    "        \"url\" : \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "    },\n",
    "    \"deepseek-chat\" : {\n",
    "        \"api_key\" : deepseek_api_key,\n",
    "        \"url\" : \"https://api.deepseek.com/v1\"\n",
    "    },\n",
    "    \"llama3.2\" : {\n",
    "        \"api_key\" : \"ollama\",\n",
    "        \"url\" : \"http://localhost:11434/v1\"\n",
    "    },\n",
    "    \"mistral\" : {\n",
    "        \"api_key\" : \"ollama\",\n",
    "        \"url\" : \"http://localhost:11434/v1\"\n",
    "    }\n",
    "}\n",
    "\n",
    "def create_model(model_name):\n",
    "    spec = model_specs[model_name]\n",
    "    if spec is None:\n",
    "        return OpenAI()\n",
    "    \n",
    "    return OpenAI(api_key=spec[\"api_key\"], base_url=spec[\"url\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "orchestrator_model = \"gemini-2.0-flash\"\n",
    "generator = create_model(orchestrator_model)\n",
    "router = create_model(orchestrator_model)\n",
    "\n",
    "qa_models = {\n",
    "    model_name : create_model(model_name) \n",
    "    for model_name in model_specs.keys()\n",
    "}\n",
    "\n",
    "judges = {\n",
    "    model_name : create_model(model_name) \n",
    "    for model_name, specs in model_specs.items() \n",
    "    if not(specs) or specs[\"api_key\"] != \"ollama\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Orchestrator-Worker Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we generate a question to evaluate the intelligence of each LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = \"Please come up with a challenging, nuanced question that I can ask a number of LLMs \"\n",
    "request += \"to evaluate and rank them based on their intelligence. \" \n",
    "request += \"Answer **only** with the question, no explanation or preamble.\"\n",
    "\n",
    "messages = [{\"role\": \"user\", \"content\": request}]\n",
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = generator.chat.completions.create(\n",
    "    model=orchestrator_model,\n",
    "    messages=messages,\n",
    ")\n",
    "eval_question = response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(eval_question))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Task Parallelization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, having the question and all the models instantiated it's time to see what each model has to say about the complex task it was given."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = [{\"role\": \"user\", \"content\": eval_question}]\n",
    "answers = []\n",
    "competitors = []\n",
    "\n",
    "for name, model in qa_models.items():\n",
    "    response = model.chat.completions.create(model=name, messages=question)\n",
    "    answer = response.choices[0].message.content\n",
    "    competitors.append(name)\n",
    "    answers.append(answer)\n",
    "\n",
    "answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "report = \"# Answer report for each of the 5 models\\n\\n\"\n",
    "report += \"\\n\\n\".join([f\"## **Model: {model}**\\n\\n{answer}\" for model, answer in zip(competitors, answers)])\n",
    "display(Markdown(report))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Synthetizer/Judge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Judge Agents ranks the LLM responses based on coherence and relevance to the evaluation prompt. Judges vote and the final LLM ranking is based on the aggregated ranking of all three judges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "together = \"\"\n",
    "for index, answer in enumerate(answers):\n",
    "    together += f\"# Response from competitor {index+1}\\n\\n\"\n",
    "    together += answer + \"\\n\\n\"\n",
    "\n",
    "together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "judge_prompt = f\"\"\"\n",
    "                You are judging a competition between {len(competitors)} LLM competitors.\n",
    "                Each model has been given this nuanced question to evaluate their intelligence:\n",
    "\n",
    "                {eval_question}\n",
    "\n",
    "                Your job is to evaluate each response for clarity and strength of argument, and rank them in order of best to worst.\n",
    "                Respond with JSON, and only JSON, with the following format:\n",
    "                {{\"results\": [\"best competitor number\", \"second best competitor number\", \"third best competitor number\", ...]}}\n",
    "                With 'best competitor number being ONLY the number', for instance:\n",
    "                {{\"results\": [\"5\", \"2\", \"4\", ...]}}\n",
    "                Here are the responses from each competitor:\n",
    "\n",
    "                {together}\n",
    "\n",
    "                Now respond with the JSON with the ranked order of the competitors, nothing else. Do NOT include MARKDOWN FORMATTING or CODE BLOCKS. ONLY the JSON\n",
    "                \"\"\"\n",
    "\n",
    "judge_messages = [{\"role\": \"user\", \"content\": judge_prompt}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import re\n",
    "\n",
    "N = len(competitors)\n",
    "scores = defaultdict(int)\n",
    "for judge_name, judge in judges.items():\n",
    "    response = judge.chat.completions.create(\n",
    "        model=judge_name,\n",
    "        messages=judge_messages,\n",
    "    )\n",
    "    response = response.choices[0].message.content\n",
    "    response_json = re.findall(r'\\{.*?\\}', response)[0]\n",
    "    results = json.loads(response_json)[\"results\"]\n",
    "    ranks = [int(result) for result in results]\n",
    "    print(f\"Judge {judge_name} ranking:\")\n",
    "    for i, c in enumerate(ranks):\n",
    "        model_name = competitors[c - 1]\n",
    "        print(f\"#{i+1} : {model_name}\")\n",
    "        scores[c - 1] += (N - i)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_indices = sorted(scores, key=scores.get)\n",
    "\n",
    "# Convert to model names\n",
    "ranked_model_names = [competitors[i] for i in sorted_indices]\n",
    "\n",
    "print(\"Final ranking from best to worst:\")\n",
    "for i, name in enumerate(ranked_model_names[::-1], 1):\n",
    "    print(f\"#{i}: {name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Routing Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define a routing agent responsible for classifying task complexity and delegating the prompt to the most appropriate model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_question_complexity(question: str, routing_agent, routing_model) -> int:\n",
    "    \"\"\"\n",
    "    Ask an LLM to classify the question complexity from 1 (easy) to 5 (very hard).\n",
    "    \"\"\"\n",
    "    prompt = f\"\"\"\n",
    "        You are a classifier responsible for assigning a complexity level to user questions, based on how difficult they would be for a language model to answer.\n",
    "\n",
    "        Please read the question below and assign a complexity score from 1 to 5:\n",
    "\n",
    "        - Level 1: Very simple factual or definitional question (e.g., “What is the capital of France?”)\n",
    "        - Level 2: Slightly more involved, requiring basic reasoning or comparison\n",
    "        - Level 3: Moderate complexity, requiring synthesis, context understanding, or multi-part answers\n",
    "        - Level 4: High complexity, requiring abstract thinking, ethical judgment, or creative generation\n",
    "        - Level 5: Extremely challenging, requiring deep reasoning, philosophical reflection, or long-term multi-step inference\n",
    "\n",
    "        Respond ONLY with a single integer between 1 and 5 that best reflects the complexity of the question.\n",
    "\n",
    "        Question:\n",
    "        {question}\n",
    "        \"\"\"\n",
    "\n",
    "    response = routing_agent.chat.completions.create(\n",
    "        model=routing_model,\n",
    "        messages=[{\"role\": \"user\", \"content\": prompt}]\n",
    "    )\n",
    "    try:\n",
    "        return int(response.choices[0].message.content.strip())\n",
    "    except Exception:\n",
    "        return 3  # default to medium complexity on error\n",
    "    \n",
    "def route_question_to_model(question: str, models_by_rank, classifier_model=router, model_name=orchestrator_model):\n",
    "    level = classify_question_complexity(question, classifier_model, model_name)\n",
    "    selected_model_name = models_by_rank[level - 1]\n",
    "    return selected_model_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "difficulty_prompts = [\n",
    "    \"Generate a very basic, factual question that a small or entry-level language model could answer easily. It should require no reasoning, just direct knowledge lookup.\",\n",
    "    \"Generate a slightly involved question that requires basic reasoning, comparison, or combining two known facts. Still within the grasp of small models but not purely factual.\",\n",
    "    \"Generate a moderately challenging question that requires some synthesis of ideas, multi-step reasoning, or contextual understanding. A mid-tier model should be able to answer it with effort.\",\n",
    "    \"Generate a difficult question involving abstract thinking, open-ended reasoning, or ethical tradeoffs. The question should challenge large models to produce thoughtful and coherent responses.\",\n",
    "    \"Generate an extremely complex and nuanced question that tests the limits of current language models. It should require deep reasoning, long-term planning, philosophy, or advanced multi-domain knowledge.\"\n",
    "]\n",
    "def generate_question(level, generator=generator, generator_model=orchestrator_model):\n",
    "    prompt = (\n",
    "        f\"{difficulty_prompts[level - 1]}\\n\"\n",
    "        \"Answer only with the question, no explanation.\"\n",
    "    )\n",
    "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    response = generator.chat.completions.create(\n",
    "        model=generator_model,  # or your planner model\n",
    "        messages=messages\n",
    "    )\n",
    "    \n",
    "    return response.choices[0].message.content\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing Routing Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, to test the routing workflow, we create a function that accepts a task complexity level and triggers the full routing process.\n",
    "\n",
    "*Note: A level-N prompt isn't always assigned to the Nth-most capable model due to the classifier's subjective decisions.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_generation_routing(level):\n",
    "    question = generate_question(level=level)\n",
    "    answer_model = route_question_to_model(question, ranked_model_names)\n",
    "    messages = [{\"role\": \"user\", \"content\": question}]\n",
    "\n",
    "    response =qa_models[answer_model].chat.completions.create(\n",
    "            model=answer_model,  # or your planner model\n",
    "            messages=messages\n",
    "        )\n",
    "    print(f\"Question : {question}\")\n",
    "    print(f\"Routed to {answer_model}\")\n",
    "    display(Markdown(response.choices[0].message.content))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_generation_routing(level=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_generation_routing(level=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_generation_routing(level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_generation_routing(level=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_generation_routing(level=5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}