From a01d9cb1cf47590921864696619c254570e8fc87 Mon Sep 17 00:00:00 2001 From: denis-samatov Date: Sun, 30 Aug 2026 23:23:53 +0700 Subject: [PATCH] Extract resource-aware optimization logic into a tested module Extracts the classify-route-respond logic from the resource-aware optimization notebook into notebooks/resource_aware_optimization.py, adds a comprehensive test suite (18 tests, all mocked -- no live API calls), and fixes several issues found along the way: - Remove a hardcoded fallback API key. - Fix a prompt-injection vulnerability. - Make JSON parsing of LLM responses robust to markdown code fences and leading/trailing text, with a safe fallback instead of raising. - Add a timeout to the Google Custom Search HTTP request. - Handle search errors in handle_prompt instead of letting them propagate uncaught. - Remove dead/commented-out code. Verified: all 18 tests pass (pytest, mocked openai/requests/dotenv, no network access needed). --- .gitignore | 5 + notebooks/resource_aware_optimization.py | 186 +++++++++ notebooks/test_resource_aware_optimization.py | 356 ++++++++++++++++++ 3 files changed, 547 insertions(+) create mode 100644 .gitignore create mode 100644 notebooks/resource_aware_optimization.py create mode 100644 notebooks/test_resource_aware_optimization.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..2345c9c2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.py[cod] +*$py.class +.env +.pytest_cache/ diff --git a/notebooks/resource_aware_optimization.py b/notebooks/resource_aware_optimization.py new file mode 100644 index 00000000..154dbc6f --- /dev/null +++ b/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"]) diff --git a/notebooks/test_resource_aware_optimization.py b/notebooks/test_resource_aware_optimization.py new file mode 100644 index 00000000..891053d4 --- /dev/null +++ b/notebooks/test_resource_aware_optimization.py @@ -0,0 +1,356 @@ +import unittest +from unittest.mock import MagicMock, patch +import json +import sys + +# Define mocks +mock_requests = MagicMock() + +# Define a mock exception class for requests.exceptions.RequestException +class MockRequestException(Exception): + pass + +mock_requests.exceptions.RequestException = MockRequestException +mock_dotenv = MagicMock() +mock_openai = MagicMock() + +class TestSafeJsonParse(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # Mock external dependencies before importing the module under test + cls.original_modules = { + 'requests': sys.modules.get('requests'), + 'dotenv': sys.modules.get('dotenv'), + 'openai': sys.modules.get('openai') + } + sys.modules['requests'] = mock_requests + sys.modules['dotenv'] = mock_dotenv + sys.modules['openai'] = mock_openai + + # Now it's safe to import the function + from notebooks.resource_aware_optimization import _safe_json_parse + cls._safe_json_parse = staticmethod(_safe_json_parse) + + @classmethod + def tearDownClass(cls): + # Restore original modules + for name, module in cls.original_modules.items(): + if module is None: + del sys.modules[name] + else: + sys.modules[name] = module + + def test_safe_json_parse_plain(self): + text = '{"key": "value"}' + result = self._safe_json_parse(text) + self.assertEqual(result, {"key": "value"}) + + def test_safe_json_parse_markdown_json(self): + text = '```json\n{"key": "value"}\n```' + result = self._safe_json_parse(text) + self.assertEqual(result, {"key": "value"}) + + def test_safe_json_parse_markdown_plain(self): + text = '```\n{"key": "value"}\n```' + result = self._safe_json_parse(text) + self.assertEqual(result, {"key": "value"}) + + def test_safe_json_parse_invalid_json(self): + text = 'invalid json' + result = self._safe_json_parse(text) + self.assertIsNone(result) + + def test_safe_json_parse_none(self): + result = self._safe_json_parse(None) + self.assertIsNone(result) + + def test_safe_json_parse_empty(self): + result = self._safe_json_parse("") + self.assertIsNone(result) + +class TestClassifyPrompt(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # Mock external dependencies before importing the module under test + cls.original_modules = { + 'requests': sys.modules.get('requests'), + 'dotenv': sys.modules.get('dotenv'), + 'openai': sys.modules.get('openai') + } + sys.modules['requests'] = mock_requests + sys.modules['dotenv'] = mock_dotenv + sys.modules['openai'] = mock_openai + + # Now it's safe to import the function + from notebooks.resource_aware_optimization import classify_prompt + cls.classify_prompt = staticmethod(classify_prompt) + + @classmethod + def tearDownClass(cls): + # Restore original modules + for name, module in cls.original_modules.items(): + if module is None: + del sys.modules[name] + else: + sys.modules[name] = module + + @patch('notebooks.resource_aware_optimization.client.chat.completions.create') + def test_classify_prompt_simple(self, mock_create): + # Setup mock response + mock_response = MagicMock() + mock_response.choices = [ + MagicMock(message=MagicMock(content='{"classification": "simple"}')) + ] + mock_create.return_value = mock_response + + # Call the function + prompt = "What is the capital of France?" + result = self.classify_prompt(prompt) + + # Assertions + self.assertEqual(result, {"classification": "simple"}) + mock_create.assert_called_once() + args, kwargs = mock_create.call_args + self.assertEqual(kwargs['model'], "gpt-4o") + self.assertEqual(kwargs['messages'][1]['content'], prompt) + + @patch('notebooks.resource_aware_optimization.client.chat.completions.create') + def test_classify_prompt_reasoning(self, mock_create): + # Setup mock response + mock_response = MagicMock() + mock_response.choices = [ + MagicMock(message=MagicMock(content='{"classification": "reasoning"}')) + ] + mock_create.return_value = mock_response + + # Call the function + prompt = "Solve for x: 2x + 5 = 15" + result = self.classify_prompt(prompt) + + # Assertions + self.assertEqual(result, {"classification": "reasoning"}) + mock_create.assert_called_once() + + @patch('notebooks.resource_aware_optimization.client.chat.completions.create') + def test_classify_prompt_internet_search(self, mock_create): + # Setup mock response + mock_response = MagicMock() + mock_response.choices = [ + MagicMock(message=MagicMock(content='{"classification": "internet_search"}')) + ] + mock_create.return_value = mock_response + + # Call the function + prompt = "Who won the Super Bowl in 2024?" + result = self.classify_prompt(prompt) + + # Assertions + self.assertEqual(result, {"classification": "internet_search"}) + mock_create.assert_called_once() + + @patch('notebooks.resource_aware_optimization.client.chat.completions.create') + def test_classify_prompt_invalid_json(self, mock_create): + # Setup mock response with invalid JSON + mock_response = MagicMock() + mock_response.choices = [ + MagicMock(message=MagicMock(content='invalid json')) + ] + mock_create.return_value = mock_response + + # Call the function and assert it returns fallback + prompt = "test prompt" + result = self.classify_prompt(prompt) + self.assertEqual(result, {"classification": "simple"}) + + @patch('notebooks.resource_aware_optimization.client.chat.completions.create') + def test_classify_prompt_markdown_json(self, mock_create): + # Setup mock response with markdown JSON + mock_response = MagicMock() + mock_response.choices = [ + MagicMock(message=MagicMock(content='```json\n{"classification": "reasoning"}\n```')) + ] + mock_create.return_value = mock_response + + # Call the function + prompt = "test prompt" + result = self.classify_prompt(prompt) + + # Assertions + self.assertEqual(result, {"classification": "reasoning"}) + +class TestGoogleSearch(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # Mock external dependencies before importing the module under test + cls.original_modules = { + 'requests': sys.modules.get('requests'), + 'dotenv': sys.modules.get('dotenv'), + 'openai': sys.modules.get('openai') + } + sys.modules['requests'] = mock_requests + sys.modules['dotenv'] = mock_dotenv + sys.modules['openai'] = mock_openai + + # Now it's safe to import the function + from notebooks.resource_aware_optimization import google_search + cls.google_search = staticmethod(google_search) + + @classmethod + def tearDownClass(cls): + # Restore original modules + for name, module in cls.original_modules.items(): + if module is None: + del sys.modules[name] + else: + sys.modules[name] = module + + def setUp(self): + mock_requests.get.reset_mock() + mock_requests.get.side_effect = None + mock_requests.get.return_value = MagicMock() + + def test_google_search_success(self): + # Setup mock response + mock_response = MagicMock() + mock_response.json.return_value = { + "items": [ + { + "title": "Test Title", + "snippet": "Test Snippet", + "link": "https://test.com" + } + ] + } + mock_requests.get.return_value = mock_response + + # Call the function + result = self.google_search("test query") + + # Assertions + expected = [ + { + "title": "Test Title", + "snippet": "Test Snippet", + "link": "https://test.com" + } + ] + self.assertEqual(result, expected) + mock_requests.get.assert_called_once() + + def test_google_search_no_items(self): + # Setup mock response with no items + mock_response = MagicMock() + mock_response.json.return_value = {} + mock_requests.get.return_value = mock_response + + # Call the function + result = self.google_search("test query") + + # Assertions + self.assertEqual(result, []) + mock_requests.get.assert_called_once() + + def test_google_search_exception(self): + # Setup mock to raise RequestException + mock_requests.get.side_effect = MockRequestException("Connection error") + + # Call the function + result = self.google_search("test query") + + # Assertions + self.assertEqual(result, {"error": "Connection error"}) + mock_requests.get.assert_called_once() + +class TestHandlePrompt(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # Mock external dependencies before importing the module under test + cls.original_modules = { + 'requests': sys.modules.get('requests'), + 'dotenv': sys.modules.get('dotenv'), + 'openai': sys.modules.get('openai') + } + sys.modules['requests'] = mock_requests + sys.modules['dotenv'] = mock_dotenv + sys.modules['openai'] = mock_openai + + # Now it's safe to import the function + from notebooks.resource_aware_optimization import handle_prompt + cls.handle_prompt = staticmethod(handle_prompt) + + @classmethod + def tearDownClass(cls): + # Restore original modules + for name, module in cls.original_modules.items(): + if module is None: + del sys.modules[name] + else: + sys.modules[name] = module + + @patch('notebooks.resource_aware_optimization.classify_prompt') + @patch('notebooks.resource_aware_optimization.generate_response') + def test_handle_prompt_simple(self, mock_generate, mock_classify): + mock_classify.return_value = {"classification": "simple"} + mock_generate.return_value = ("Simple answer", "gpt-4o-mini") + + result = self.handle_prompt("Simple question") + + self.assertEqual(result["classification"], "simple") + self.assertEqual(result["response"], "Simple answer") + self.assertEqual(result["model"], "gpt-4o-mini") + mock_generate.assert_called_once_with("Simple question", "simple", None) + + @patch('notebooks.resource_aware_optimization.classify_prompt') + @patch('notebooks.resource_aware_optimization.generate_response') + def test_handle_prompt_reasoning(self, mock_generate, mock_classify): + mock_classify.return_value = {"classification": "reasoning"} + mock_generate.return_value = ("Reasoning answer", "o1-mini") + + result = self.handle_prompt("Complex question") + + self.assertEqual(result["classification"], "reasoning") + self.assertEqual(result["response"], "Reasoning answer") + self.assertEqual(result["model"], "o1-mini") + mock_generate.assert_called_once_with("Complex question", "reasoning", None) + + @patch('notebooks.resource_aware_optimization.classify_prompt') + @patch('notebooks.resource_aware_optimization.google_search') + @patch('notebooks.resource_aware_optimization.generate_response') + def test_handle_prompt_internet_search_success(self, mock_generate, mock_search, mock_classify): + mock_classify.return_value = {"classification": "internet_search"} + search_results = [{"title": "Result", "snippet": "Snippet", "link": "url"}] + mock_search.return_value = search_results + mock_generate.return_value = ("Search answer", "gpt-4o") + + result = self.handle_prompt("Current event question") + + self.assertEqual(result["classification"], "internet_search") + self.assertEqual(result["response"], "Search answer") + self.assertEqual(result["model"], "gpt-4o") + mock_search.assert_called_once_with("Current event question") + mock_generate.assert_called_once_with("Current event question", "internet_search", search_results) + + @patch('notebooks.resource_aware_optimization.classify_prompt') + @patch('notebooks.resource_aware_optimization.google_search') + @patch('notebooks.resource_aware_optimization.generate_response') + def test_handle_prompt_internet_search_error(self, mock_generate, mock_search, mock_classify): + mock_classify.return_value = {"classification": "internet_search"} + search_error = {"error": "API Key Invalid"} + mock_search.return_value = search_error + mock_generate.return_value = ("Error-based answer", "gpt-4o") + + result = self.handle_prompt("Current event question") + + self.assertEqual(result["classification"], "internet_search") + self.assertEqual(result["response"], "Error-based answer") + self.assertEqual(result["model"], "gpt-4o") + mock_search.assert_called_once_with("Current event question") + # Here we check if the error was passed to generate_response + mock_generate.assert_called_once_with("Current event question", "internet_search", search_error) + +if __name__ == '__main__': + unittest.main()