import os import sys # Dodajemy projekt do PATH, aby importy działały sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from core.subscription.callbacks import TokenUsageCallback from core.subscription.tracker import get_used_tokens from langchain_core.outputs import LLMResult def test_token_callback(): user_id = "test_dev_user" print("--- TEST 1: Inkrementacja z Callbacku ---") # Symulacja pustego callbacka callback = TokenUsageCallback(user_id=user_id) # Rekonstrukcja fikcyjnego responsu z modelem # (LLMResult ma atrybut llm_output, który uzupełnia pole total_tokens za nas po API wywołaniu) fake_response = LLMResult( generations=[], llm_output={ "token_usage": { "total_tokens": 125, "prompt_tokens": 75, "completion_tokens": 50, } }, ) initial_tokens = get_used_tokens(user_id) print(f"Początkowa pula tokenów dla {user_id}: {initial_tokens}") # Wywołanie sztucznego hook'a callback.on_llm_end(fake_response) final_tokens = get_used_tokens(user_id) print(f"Końcowa pula tokenów w bazie dla {user_id}: {final_tokens}") if final_tokens == initial_tokens + 125: print( "[SUKCES] Callback poprawnie zrzucił użycie 125 tokenów do bazy PostgreSQL!" ) else: print("[BŁĄD] Różnica w zliczaniu tokenów.") if __name__ == "__main__": print("Rozpoczęcie modułu sprawdzania poprawności...") test_token_callback()