Spaces:
Sleeping
Sleeping
| from repograg.answer import ( | |
| AnswerGenerator, | |
| build_generator, | |
| extract_citations, | |
| validate_citations, | |
| ) | |
| from repograg.retrieval import RetrievalResult | |
| def make_result(file_path="src/elo.py", start=10, end=20, name="update_elo"): | |
| return RetrievalResult( | |
| chunk={ | |
| "id": 1, | |
| "file_path": file_path, | |
| "start_line": start, | |
| "end_line": end, | |
| "name": name, | |
| "parent": "", | |
| "kind": "function", | |
| "language": "python", | |
| "text": "def update_elo(): ...", | |
| "header": f"# {file_path}:{start}-{end}", | |
| "calls": "[]", | |
| }, | |
| score=1.0, | |
| sources=["bm25"], | |
| ) | |
| class StubResponse: | |
| def __init__(self, text): | |
| self.text = text | |
| class StubClient: | |
| """Mimics google-genai's Client: returns queued responses in order.""" | |
| def __init__(self, *texts): | |
| self._texts = list(texts) | |
| self.calls = 0 | |
| self.last_kwargs = None | |
| outer = self | |
| class Models: | |
| def generate_content(self, **kwargs): | |
| outer.calls += 1 | |
| outer.last_kwargs = kwargs | |
| return StubResponse(outer._texts.pop(0)) | |
| self.models = Models() | |
| def test_extract_citations(): | |
| cits = extract_citations( | |
| "Ratings update in update_elo [src/elo.py:10-20], " | |
| "using expected_score [src/elo.py:3]." | |
| ) | |
| assert len(cits) == 2 | |
| assert cits[0].file_path == "src/elo.py" | |
| assert (cits[0].start_line, cits[0].end_line) == (10, 20) | |
| assert (cits[1].start_line, cits[1].end_line) == (3, 3) | |
| def test_validate_citations_rejects_fabricated(): | |
| results = [make_result()] | |
| cits = extract_citations( | |
| "Real [src/elo.py:12-15]. Fabricated file [src/fake.py:1-5]. " | |
| "Out of range [src/elo.py:100-110]." | |
| ) | |
| valid = validate_citations(cits, results) | |
| assert len(valid) == 1 | |
| assert valid[0].file_path == "src/elo.py" | |
| def test_answer_with_valid_citations(): | |
| gen = AnswerGenerator(client=StubClient( | |
| "update_elo mutates the ratings dict [src/elo.py:10-20]." | |
| )) | |
| ans = gen.answer("how do ratings update", [make_result()]) | |
| assert ans.citations | |
| assert gen.client.calls == 1 | |
| def test_answer_retries_when_uncited(): | |
| gen = AnswerGenerator(client=StubClient( | |
| "Ratings update in place.", # no citations → retry | |
| "Ratings update in place [src/elo.py:10-20].", | |
| )) | |
| ans = gen.answer("how do ratings update", [make_result()]) | |
| assert gen.client.calls == 2 | |
| assert ans.citations | |
| assert "[src/elo.py:10-20]" in ans.text | |
| def test_answer_handles_blocked_response(): | |
| # google-genai returns text=None when the response is blocked/empty | |
| gen = AnswerGenerator(client=StubClient(None)) | |
| ans = gen.answer("how do ratings update", [make_result()]) | |
| assert ans.text == "" | |
| assert not ans.citations | |
| assert gen.client.calls == 1 # no retry on empty text | |
| def test_conversation_roles_mapped_for_gemini(): | |
| # the corrective retry sends the assistant turn back as role "model" | |
| stub = StubClient( | |
| "Uncited first answer.", | |
| "Cited [src/elo.py:10-20].", | |
| ) | |
| AnswerGenerator(client=stub).answer("q", [make_result()]) | |
| roles = [c["role"] for c in stub.last_kwargs["contents"]] | |
| assert roles == ["user", "model", "user"] | |
| def test_build_generator_disabled_without_key(monkeypatch): | |
| monkeypatch.delenv("GEMINI_API_KEY", raising=False) | |
| monkeypatch.delenv("GOOGLE_API_KEY", raising=False) | |
| assert build_generator() is None | |