| import json |
| import time |
| import logging |
| from pathlib import Path |
| from typing import List, Dict, Any |
| from tqdm import tqdm |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class TextToMongoGenerator: |
| def __init__(self, pipeline, config: Any = None): |
| self.pipeline = pipeline |
| self.config = config |
|
|
| def _build_prompt(self, question: str, schema: str = "") -> str: |
| prompt = "### Task: Generate MongoDB Query (MQL)\n" |
| if schema: |
| prompt += f"### Collection Schema / Structure:\n{schema}\n\n" |
| prompt += f"### Question:\n{question}\n\n### MongoDB Query:\n db." |
| return prompt |
|
|
| def generate_single(self, question: str, schema: str = "") -> Dict[str, Any]: |
| """Generates a single MongoDB query and measures latency.""" |
| prompt = self._build_prompt(question, schema) |
| start_time = time.perf_counter() |
| raw_output = self.pipeline.generate(prompt) |
| end_time = time.perf_counter() |
| latency_ms = (end_time - start_time) * 1000 |
|
|
| |
| generated_query = self._clean_mongo_output(raw_output) |
|
|
| return { |
| "generated_query": generated_query, |
| "latency_ms": round(latency_ms, 2) |
| } |
|
|
| @staticmethod |
| def clean_mongo_output(raw_output: str) -> str: |
| stripped = raw_output.strip() |
| if stripped.lower().startswith("db."): |
| |
| restored = stripped |
| else: |
| restored = "db." + stripped |
|
|
| cleaned = ( |
| restored.replace("```javascript", "") |
| .replace("```json", "") |
| .replace("```js", "") |
| .replace("```", "") |
| .strip() |
| ) |
|
|
| |
| stop_markers = ["--->", "\n\n", "\nUser:", "\nQuestion:", "```", "###", "Human:"] |
| for marker in stop_markers: |
| idx = cleaned.find(marker) |
| if idx != -1: |
| cleaned = cleaned[:idx] |
|
|
| |
| lines = cleaned.split('\n') |
| for line in lines: |
| line = line.strip() |
| if line.startswith("db."): |
| return line.rstrip(';').strip() |
|
|
| |
| if "db." in cleaned: |
| idx = cleaned.find("db.") |
| end_idx = len(cleaned) |
| for marker in stop_markers + ["\n"]: |
| pos = cleaned.find(marker, idx) |
| if pos != -1: |
| end_idx = min(end_idx, pos) |
| return cleaned[idx:end_idx].strip().rstrip(';') |
|
|
| return cleaned.strip() |
|
|
| _clean_mongo_output = clean_mongo_output |
|
|
| def generate_batch(self, dataset: List[Dict[str, Any]], output_path: str = None) -> List[Dict[str, Any]]: |
| results = [] |
| logger.info(f"🚀 Starting batch MongoDB query generation for {len(dataset)} examples...") |
|
|
| for item in tqdm(dataset, desc="Generating MongoDB Queries", unit="query"): |
| query_id = item.get("query_id", f"sample_{len(results)}") |
| database_id = item.get("database_id", "unknown") |
| question = item.get("question", "") |
| schema = item.get("schema", "") |
| gold_query = item.get("generated_mongodb_query", item.get("gold_query", "")) |
| if isinstance(gold_query, list) and len(gold_query) > 0: |
| gold_query = gold_query[0] |
|
|
| if not question: |
| continue |
|
|
| gen_result = self.generate_single(question, schema) |
|
|
| result_entry = { |
| "query_id": query_id, |
| "database_id": database_id, |
| "question": question, |
| "gold_query": gold_query, |
| "generated_query": gen_result["generated_query"], |
| "latency_ms": gen_result["latency_ms"], |
| } |
| results.append(result_entry) |
|
|
| if output_path: |
| self.save_results(results, output_path) |
|
|
| return results |
|
|
| def save_results(self, results: List[Dict[str, Any]], output_path: str): |
| path = Path(output_path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| output_data = {"results": results} |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump(output_data, f, indent=4, ensure_ascii=False) |
|
|
| logger.info(f"✅ Successfully saved {len(results)} generated queries to {path}") |
|
|
