from datetime import date
from aiatlas.ids import kind_of, new_id, normalize_alias, slugify
from aiatlas.sdk.extract.dates import parse_date, parse_datetime
from aiatlas.sdk.extract.html import parse_html
from aiatlas.sdk.extract.markdown import parse_markdown
from aiatlas.sdk.extract.numbers import (
parse_active_params,
parse_context_length,
parse_money_per_mtok,
parse_param_count,
parse_percent,
)
from aiatlas.sdk.fetch import canonicalize_url
from aiatlas.services.search import compile_query
def test_ids_and_slugs():
mid = new_id("model")
assert mid.startswith("model_") and kind_of(mid) == "model"
assert slugify("Qwen3-235B-A22B") == "qwen3-235b-a22b"
assert slugify("GPT-4.1 mini") == "gpt-4.1-mini"
assert normalize_alias("Claude 3.5 Haiku") == normalize_alias("claude-3-5-haiku") == "claude35haiku"
def test_numbers():
assert parse_param_count("Qwen3-235B-A22B") == 235_000_000_000
assert parse_active_params("Qwen3-235B-A22B") == 22_000_000_000
assert parse_param_count("a 7.6 billion parameters model") == 7_600_000_000
assert parse_param_count("Llama 3.1 8B Instruct") == 8_000_000_000
assert parse_param_count("version 2.0") is None
assert parse_context_length("128K tokens") == 128_000
assert parse_context_length("1M") == 1_000_000
assert parse_context_length("200,000 tokens") == 200_000
assert parse_context_length("32768") == 32768
assert parse_money_per_mtok("$3.00 / 1M tokens") == 3.0
assert parse_money_per_mtok("$0.15/M") == 0.15
assert parse_money_per_mtok("$2 per 1K tokens") == 2000.0
assert parse_percent("72.4%") == 72.4
def test_dates():
assert parse_datetime("2026-07-24T10:00:00Z").year == 2026
assert parse_date("July 24, 2026") == (date(2026, 7, 24), "day")
assert parse_date("March 2025") == (date(2025, 3, 1), "month")
assert parse_date("2024") == (date(2024, 1, 1), "year")
assert parse_date("") == (None, "none")
def test_canonical_url():
assert canonicalize_url("https://Example.com/a/?utm_source=x&b=1#frag") == "https://example.com/a?b=1"
assert canonicalize_url("https://example.com/") == "https://example.com/"
def test_html_parse():
html = """T
Head
linkBody text
"""
doc = parse_html(html, "https://example.com/p")
assert doc.title == "T" and doc.description == "D" and doc.canonical == "https://example.com/c" and doc.lang == "en"
assert doc.json_ld[0]["name"] == "X"
assert doc.embedded_json["__NEXT_DATA__"]["props"]["models"][0]["id"] == "m1"
assert doc.tables[0]["headers"] == ["Model", "Ctx"] and doc.tables[0]["rows"] == [["A", "128K"]]
assert doc.links == [("https://example.com/x", "link")]
assert "Body text" in doc.text and "Head" in doc.text
def test_markdown_parse():
md = parse_markdown("---\ntitle: X\nlicense: apache-2.0\n---\n# H1\n\n| a | b |\n|---|---|\n| 1 | 2 |\n\nSee [l](https://e.com).\n\n## Sec\ntext")
assert md.front_matter == {"title": "X", "license": "apache-2.0"}
assert md.tables[0]["headers"] == ["a", "b"] and md.tables[0]["rows"] == [["1", "2"]]
assert md.links == [("https://e.com", "l")]
assert md.section("Sec") == "text"
def test_compile_query():
q = compile_query("open models released in 2026 with more than 100B parameters and 128k context")
assert q.entity_type == "model" and q.openness == "open" and q.year_from == 2026 and q.params_min == 100_000_000_000 and q.context_min == 128_000
q2 = compile_query("vision models by Mistral")
assert q2.entity_type == "model" and "image" in q2.modalities and q2.organization == "Mistral"
assert compile_query("nvidia").filters["residual"] == "nvidia"