Add similar-problem search and fix vector database pipeline

autosetter/cli.py

@@ -6,9 +6,10 @@
 Executes:
 1. Intake: Read statement image/PDF.
 2. Extraction: Extract and validate `problem.json` via Qwen vision model (1st model — UNCHANGED).
-3. Generation: Produce artifacts via Ollama (validator.cpp, generator.cpp, solution.cpp, etc.)
-4. Validation: Sandboxed compilation, sample verification, test case generation, and checker probing.
-5. Packaging: Assemble release package bundle with manifest.
+3. Similarity: Look up similar existing problems in the vector database (optional).
+4. Generation: Produce artifacts via Ollama (validator.cpp, generator.cpp, solution.cpp, etc.)
+5. Validation: Sandboxed compilation, sample verification, test case generation, and checker probing.
+6. Packaging: Assemble release package bundle with manifest.
 """
 
 from __future__ import annotations
@@ -18,9 +19,9 @@
 import logging
 import os
 import sys
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from autosetter.config import (
     DEFAULT_NUM_TESTS,
@@ -29,6 +30,8 @@
     DEFAULT_TEXT_MODEL,
     DEFAULT_VISION_MODEL,
     PROMPTS_DIR,
+    SIMILARITY_ENABLED,
+    SIMILARITY_TOP_K,
 )
 from autosetter.extractor import (
     JSONExtractionError,
@@ -41,6 +44,11 @@
 from autosetter.packager import Packager, PackagerError
 from autosetter.pipeline import PipelineError, TestPipeline, TestReport
 from autosetter.sandbox import SandboxError, SandboxLocalClient, ensure_testlib
+from autosetter.similarity import (
+    SimilaritySearchError,
+    find_similar_problems,
+    save_similar_problems,
+)
 from autosetter.vision import ImageParsingError
 
 logger = logging.getLogger(__name__)
@@ -62,6 +70,7 @@ class PipelineResult:
     package_dir: Path
     report: Optional[TestReport] = None
     validation_error: str = ""
+    similar_problems: List[Dict[str, Any]] = field(default_factory=list)
 
     @property
     def ready_for_release(self) -> bool:
@@ -95,6 +104,8 @@ def generate_from_image(
     out_dir: str | Path = DEFAULT_OUT_DIR,
     prompts_dir: str | Path = PROMPTS_DIR,
     progress_callback: Optional[Callable[[str], None]] = None,
+    similarity_check: bool = SIMILARITY_ENABLED,
+    similar_k: int = SIMILARITY_TOP_K,
 ) -> PipelineResult:
     """
     Execute the end-to-end AutoSetter problem packaging pipeline.
@@ -119,6 +130,10 @@ def generate_from_image(
         Directory containing prompt templates.
     progress_callback : Optional[Callable[[str], None]]
         Progress reporting callback.
+    similarity_check : bool
+        If True, look up similar existing problems in the vector database (Qdrant).
+    similar_k : int
+        Number of similar problems to return.
 
     Returns
     -------
@@ -163,6 +178,18 @@ def generate_from_image(
     except JSONExtractionError as exc:
         raise AutoSetterError(f"Failed to save problem.json: {exc}") from exc
 
+    # 3b. Similar-problem lookup in the vector database (optional, never fatal)
+    similar_problems: List[Dict[str, Any]] = []
+    if similarity_check:
+        logger_fn("Searching vector database for similar problems...")
+        try:
+            similar_problems = find_similar_problems(problem_data, k=similar_k)
+            save_similar_problems(similar_problems, output_root / "similar_problems.json")
+            for match in similar_problems:
+                logger_fn(f"  {match['score']:.4f}  {match['title']}  {match['url'] or 'N/A'}")
+        except SimilaritySearchError as exc:
+            logger_fn(f"⚠️  Similar-problem search skipped: {exc}")
+
     # 4. Generate Downstream Artifacts & Validation Loop
     test_report = None
     validation_error = ""
@@ -287,6 +314,7 @@ def generate_from_image(
         package_dir=package_dir,
         report=test_report,
         validation_error=validation_error,
+        similar_problems=similar_problems,
     )
 
 
@@ -334,6 +362,18 @@ def build_arg_parser() -> argparse.ArgumentParser:
         default=False,
         help="Skip sandbox compilation and validation stage.",
     )
+    parser.add_argument(
+        "--no-similarity",
+        action="store_true",
+        default=not SIMILARITY_ENABLED,
+        help="Skip the similar-problem search in the vector database.",
+    )
+    parser.add_argument(
+        "--similar-k",
+        type=int,
+        default=SIMILARITY_TOP_K,
+        help=f"Number of similar problems to retrieve (default: {SIMILARITY_TOP_K}).",
+    )
     parser.add_argument(
         "--out-dir",
         type=str,
@@ -369,6 +409,8 @@ def main(argv: Optional[List[str]] = None) -> int:
             num_tests=args.num_tests,
             skip_validation=args.skip_validation,
             out_dir=args.out_dir,
+            similarity_check=not args.no_similarity,
+            similar_k=args.similar_k,
         )
     except AutoSetterError as exc:
         print(f"Error: {exc}", file=sys.stderr)
@@ -379,6 +421,10 @@ def main(argv: Optional[List[str]] = None) -> int:
 
     print(f"\nArtifacts written to: {result.generated_dir}")
     print(f"Package assembled at: {result.package_dir}")
+    if result.similar_problems:
+        print("Similar existing problems:")
+        for match in result.similar_problems:
+            print(f"  {match['score']:.4f}  {match['title']}  {match['url'] or 'N/A'}")
 
     if result.ready_for_release:
         print(f"Ready for release — {result.summary}")

autosetter/config.py

@@ -65,6 +65,15 @@
 SUPPORTED_PDF_EXTENSIONS = {".pdf"}
 SUPPORTED_EXTENSIONS = SUPPORTED_RASTER_EXTENSIONS | SUPPORTED_PDF_EXTENSIONS
 
+# Similar-problem search (vector_database/ + Qdrant)
+# Must match the model and collection used to build the vector database.
+SIMILARITY_ENABLED = os.environ.get("AUTOSETTER_SIMILARITY", "1") != "0"
+SIMILARITY_TOP_K = int(os.environ.get("AUTOSETTER_SIMILARITY_K", "5"))
+QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333")
+QDRANT_COLLECTION = os.environ.get("QDRANT_COLLECTION", "competitive_programming_problems")
+QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY") or None
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
+
 # Codeforces Polygon API
 POLYGON_API_URL = os.environ.get("POLYGON_API_URL", "https://polygon.codeforces.com/api/")
 POLYGON_API_KEY = os.environ.get("POLYGON_API_KEY", "")

autosetter/similarity.py

@@ -0,0 +1,91 @@
+"""
+autosetter.similarity
+=====================
+Similar-problem lookup against the vector database built by `vector_database/`.
+
+After extraction, the problem.json is embedded the same way the stored problems
+were and searched in Qdrant, so the setter can see whether the problem (or a
+close variant) already exists, with links to the matches.
+
+This stage is optional: the heavy dependencies (sentence-transformers, torch,
+qdrant-client) live in `vector_database/requirements.txt`, and a missing
+dependency or unreachable Qdrant only skips the check.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, List
+
+from autosetter.config import (
+    EMBEDDING_MODEL,
+    PROJECT_ROOT,
+    QDRANT_API_KEY,
+    QDRANT_COLLECTION,
+    QDRANT_URL,
+    SIMILARITY_TOP_K,
+)
+
+
+class SimilaritySearchError(Exception):
+    """Raised when the similar-problem lookup cannot run."""
+
+
+def _as_text(value: Any) -> str:
+    if isinstance(value, list):
+        return "\n".join(str(v) for v in value)
+    return str(value or "").strip()
+
+
+def find_similar_problems(
+    problem_data: Dict[str, Any],
+    k: int = SIMILARITY_TOP_K,
+) -> List[Dict[str, Any]]:
+    """
+    Return the top-k stored problems most similar to `problem_data`.
+
+    Each match is a dict with score, title, source, source_id, difficulty, tags and url.
+    """
+    # vector_database/ sits at the repository root, next to autosetter/
+    if str(PROJECT_ROOT) not in sys.path:
+        sys.path.append(str(PROJECT_ROOT))
+    try:
+        from vector_database.processing.schema import Problem
+        from vector_database.search import ProblemSearcher
+    except ImportError as exc:
+        raise SimilaritySearchError(
+            f"vector database dependencies are not installed ({exc}); "
+            "run `pip install -r vector_database/requirements.txt`"
+        ) from exc
+
+    problem = Problem(
+        id="autosetter_query",
+        source="autosetter",
+        source_id="",
+        title=_as_text(problem_data.get("title")),
+        statement=_as_text(problem_data.get("story")) or None,
+        input_format=_as_text(problem_data.get("input_format")) or None,
+        output_format=_as_text(problem_data.get("output_format")) or None,
+        constraints=_as_text(problem_data.get("constraints")) or None,
+    )
+
+    try:
+        searcher = ProblemSearcher(
+            url=QDRANT_URL,
+            collection_name=QDRANT_COLLECTION,
+            model_name=EMBEDDING_MODEL,
+            api_key=QDRANT_API_KEY,
+        )
+        return searcher.search_problem(problem, k=k)
+    except Exception as exc:
+        raise SimilaritySearchError(f"search against Qdrant at {QDRANT_URL} failed: {exc}") from exc
+
+
+def save_similar_problems(matches: List[Dict[str, Any]], output_path: str | Path) -> Path:
+    """Write the matches to a JSON file."""
+    path = Path(output_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(matches, indent=2, ensure_ascii=False), encoding="utf-8")
+    return path

tests/test_cli.py

@@ -85,6 +85,7 @@ def test_generate_from_image_end_to_end(tmp_path: Path, monkeypatch, stub_client
         image_path=img_path,
         out_dir=out_dir,
         skip_validation=True,
+        similarity_check=False,
     )
 
     assert result.generated_dir == out_dir / "generated"

tests/test_similarity.py

@@ -0,0 +1,91 @@
+"""
+Unit tests for the similar-problem lookup (autosetter.similarity) and its
+wiring into the CLI pipeline. Qdrant and the embedding model are stubbed.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from PIL import Image
+
+from autosetter.cli import build_arg_parser, generate_from_image
+from autosetter.similarity import SimilaritySearchError, save_similar_problems
+from tests.test_cli import SAMPLE_PROBLEM
+
+MATCHES = [
+    {
+        "score": 0.93,
+        "title": "Theatre Square",
+        "source": "codeforces",
+        "source_id": "1A",
+        "difficulty": "1000",
+        "tags": ["math"],
+        "url": "https://codeforces.com/problemset/problem/1/A",
+    }
+]
+
+
+def _run_pipeline(tmp_path: Path, monkeypatch, stub_client, **kwargs):
+    img_path = tmp_path / "problem.png"
+    Image.new("RGB", (50, 50), color="white").save(img_path, format="PNG")
+    problem_reply = "```json\n" + json.dumps(SAMPLE_PROBLEM) + "\n```"
+    client = stub_client(
+        replies={
+            "OCR specialist": problem_reply,
+            "competitive programming problem statement": problem_reply,
+        },
+        default="```cpp\nint main() { return 0; }\n```",
+    )
+    monkeypatch.setattr("autosetter.cli.OllamaClient", lambda **_: client)
+    return generate_from_image(
+        image_path=img_path, out_dir=tmp_path / "out", skip_validation=True, **kwargs
+    )
+
+
+def test_arg_parser_similarity_flags():
+    args = build_arg_parser().parse_args(["p.png", "--no-similarity", "--similar-k", "3"])
+    assert args.no_similarity is True
+    assert args.similar_k == 3
+
+
+def test_save_similar_problems(tmp_path: Path):
+    path = save_similar_problems(MATCHES, tmp_path / "similar_problems.json")
+    assert json.loads(path.read_text()) == MATCHES
+
+
+def test_pipeline_saves_similar_problems(tmp_path: Path, monkeypatch, stub_client):
+    seen = {}
+
+    def fake_find(problem_data, k):
+        seen["title"], seen["k"] = problem_data["title"], k
+        return MATCHES
+
+    monkeypatch.setattr("autosetter.cli.find_similar_problems", fake_find)
+    result = _run_pipeline(tmp_path, monkeypatch, stub_client, similar_k=3)
+
+    assert seen == {"title": SAMPLE_PROBLEM["title"], "k": 3}
+    assert result.similar_problems == MATCHES
+    saved = json.loads((tmp_path / "out" / "similar_problems.json").read_text())
+    assert saved[0]["url"] == "https://codeforces.com/problemset/problem/1/A"
+
+
+def test_pipeline_continues_when_search_fails(tmp_path: Path, monkeypatch, stub_client):
+    def failing_find(problem_data, k):
+        raise SimilaritySearchError("Qdrant unreachable")
+
+    monkeypatch.setattr("autosetter.cli.find_similar_problems", failing_find)
+    result = _run_pipeline(tmp_path, monkeypatch, stub_client)
+
+    assert result.similar_problems == []
+    assert (tmp_path / "out" / "package" / "manifest.json").exists()
+
+
+def test_pipeline_skips_search_when_disabled(tmp_path: Path, monkeypatch, stub_client):
+    def should_not_run(problem_data, k):
+        raise AssertionError("search should be skipped")
+
+    monkeypatch.setattr("autosetter.cli.find_similar_problems", should_not_run)
+    result = _run_pipeline(tmp_path, monkeypatch, stub_client, similarity_check=False)
+    assert result.similar_problems == []

vector_database/.gitignore

@@ -1,3 +1,4 @@
 __pycache__
 dataset/
 .env
+qdrant_storage/

vector_database/config.py

@@ -32,7 +32,7 @@
 QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
 QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "competitive_programming_problems")
 QDRANT_BATCH_SIZE = int(os.getenv("QDRANT_BATCH_SIZE", "64"))
-QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None)
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY") or None
 
 STATE_FILE = DATASET_DIR / "state.json"
 STATS_FILE = DATASET_DIR / "stats.json"

vector_database/embeddings/embedder.py

@@ -38,11 +38,12 @@ def build_embedding_text(problem: Problem) -> str:
 
 
 class ProblemEmbedder:
-    def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
+    def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5", batch_size: int = 32):
         if SentenceTransformer is None:
             raise ImportError("sentence-transformers is not installed. Please install it.")
 
         self.model_name = model_name
+        self.batch_size = batch_size
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading embedding model {model_name} on {self.device}...")
         self.model = SentenceTransformer(model_name, device=self.device)
@@ -59,6 +60,8 @@ def embed_batch(self, problems: List[Problem]) -> List[List[float]]:
         texts = [build_embedding_text(p) for p in problems]
 
         # We normalize embeddings for cosine similarity
-        embeddings = self.model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
+        embeddings = self.model.encode(
+            texts, batch_size=self.batch_size, normalize_embeddings=True, show_progress_bar=False
+        )
 
         return embeddings.tolist()