import random from pydantic import ValidationError from neo4j_driver import driver from pydantic_restrictions import AltReviewOut, AltCloseOut, AltInviteOut, SocraticEval, SocraticQuestion from bielik import llm, modelLanguage from prompts import ALT_CLOSE_SYSTEM, ALT_INVITE_SYSTEM, EVAL_SYSTEM, build_system_prompt_introduction_chapter_ellis_distortion, build_eval_user_prompt from helpful_functions import getQuestions, get_last_user_message, introduction_talk, create_interview, check_situation from classifier import predict_raw, predict_raw1 from guardian import check_input from src.helpful_functions import beliefs_check_function from state import ChatState from langgraph.graph import StateGraph, START, END def _short_context(state: ChatState) -> str: lines = [] if state.get("distortion"): lines.append(f"- Zniekształcenie: {state['distortion']}") if state.get("distortion_def"): lines.append(f"- Definicja: {state['distortion_def']}") if state.get("current_intention"): lines.append(f"- Intencja pytań: {state['current_intention']}") if state.get("cel"): lines.append(f"- Cel intencji: {state['cel']}") if state.get("wniosek"): lines.append(f"- Wniosek ktory musi być zawarty w stworzonej alternatywnej myśli: {state['wniosek']}") if state.get("socratic_question"): lines.append(f"- Ostatnie pytanie sokratejskie: {state['socratic_question']}") hist = state.get("messages_detect") or state.get("messages") or [] tail = hist[-6:] if len(hist) > 6 else hist if tail: lines.append("- Fragment rozmowy:") for m in tail: who = "U" if m["role"] == "user" else "A" lines.append(f" {who}: {m['content']}") return "\n".join(lines) if lines else "(brak kontekstu)" def altthought_invite(state: ChatState) -> str: user_prompt = f"""Kontekst: {_short_context(state)} Historia chatu {state["messages"]} Zadanie: Napisz krótką, życzliwą prośbę, by użytkownik spróbował sformułować myśl alternatywną.""" out = llm.with_structured_output(AltInviteOut).invoke([ {"role": "system", "content": ALT_INVITE_SYSTEM}, {"role": "user", "content": user_prompt}, ]) return out.assistant_message # ───────────────────────────────────────── # KROK 2: Review + feedback / poprawka (LLM) # ───────────────────────────────────────── def altthought_review(state: ChatState, user_sentence: str) -> AltReviewOut: query = "MATCH (i:Intent {name:$intent})-[:HAS_CONCLUSION]->(c:Conclusion) RETURN c.must_include AS wniosek, c.example AS przyklad;" records, _, _ = driver.execute_query( query, parameters_={"intent": state["current_intention"]}, ) wniosek = records[0]["wniosek"] przyklad = records[0]["przyklad"] ALT_REVIEW_SYSTEM = f""" Jesteś empatycznym asystentem CBT. Oceń zdanie alternatywnej myśli stworzone przez użytkownika. DANE WEJSCIOWE: Stworzona alternatywna myśl przez użytkownika: {user_sentence} Wniosek, który musi się pojawić w alternatywnej myśli {wniosek} Przykładowa myśl alternatywna {przyklad} Historia rozmowy {state["messages"]} ZADANIE: 1) Oceń ALT pod kątem obecności powyższego WNIOSKU. 2) Jeśli WNIOSEK jest obecny → is_ok = true. 3) Jeśli WNIOSEK jest nieobecny lub niejednoznaczny → is_ok = false i: - zidentyfikuj, czego brak (konkretne elementy), - podaj zwięzłe wskazówki JAK użytkownik może ulepszyć ALT, aby zawierała WNIOSEK. - Nie proponuj gotowej treści; formułuj wskazówki („Zachęcam, abyś…” / „Dodaj…” / „Doprecyzuj…”). 4) Możesz odwołać się do poprzednich wypowiedzi i/lub zniekształcenia, ale nie podawaj przykładowej nowej myśli. Zwróć WYŁĄCZNIE JSON zgodny z AltReviewOut. """ out = llm.with_structured_output(AltReviewOut).invoke([ {"role": "system", "content": ALT_REVIEW_SYSTEM} ]) return out # ───────────────────────────────────────── # KROK 3: Komunikat końcowy (LLM) # ───────────────────────────────────────── def altthought_close(state: ChatState, final_sentence: str) -> str: user_prompt = f"""Kontekst: {_short_context(state)} Zatwierdzone zdanie AltThought: "{final_sentence}" """ out = llm.with_structured_output(AltCloseOut).invoke([ {"role": "system", "content": ALT_CLOSE_SYSTEM}, {"role": "user", "content": user_prompt}, ]) return out.assistant_message def call_eval_llm(state: ChatState, intent_name, messages): print(intent_name) classifier_llm = llm.with_structured_output(SocraticEval) result = classifier_llm.invoke([ { "role": "system", "content": EVAL_SYSTEM, }, {"role": "user", "content": build_eval_user_prompt(state, intent_name, messages)} ]) return result.cue_hit, result.route, result.explanation, result.proposition def detect_distortion(state: ChatState): if not state.get("messages"): print("Siema") state["messages"] = [{ "role": "assistant", "content": "Cześć! Cieszę się, że jesteś. Co u ciebie, czy masz jakiś problem? Z checią ci pomogę!" }] state["awaitingUser"] = True state["stage"] = "detect_distortion" return state else: state["first_stage_iterations"] += 1 print(state["first_stage_iterations"]) print("Siema1") last_message = get_last_user_message(state) user_text = (last_message["content"] or "").strip() if state["distortion"] is None: result = predict_raw(user_text) if result != "No Distortion": thought = beliefs_check_function(user_text) if thought: distortion = predict_raw1(user_text) print(distortion) state["distortion"] = distortion state["distortion_text"] = user_text print("Siema2") system_prompt = build_system_prompt_introduction_chapter_ellis_distortion(state["distortion"], state["situation"], state["think"], state["emotion"]) result = introduction_talk(state["messages"], system_prompt) if state["situation"] == "": state["situation"] = result.situation else: if result.situation != "": state["situation"] = create_interview(result.situation, state["situation"]) if state["emotion"] == "": state["emotion"] = result.emotion else: if result.emotion != "": state["emotion"] = create_interview(result.emotion, state["emotion"]) if state["think"] == "": state["think"] = result.think else: if result.think != "": state["think"] = create_interview(result.think, state["think"]) state["introduction_end_flag"] = result.chapter_end if state["distortion"] is not None and state["situation"] != "" and state["think"] != "" and state["emotion"] != "": print("Next") state["awaitingUser"] = False state["messages_detect"] = state["messages"] state["stage"] = "get_distortion_def" return state else: state["messages"].append({"role":"assistant", "content": result.model_output}) state["awaitingUser"] = True state["stage"] = "detect_distortion" return state def get_distortion_def(state: ChatState): print("Siema4") distortion = state["distortion"] query = """ MATCH (d:Distortion {name: $name}) RETURN d.definicja AS definicja """ records, summary, keys = driver.execute_query( query, parameters_={"name": distortion}, ) state["distortion_def"] = records[0]["definicja"] if records else None state["stage"] = "talk_about_distortion" state["awaitingUser"] = False return state def talk_about_distortion(state: ChatState): distortion = state["distortion"] distortion_def = state["distortion_def"] print("Siema5") if not state.get("distortion_explained"): print("Siema6") system_prompt_talk = f""" Jesteś empatycznym asystentem CBT. Użytkownikowi wykryto zniekształcenie poznawcze: Nazwa: {distortion} Definicja: {distortion_def} Przedstaw mu, że wykryłeś u niego zniekształcenie i wyjaśnij je w prosty, życzliwy sposób i zapytaj, czy chce, abyś pomógł mu to wspólnie przepracować. Język: polski, maksymalnie 2–3 zdania. """ llm_reply = llm.invoke([ { "role": "system", "content": system_prompt_talk, }, ]) follow_text = ( llm_reply if isinstance(llm_reply, str) else getattr(llm_reply, "content", str(llm_reply)) ) state["messages"].append({"role": "assistant", "content": follow_text}) state["awaitingUser"] = True state["stage"] = "talk_about_distortion" state["distortion_explained"] = True return state else: print("Siema7") last_user_msg = get_last_user_message(state) if not last_user_msg: state["awaitingUser"] = True return state classify_result = check_situation(last_user_msg["content"]) state["classify_result"] = classify_result if classify_result == "understand": print("Siema8") state["messages"].append({ "role": "assistant", "content": "Super! To przejdźmy teraz do kolejnego kroku" }) state["stage"] = "get_intention" state["awaitingUser"] = False return state # elif classify_result == "low_expression": # system_prompt = f""" # WEJSCIE # Historia wiadomości - {state["messages"]} # # Użytkownik jest mało wylewny i odpowiada krótko. # Twoim zadaniem jest napisać 2–3 empatyczne zdania po polsku, które spokojnie i nienachalnie zachęcą go do kontynuowania rozmowy. # Brzmi naturalnie, bez punktów, presji ani oceniania. # Na końcu zapytaj czy możemy możemy przejść do działania # Twoją rolą jest tylko i wyłącznie zachęcenie do działania nie pisz nic innego # """ # llm_reply = llm.invoke([ # { # "role": "system", # "content": system_prompt, # }, # ]) # follow_text = ( # llm_reply if isinstance(llm_reply, str) # else getattr(llm_reply, "content", str(llm_reply)) # ) # state["messages"].append({"role": "assistant", "content": follow_text}) # state["awaitingUser"] = True # state["stage"] = "talk_about_distortion" else: print("Siema9") system_prompt = f""" WEJSCIE Historia wiadomości - {state["messages"]} Użytkownik nie zrozumiał wyjaśnienia zniekształcenia. Nazwa: {distortion} Definicja: {distortion_def} Język tylko polski. Twoje zadanie: - Wyjaśnij prostszymi słowami (1–2 zdania). - Dodaj przykład z życia (1–2 zdania). - Zapytaj, czy teraz jest to jasne i czy możemy przejść do działania. Maksymalnie 3-4 zdania """ llm_reply = llm.invoke([ { "role": "system", "content": system_prompt, }, ]) follow_text = ( llm_reply if isinstance(llm_reply, str) else getattr(llm_reply, "content", str(llm_reply)) ) state["messages"].append({"role": "assistant", "content": follow_text}) state["awaitingUser"] = True state["stage"] = "talk_about_distortion" return state def get_intention(state: ChatState): distortion = state["distortion"] take_intent = """ MATCH (d:Distortion {name:$distortion})<-[:TARGETS]-(i:Intent) RETURN i.name AS nazwa ORDER BY nazwa """ records, summary, keys = driver.execute_query(take_intent, parameters_={"distortion": distortion}) result = [] for record in records: result.append(record["nazwa"]) state["priority_check"] = result state["stage"] = "select_intention" state["awaitingUser"] = False return state def select_intention(state: ChatState): state["messages_socratic"] = [] element = random.choice(state["priority_check"]) state["priority_check"].remove(element) state["current_intention"] = element state["question"] = 1 state["stage"] = "create_socratic_question" state["awaitingUser"] = False return state def create_socratic_question(state: ChatState): query = """ MATCH (i:Intent {name:$intencja}) RETURN i.name AS nazwa, i.aim AS cel, i.model_hint AS hint; """ records, _, _ = driver.execute_query( query, parameters_={"intencja":state["current_intention"]}, ) questions = getQuestions(records[0]["nazwa"]) socratic = state["messages_socratic"] if not socratic: creating_question_prompt = f""" Jesteś chatbotem terapeutycznym prowadzącym dialog sokratejski. ZADANIE: Wygeneruj DOKŁADNIE jedno krótkie pytanie po polsku. WEJŚCIE: - Zniekształcenie: {state["distortion"]} - Definicja: {state["distortion_def"]} - Błąd (cytat): {state["distortion_text"]} - Historia (P→U): {socratic} ← ostatnia odpowiedź to ostatnia linia zaczynająca się od „U:” - Intencja: {records[0]["nazwa"]} - Cel: {records[0]["cel"]} - Hint: {records[0]["hint"]} - Pytania referencyjne: {questions} REGUŁY: 1) Oprzyj pytanie przede wszystkim na hint + pytaniach referencyjnych. 2) Pytanie ma przybliżać do celu: {records[0]["cel"]}. 3) Nawiąż neutralnie do błędu „{state["distortion_text"]}”, eksplorując dowody/zakres/wyjątki/realistyczne alternatywy. Unikaj słowa „Dlaczego”. 4) Jedno pytanie; bez diagnoz, porad, definicji; bez kilku pytań naraz. 5) Nie powtarzaj dosłownie wcześniejszych pytań z {questions} ani pytań asystenta z {socratic}; parafrazuj i personalizuj wobec „{state["distortion_text"]}”. FORMAT WYJŚCIA: - Zwróć wyłącznie jedno zdanie zakończone „?” — bez cudzysłowów, markdown i etykiet; zero tekstu po „?”. AUTOKOREKTA: - Jeśli wygenerowano więcej niż jedno zdanie/linia, zwróć tylko pierwsze do pierwszego „?” włącznie. - Usuń frazy: "Wyjaśnienie:", "Explanation:", "Uzasadnienie:", "Dlaczego:", "Komentarz:". - Jeśli >140 znaków, skróć z zachowaniem sensu i „?” na końcu. """ else: creating_question_prompt = f""" Jesteś chatbotem terapeutycznym prowadzącym dialog sokratejski. ZADANIE: Wygeneruj DOKŁADNIE jedno krótkie pytanie po polsku. WEJŚCIE: - Zniekształcenie: {state["distortion"]} - Definicja: {state["distortion_def"]} - Błąd (cytat): {state["distortion_text"]} - Historia (P→U): {socratic} ← ostatnia odpowiedź to ostatnia linia zaczynająca się od „U:” - Intencja: {records[0]["nazwa"]} - Cel: {records[0]["cel"]} - Hint: {records[0]["hint"]} - Braki do celu: {state["decision_explanation"]} - Wskazówki do kolejnego pytania: {state["proposition"]} REGUŁY: 1) Oprzyj pytanie na ostatniej odpowiedzi użytkownika. Jeśli {state["decision_explanation"]} lub {state["proposition"]} nie są puste, wykorzystaj je do domknięcia brakujących informacji prowadzących do celu. 2) Pytanie ma przybliżać do celu: {records[0]["cel"]}. 3) Nawiąż neutralnie do błędu „{state["distortion_text"]}”, eksplorując dowody/zakres/wyjątki/alternatywy. Unikaj słowa „Dlaczego”. 4) Jedno pytanie; bez diagnoz, porad, definicji; bez kilku pytań naraz. 5) Nie powtarzaj dosłownie pytań z {questions} ani wcześniejszych pytań asystenta z {socratic}; parafrazuj i personalizuj wobec „{state["distortion_text"]}”. FORMAT WYJŚCIA: - Zwróć wyłącznie jedno zdanie zakończone „?” — bez cudzysłowów, markdown i etykiet; zero tekstu po „?”. AUTOKOREKTA: - Jeśli wygenerowano więcej niż jedno zdanie/linia, zwróć tylko pierwsze do pierwszego „?” włącznie. - Usuń frazy: "Wyjaśnienie:", "Explanation:", "Uzasadnienie:", "Dlaczego:", "Komentarz:". - Jeśli >140 znaków, skróć z zachowaniem sensu i „?” na końcu. """ question_llm = llm.with_structured_output(SocraticQuestion) result = question_llm.invoke([ { "role": "system", "content": creating_question_prompt, }, ]) state["messages"].append({"role":"assistant", "content": result.question}) state["messages_socratic"].append({"role": "assistant", "content": result.question}) state["stage"] = "analyze_output" state["awaitingUser"] = True return state def analyze_output(state: ChatState): state["messages_socratic"].append({"role": "user", "content": state["messages"][-1].get("content")}) cue_hit, confidence, explanation, proposition = call_eval_llm(state, state["current_intention"], state["messages_socratic"]) state["cue_hit"] = bool(cue_hit) state["confidence"] = confidence if cue_hit and confidence == "advance": state["stage"] = "enter_alt_thought" return state elif (not cue_hit) and confidence == "switch": state["decision_explanation"] = "" state["proposition"] = "" state["stage"] = "get_intention" return state else: state["stage"] = "create_socratic_question" state["decision_explanation"] = explanation state["proposition"] = proposition state["question"] = state.get("question") + 1 return state def validate_input(state: ChatState): stage = state.get("stage") if stage == "detect_distortion": chapter = "ETAP 1" elif stage == "talk_about_distortion" or stage == "get_distortion_def": chapter = "ETAP 2" elif stage == "create_socratic_question" or stage == "get_intention" or stage == "select_intention" or stage == "analyze_output": chapter = "ETAP 3" elif stage == "enter_alt_thought" or stage == "enter_alt_thought" or stage == "handle_alt_thought_input" or stage == "handle_alt_thought_input": chapter = "ETAP 4" else: chapter = "None" last_user_msg = state.get("last_user_msg_content") result = check_input(state["messages"], chapter, last_user_msg) state["last_user_msg"] = False if result.decision: state["validated"] = True state["awaitingUser"] = False else: state["noValidated"] = f"{chapter} - {last_user_msg}" state["explanation"] = result.explanation state["messages"].append({"role": "assistant", "content": result.message_to_user}) state["awaitingUser"] = True return state def enter_alt_thought(state: ChatState): result = altthought_invite(state) #TODO zmiana tekstu wprowadzającego do zrównoważonej myśli state.setdefault("messages", []).append({"role": "assistant", "content": result}) state["stage"] = "handle_alt_thought_input" state["awaitingUser"] = True return state def handle_alt_thought_input(state: ChatState): user_msg = next((m for m in reversed(state.get("messages", [])) if m["role"] == "user"), None) if not user_msg: state["awaitingUser"] = True return state user_sentence = (user_msg["content"] or "").strip() try: review = altthought_review(state, user_sentence) except ValidationError: msg = altthought_invite(state) state["messages"].append({"role": "assistant", "content": msg}) state["stage"] = "handle_alt_thought_input" state["awaitingUser"] = True return state if review.is_ok: final_sentence = user_sentence closing = altthought_close(state, final_sentence) state["messages"].append({"role": "assistant", "content": f"Zatwierdzona myśl: „{final_sentence}”"}) state["messages"].append({"role": "assistant", "content": closing}) state["stage"] = "end" state["awaitingUser"] = False return state else: state["messages"].append({"role": "assistant", "content": review.assistant_message}) state["stage"] = "handle_alt_thought_input" state["awaitingUser"] = True return state def global_router(state: ChatState) -> str: if state.get("awaitingUser"): print("[ROUTER] awaitingUser=True → __end__") return "__end__" stage = state.get("stage") print(f"[ROUTER] stage={stage} (fallback)") if not state.get("validated") and state.get("last_user_msg"): return "validate_input" if stage == "end": return "__end__" if stage == "get_distortion_def": return "get_distortion_def" if stage == "talk_about_distortion": return "talk_about_distortion" if stage == "get_intention": return "get_intention" # if stage == "get_socratic_question": # return "get_socratic_question" if stage == "select_intention": return "select_intention" if stage == "create_socratic_question": return "create_socratic_question" if stage == "analyze_output": return "analyze_output" if stage == "enter_alt_thought": return "enter_alt_thought" if stage == "handle_alt_thought_input": return "handle_alt_thought_input" print("[ROUTER] default → detect_distortion") return "detect_distortion" graph_builder = StateGraph(ChatState) graph_builder.add_node("detect_distortion", detect_distortion) graph_builder.add_node("get_distortion_def", get_distortion_def) graph_builder.add_node("talk_about_distortion", talk_about_distortion) graph_builder.add_node("get_intention", get_intention) graph_builder.add_node("select_intention", select_intention) # graph_builder.add_node("get_socratic_question", get_socratic_question) graph_builder.add_node("create_socratic_question", create_socratic_question) graph_builder.add_node("analyze_output", analyze_output) graph_builder.add_node("enter_alt_thought", enter_alt_thought) graph_builder.add_node("handle_alt_thought_input", handle_alt_thought_input) graph_builder.add_node("validate_input", validate_input) graph_builder.add_conditional_edges(START, global_router, { "detect_distortion": "detect_distortion", "get_distortion_def": "get_distortion_def", "talk_about_distortion": "talk_about_distortion", "get_intention": "get_intention", "select_intention": "select_intention", # "get_socratic_question": "get_socratic_question", "create_socratic_question": "create_socratic_question", "analyze_output": "analyze_output", "enter_alt_thought": "enter_alt_thought", "handle_alt_thought_input": "handle_alt_thought_input", "validate_input": "validate_input", "__end__": END, }) edge_map = { "detect_distortion": "detect_distortion", "get_distortion_def": "get_distortion_def", "talk_about_distortion": "talk_about_distortion", "get_intention": "get_intention", "select_intention": "select_intention", # "get_socratic_question": "get_socratic_question", "create_socratic_question": "create_socratic_question", "analyze_output": "analyze_output", "enter_alt_thought": "enter_alt_thought", "handle_alt_thought_input": "handle_alt_thought_input", "validate_input": "validate_input", "__end__": END, } for node in ["detect_distortion", "get_distortion_def","talk_about_distortion","get_intention","select_intention", "create_socratic_question", "analyze_output", "enter_alt_thought", "handle_alt_thought_input", "validate_input"]: graph_builder.add_conditional_edges(node, global_router, edge_map) graph = graph_builder.compile()