import types import app from chat_state import ConversationState from intent import CREATE_TABLE, EDIT_TABLE, SMALLTALK, SQL_QUERY, UNKNOWN, classify_intent def test_conversation_state_roundtrip_dict(): state = ConversationState( active_schema="CREATE TABLE zoologico (id INTEGER)", last_intent=SQL_QUERY, debug={"intent": SQL_QUERY, "confidence": 0.86, "reason": "sql_query_terms"}, ) restored = ConversationState.from_value(state.to_dict()) assert restored.active_schema == "CREATE TABLE zoologico (id INTEGER)" assert restored.last_intent == SQL_QUERY assert restored.debug["reason"] == "sql_query_terms" def test_intent_smalltalk_with_active_schema_is_not_sql(): state = ConversationState(active_schema="CREATE TABLE employees (id INTEGER)") result = classify_intent("como voce esta hoje?", state) assert result.intent == SMALLTALK def test_schema_request_without_columns_is_unknown_not_model_schema_task(): result = classify_intent("preciso de uma tabela sobre zoologico", ConversationState()) assert result.intent == UNKNOWN def test_intent_create_edit_and_sql_query(): empty_state = ConversationState() schema_state = ConversationState(active_schema="CREATE TABLE zoologico (id INTEGER, cidade TEXT)") create = classify_intent("crie tabela zoologico com id nome cidade", empty_state) edit = classify_intent("troca cidade por municipio", schema_state) query = classify_intent("liste zoologicos por municipio", schema_state) assert create.intent == CREATE_TABLE assert edit.intent == EDIT_TABLE assert query.intent == SQL_QUERY def test_destructive_data_mutation_is_unknown_with_active_schema(): state = ConversationState(active_schema="CREATE TABLE animals (id INTEGER, name TEXT, species TEXT)") assert classify_intent("delete all animals", state).intent == UNKNOWN assert classify_intent("DELETE FROM animals", state).intent == UNKNOWN assert classify_intent("UPDATE animals SET name = 'x'", state).intent == UNKNOWN assert classify_intent("INSERT INTO animals VALUES (1)", state).intent == UNKNOWN assert classify_intent("insert animal", state).intent == UNKNOWN assert classify_intent("insert row into animals", state).intent == UNKNOWN assert classify_intent("add row to animals", state).intent == UNKNOWN assert classify_intent("add animal record", state).intent == UNKNOWN assert classify_intent("drop animals table", state).intent == UNKNOWN assert classify_intent("drop the animals table", state).intent == UNKNOWN assert classify_intent("delete animals", state).intent == UNKNOWN assert classify_intent("drop animals", state).intent == UNKNOWN singular_state = ConversationState(active_schema="CREATE TABLE animal (id INTEGER, species TEXT)") assert classify_intent("delete animals", singular_state).intent == UNKNOWN assert classify_intent("drop animals", singular_state).intent == UNKNOWN def test_data_mutation_uses_schema_from_history_when_active_schema_empty(): history = [ { "role": "assistant", "content": "```sql\nCREATE TABLE animals (id INTEGER, species TEXT);\n```", } ] assert classify_intent("delete animals", ConversationState(), history).intent == UNKNOWN assert classify_intent("drop animals", ConversationState(), history).intent == UNKNOWN def test_delete_existing_column_stays_schema_edit(): state = ConversationState(active_schema="CREATE TABLE animals (id INTEGER, species TEXT)") assert classify_intent("delete species", state).intent == EDIT_TABLE assert classify_intent("drop coluna id", state).intent == EDIT_TABLE assert classify_intent("add habitat", state).intent == EDIT_TABLE def test_create_tables_named_rows_or_records_are_not_data_mutation(): assert classify_intent("create table rows with id name", ConversationState()).intent == CREATE_TABLE assert classify_intent("create table records with id name", ConversationState()).intent == CREATE_TABLE def test_zoologico_transcript_with_mocked_sql_model(monkeypatch): app._model = types.SimpleNamespace(generation_config=types.SimpleNamespace(eos_token_id=0)) app._tokenizer = types.SimpleNamespace(eos_token_id=0, pad_token_id=0) app._current_model_id = app.FINE_TUNED_MODEL_ID def fake_generate(prompt, generation_kind): assert generation_kind == app.model_core.SQL_GENERATION assert "CREATE TABLE zoologico" in prompt return "SELECT * FROM zoologico WHERE city = 'Sao Paulo';", 1 monkeypatch.setattr(app, "_generate_model_text", fake_generate) r1 = app.generate_response("oi", [], "", app.FINE_TUNED_MODEL_KEY) assert r1[4] == "" assert app.FALLBACK_RESPONSE in r1[0][-1]["content"] r2 = app.generate_response( "create table zoologico with id name city capacity", r1[0], r1[2], app.FINE_TUNED_MODEL_KEY, r1[7], ) assert "CREATE TABLE zoologico" in r2[4] assert "CREATE TABLE zoologico" in r2[2] r3 = app.generate_response( "change capacity to animal_count", r2[0], r2[2], app.FINE_TUNED_MODEL_KEY, r2[7], ) assert "animal_count TEXT" in r3[4] assert "capacidade" not in r3[4] r4 = app.generate_response( "list zoologicos from Sao Paulo", r3[0], r3[2], app.FINE_TUNED_MODEL_KEY, r3[7], ) assert "SELECT * FROM zoologico" in r4[4]