Spaces:
Running
Running
File size: 3,561 Bytes
8429e5e ad8201c 8429e5e 3245c3b 8429e5e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Tests for gazet.lm — prompt templates, postprocessing, and GGUF helpers."""
from gazet.lm import (
_PLACES_SYSTEM_PROMPT,
_SYSTEM_PROMPT_TEMPLATE,
_USER_PROMPT_TEMPLATE,
_postprocess_sql,
)
from gazet.schemas import PlacesResult
class TestPostprocessSql:
def test_plain_sql(self):
assert _postprocess_sql("SELECT * FROM foo") == "SELECT * FROM foo"
def test_sql_fence_with_lang(self):
raw = "```sql\nSELECT id FROM bar\n```"
assert _postprocess_sql(raw) == "SELECT id FROM bar"
def test_sql_fence_no_lang(self):
raw = "```\nSELECT 1\n```"
assert _postprocess_sql(raw) == "SELECT 1"
def test_fence_with_trailing_text(self):
raw = "```sql\nSELECT * FROM foo\n```\nDone!"
result = _postprocess_sql(raw)
assert "SELECT" in result
assert "FROM" in result
def test_nested_backticks_preserved(self):
raw = "```sql\nSELECT `col`\n```"
assert _postprocess_sql(raw) == "SELECT `col`"
def test_whitespace_only(self):
assert _postprocess_sql(" ") == ""
def test_empty_string(self):
assert _postprocess_sql("") == ""
def test_multiline_query(self):
raw = """```sql
SELECT
id,
names."primary" AS name
FROM read_parquet('divisions_area')
WHERE country = 'IN'
```"""
result = _postprocess_sql(raw)
assert "SELECT" in result
assert "WHERE" in result
assert "```" not in result
def test_no_fence_just_whitespace_surround(self):
raw = " SELECT * FROM table "
assert _postprocess_sql(raw) == "SELECT * FROM table"
class TestPromptTemplates:
def test_system_prompt_has_schema_placeholder(self):
assert "{schema}" in _SYSTEM_PROMPT_TEMPLATE
def test_system_prompt_mentions_st_asgeojson(self):
assert "ST_AsGeoJSON" in _SYSTEM_PROMPT_TEMPLATE
def test_user_prompt_has_candidates_placeholder(self):
assert "{candidates_csv}" in _USER_PROMPT_TEMPLATE
def test_user_prompt_has_question_placeholder(self):
assert "{question}" in _USER_PROMPT_TEMPLATE
def test_system_prompt_formatting(self):
formatted = _SYSTEM_PROMPT_TEMPLATE.format(schema="dummy_schema")
assert "dummy_schema" in formatted
assert "{schema}" not in formatted
def test_user_prompt_formatting(self):
formatted = _USER_PROMPT_TEMPLATE.format(
candidates_csv="source,id,name",
question="get boundary",
)
assert "source,id,name" in formatted
assert "get boundary" in formatted
def test_places_system_prompt_has_examples(self):
assert "EXAMPLES:" in _PLACES_SYSTEM_PROMPT
assert "Puri, Odisha" in _PLACES_SYSTEM_PROMPT
assert "Amazon basin" in _PLACES_SYSTEM_PROMPT
class TestGeneratePlacesFallback:
"""Test that generate_places handles parse errors gracefully.
We can't test the real llama-server path without a running server,
but we can verify the PlacesResult model behavior used in fallback.
"""
def test_places_result_from_raw_query(self):
# Simulates the fallback path in generate_places
query = "get me the boundary of Paris"
result = PlacesResult(places=[{"place": query}])
assert len(result.places) == 1
assert result.places[0].place == query
def test_places_result_validation(self):
data = {"places": [{"place": "Chad"}, {"place": "Lake Chad"}]}
r = PlacesResult.model_validate(data)
assert len(r.places) == 2
|