Extract resource-aware optimization logic into a tested module

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.env
+.pytest_cache/

notebooks/resource_aware_optimization.py

@@ -0,0 +1,186 @@
+import os
+import requests
+import json
+from dotenv import load_dotenv
+from openai import OpenAI
+
+
+# Load environment variables
+load_dotenv()
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+GOOGLE_CUSTOM_SEARCH_API_KEY = os.getenv("GOOGLE_CUSTOM_SEARCH_API_KEY")
+GOOGLE_CSE_ID = os.getenv("GOOGLE_CSE_ID")
+
+# Only raise ValueError if running as main
+if __name__ == "__main__":
+    if not OPENAI_API_KEY or not GOOGLE_CUSTOM_SEARCH_API_KEY or not GOOGLE_CSE_ID:
+        raise ValueError(
+            "Please set OPENAI_API_KEY, GOOGLE_CUSTOM_SEARCH_API_KEY, and GOOGLE_CSE_ID in your .env file."
+        )
+
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+
+def _safe_json_parse(text: str) -> dict:
+    """
+    Safely parse JSON from LLM response, handling markdown code blocks and leading/trailing text.
+    """
+    if not isinstance(text, str):
+        return None
+
+    try:
+        # 1. Try direct parsing
+        return json.loads(text.strip())
+    except (json.JSONDecodeError, AttributeError):
+        pass
+
+    # 2. Try to extract JSON from markdown blocks
+    # Use regex to find content between ```json and ``` or ``` and ```
+    import re
+    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1).strip())
+        except (json.JSONDecodeError, AttributeError):
+            pass
+
+    # 3. Fallback: try to find anything that looks like a JSON object
+    match = re.search(r"(\{.*?\})", text, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1).strip())
+        except (json.JSONDecodeError, AttributeError):
+            pass
+
+    return None
+
+
+# --- Step 1: Classify the Prompt ---
+def classify_prompt(prompt: str) -> dict:
+    system_message = {
+        "role": "system",
+        "content": (
+            "You are a classifier that analyzes user prompts and returns one of three categories ONLY:\n\n"
+            "- simple\n"
+            "- reasoning\n"
+            "- internet_search\n\n"
+            "Rules:\n"
+            "- Use 'simple' for direct factual questions that need no reasoning or current events.\n"
+            "- Use 'reasoning' for logic, math, or multi-step inference questions.\n"
+            "- Use 'internet_search' if the prompt refers to current events, recent data, or things not in your training data.\n\n"
+            "Respond ONLY with JSON like:\n"
+            '{ "classification": "simple" }'
+        ),
+    }
+
+    user_message = {"role": "user", "content": prompt}
+
+    # Added temperature=1 to match the pattern in other parts of the codebase
+    response = client.chat.completions.create(
+        model="gpt-4o", messages=[system_message, user_message], temperature=1
+    )
+
+    reply = response.choices[0].message.content
+    parsed = _safe_json_parse(reply)
+    if parsed is None:
+        return {"classification": "simple"}
+    return parsed
+
+
+# --- Step 2: Google Search ---
+def google_search(query: str, num_results=1):
+    url = "https://www.googleapis.com/customsearch/v1"
+    params = {
+        "key": GOOGLE_CUSTOM_SEARCH_API_KEY,
+        "cx": GOOGLE_CSE_ID,
+        "q": query,
+        "num": num_results,
+    }
+
+    try:
+        response = requests.get(url, params=params, timeout=10)
+        response.raise_for_status()
+        results = response.json()
+
+        if "items" in results and results["items"]:
+            return [
+                {
+                    "title": item.get("title"),
+                    "snippet": item.get("snippet"),
+                    "link": item.get("link"),
+                }
+                for item in results["items"]
+            ]
+        else:
+            return []
+    except requests.exceptions.RequestException as e:
+        return {"error": str(e)}
+
+
+# --- Step 3: Generate Response ---
+def generate_response(prompt: str, classification: str, search_results=None) -> str:
+    messages = []
+    if classification == "simple":
+        model = "gpt-4o-mini"
+        messages.append({"role": "user", "content": prompt})
+    elif classification == "reasoning":
+        model = "o1-mini"
+        messages.append({"role": "user", "content": prompt})
+    elif classification == "internet_search":
+        model = "gpt-4o"
+        # Convert each search result dict to a readable string
+        if isinstance(search_results, list) and search_results:
+            search_context = "\n".join(
+                [
+                    f"Title: {item.get('title')}\nSnippet: {item.get('snippet')}\nLink: {item.get('link')}"
+                    for item in search_results
+                ]
+            )
+        elif isinstance(search_results, dict) and "error" in search_results:
+            search_context = f"Error during search: {search_results['error']}"
+        else:
+            search_context = "No search results found."
+
+        messages.append(
+            {
+                "role": "system",
+                "content": f"Use the following web results to answer the user query:\n\n{search_context}",
+            }
+        )
+        messages.append({"role": "user", "content": prompt})
+    else:
+        # Default or error case
+        model = "gpt-4o-mini"
+        messages.append({"role": "user", "content": prompt})
+
+    response = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        temperature=1,
+    )
+
+    return response.choices[0].message.content, model
+
+
+# --- Step 4: Combined Router ---
+def handle_prompt(prompt: str) -> dict:
+    classification_result = classify_prompt(prompt)
+    classification = classification_result["classification"]
+
+    search_results = None
+    if classification == "internet_search":
+        try:
+            search_results = google_search(prompt)
+        except Exception as e:
+            search_results = {"error": str(e)}
+
+    answer, model = generate_response(prompt, classification, search_results)
+    return {"classification": classification, "response": answer, "model": model}
+
+if __name__ == "__main__":
+    test_prompt = "What is the capital of Australia?"
+
+    result = handle_prompt(test_prompt)
+    print("🔍 Classification:", result["classification"])
+    print("🧠 Model Used:", result["model"])
+    print("🧠 Response:\n", result["response"])