Spaces:

cereeenn120
/

CharacterChatBot

Running

Refactor model call according to featherless ai

17ba573 1 day ago

10.6 kB

"""
Chat engine — core conversation loop.
Connects prompt_builder + memory + parser into a single streaming generator
that app.py (or any UI layer) can consume with `yield from`.
Separation of concerns:
    - app.py      : Gradio UI, OAuth token, slider values
    - chat.py     : LLM call, streaming, post-stream memory save
    - prompt_builder.py : system prompt assembly
    - memory.py   : session-scoped snapshots and periodic memory extraction
    - parser.py   : raw response → typed segments → markdown
"""

import logging
from huggingface_hub import InferenceClient
from engine.memory import sync_history, append_exchange, bounded_recent_context
from engine.memory_jobs import schedule_memory_update
from engine.relationship import relationship_to_prompt
from engine.parser import parse_response, segments_to_markdown
from engine.prompt_builder import build_messages, load_character
from engine.translator import translate_from_english, translate_to_english

logger = logging.getLogger(__name__)


def _chat_completion(client, messages, **kwargs):
    """Adapt chat completion calls for HF InferenceClient or OpenAI HF Router."""
    if isinstance(kwargs.get("model"), str) and kwargs["model"].endswith(":featherless-ai"):
        model = kwargs.pop("model")
        return client.chat.completions.create(model=model, messages=messages, **kwargs)
    return client.chat_completion(messages, **kwargs)


def generate_response(
    message:            str,
    character_id:       str,
    client:             InferenceClient,
    session_id:         str,
    history:            list[dict],
    persona:            dict | None = None,
    model:              str   = "Qwen/Qwen3-235B-A22B-Instruct-2507",
    max_tokens:         int   = 512,
    temperature:        float = 0.7,
    top_p:              float = 0.95,
    language:           str   = "English",
    translation_model:  str | None = None,
    use_translator:     bool  = False,
):
    """
    Core streaming generator.
    Yields the accumulated raw response after each token arrives so Gradio
    can display it in real time. After the stream ends, yields one final
    formatted version (actions italic, thoughts em-dashed) and saves
    the exchange to session-scoped memory.
    All errors are caught and surfaced as in-chat messages so the UI
    never silently breaks.
    Translation modes (when language ≠ English):
        use_translator=True  — dedicated translation_model translates user input
                               to English and response back to user's language.
                               Main model always operates in English.
        use_translator=False — no separate model. Main model receives the language
                               instruction directly and responds in user's language.
    Args:
        message:          Current user input.
        character_id:     Which character to respond as.
        client:           Initialised HuggingFace InferenceClient.
        persona:          Optional user persona dict.
        model:            HuggingFace model string.
        max_tokens:       Maximum tokens to generate.
        temperature:      Sampling temperature.
        top_p:            Nucleus sampling probability.
        language:         UI language selected by the user.
        translation_model: Model used when use_translator=True.
        use_translator:   True = separate translator, False = main model handles language.
    """
    
    
    logger.info(f"[Chat] ▶ character='{character_id}' | language='{language}' | model='{model}'")

    # ── Load character ─────────────────────────────────────────────────────────
    
    try:
        character = load_character(character_id)
        logger.debug(f"[Chat] Character loaded: '{character.get('name', character_id)}'")
    except FileNotFoundError:
        logger.error(f"[Chat] Character not found: {character_id}")
        yield f"*(Character '{character_id}' could not be loaded.)*"
        return
    except Exception as e:
        logger.error(f"[Chat] Unexpected error loading character: {e}", exc_info=True)
        yield "*(Something went wrong loading the character. Check the logs.)*"
        return

    # Synchronize the visible transcript before assembling model context.
    try:
        conversation = sync_history(session_id, character_id, history)
    except Exception:
        logger.exception("Could not synchronize conversation memory.")
        yield "*(Conversation history could not be loaded. Please try again.)*"
        return
    recent_messages, context_omitted = bounded_recent_context(conversation)
    summary = conversation['summary']
    long_term_memory = conversation['long_term_memory']

    # ── Translation mode? ──────────────────────────────────────────────────────
    is_non_english = language.strip().lower() != "english"
    use_translation = use_translator and is_non_english and bool(translation_model)
    translation_client = client
    if use_translation and isinstance(model, str) and model.endswith(":featherless-ai"):
        from huggingface_hub import InferenceClient
        translation_client = InferenceClient(token=getattr(client, "api_key", None), timeout=60)

    if is_non_english and not use_translator:
        logger.info(f"[Chat] Translation: OFF — main model will respond in '{language}' directly")
    elif use_translation:
        logger.info(f"[Chat] Translation: ON → model='{translation_model}'")
    else:
        logger.info("[Chat] Translation: not needed (language is English)")

    if use_translation:
        # Translate user message → English before sending to main model
        logger.debug(f"[Chat] Translating user message: {message[:80]!r}")
        en_message = translate_to_english(message, language, translation_client, translation_model)
        logger.debug(f"[Chat] User message (EN): {en_message[:80]!r}")
        prompt_language = "English"   # main model always responds in English
    else:
        en_message      = message
        prompt_language = language

    # ── Build prompt ───────────────────────────────────────────────────────────
    messages = build_messages(
        character=character,
        persona=persona,
        summary=summary,
        long_term_memory=long_term_memory,
        recent_messages=recent_messages,
        current_message=en_message,
        language=prompt_language,
    )
    messages[0]['content'] += '\n\n' + relationship_to_prompt(conversation['relationship'])
    if context_omitted:
        messages[0]['content'] += ('\n\nSome unsummarized history is temporarily omitted to keep context bounded. '
                                   'Do not invent details from missing turns.')
    if use_translation:
        messages[0]['content'] += "\n\nReply in English, even if earlier messages use another language."
    logger.debug(
        f"[Chat] Prompt built: {len(messages[0]['content'])} chars system | "
        f"{len(messages)} total messages"
    )

    # ── LLM call ──────────────────────────────────────────────────────────────
    logger.info(f"[Chat] LLM call: model='{model}' max_tokens={max_tokens} temp={temperature:.2f} top_p={top_p:.2f}")
    
    raw_response = ""

    try:
        if use_translation:
            # Collect silently — yield placeholder while model generates
            yield "*(…)*"
            for chunk in _chat_completion(
                client,
                messages,
                model=model,
                max_tokens=max_tokens,
                stream=True,
                temperature=temperature,
                top_p=top_p,
                frequency_penalty=0.3,
            ):
                if chunk.choices and chunk.choices[0].delta.content:
                    raw_response += chunk.choices[0].delta.content
        else:
            # Stream tokens in real time
            for chunk in _chat_completion(
                client,
                messages,
                model=model,
                max_tokens=max_tokens,
                stream=True,
                temperature=temperature,
                top_p=top_p,
                frequency_penalty=0.3,
            ):
                if chunk.choices and chunk.choices[0].delta.content:
                    raw_response += chunk.choices[0].delta.content
                yield raw_response

    except Exception as e:
        logger.error(f"[Chat] LLM error: {e}", exc_info=True)
        error_note = "\n\n*(The model did not respond. Please try again.)*"
        yield (raw_response + error_note) if raw_response else error_note.strip()
        return

    logger.info(f"[Chat] Stream complete: {len(raw_response)} chars generated")

    if not raw_response.strip():
        yield "*(The model returned an empty response. Please try again.)*"
        return

    # ── Translate response → user's language, format, yield ───────────────────
    if use_translation:
        logger.debug(f"[Chat] Translating response → {language} ({len(raw_response)} chars)")
        final_text = translate_from_english(raw_response, language, translation_client, translation_model)
    else:
        final_text = raw_response

    segments  = parse_response(final_text)
    logger.debug(f"[Chat] Parser: {len(segments)} segments → {[s['type'] for s in segments]}")
    formatted = segments_to_markdown(segments)
    logger.info(f"[Chat] ✓ Response ready: {len(formatted)} chars, {len(segments)} segments")
    # Store exactly what the user sees, including translated text.
    try:
        saved = append_exchange(session_id, character_id, conversation, message, formatted)
    except Exception:
        logger.exception("Could not save conversation memory.")
        yield formatted + "\n\n*(This response could not be saved. Please try again.)*"
        return
    if not saved:
        yield "*(Conversation changed while generating. Please try again.)*"
        return
    yield formatted
    schedule_memory_update(session_id, character_id, saved, client, model, persona=persona)

	"""
	Chat engine — core conversation loop.

	Connects prompt_builder + memory + parser into a single streaming generator
	that app.py (or any UI layer) can consume with `yield from`.

	Separation of concerns:
	- app.py : Gradio UI, OAuth token, slider values
	- chat.py : LLM call, streaming, post-stream memory save
	- prompt_builder.py : system prompt assembly
	- memory.py : session-scoped snapshots and periodic memory extraction
	- parser.py : raw response → typed segments → markdown
	"""

	import logging
	from huggingface_hub import InferenceClient
	from engine.memory import sync_history, append_exchange, bounded_recent_context
	from engine.memory_jobs import schedule_memory_update
	from engine.relationship import relationship_to_prompt
	from engine.parser import parse_response, segments_to_markdown
	from engine.prompt_builder import build_messages, load_character
	from engine.translator import translate_from_english, translate_to_english

	logger = logging.getLogger(__name__)


	def _chat_completion(client, messages, **kwargs):
	"""Adapt chat completion calls for HF InferenceClient or OpenAI HF Router."""
	if isinstance(kwargs.get("model"), str) and kwargs["model"].endswith(":featherless-ai"):
	model = kwargs.pop("model")
	return client.chat.completions.create(model=model, messages=messages, **kwargs)
	return client.chat_completion(messages, **kwargs)


	def generate_response(
	message: str,
	character_id: str,
	client: InferenceClient,
	session_id: str,
	history: list[dict],
	persona: dict \| None = None,
	model: str = "Qwen/Qwen3-235B-A22B-Instruct-2507",
	max_tokens: int = 512,
	temperature: float = 0.7,
	top_p: float = 0.95,
	language: str = "English",
	translation_model: str \| None = None,
	use_translator: bool = False,
	):
	"""
	Core streaming generator.

	Yields the accumulated raw response after each token arrives so Gradio
	can display it in real time. After the stream ends, yields one final
	formatted version (actions italic, thoughts em-dashed) and saves
	the exchange to session-scoped memory.

	All errors are caught and surfaced as in-chat messages so the UI
	never silently breaks.

	Translation modes (when language ≠ English):
	use_translator=True — dedicated translation_model translates user input
	to English and response back to user's language.
	Main model always operates in English.
	use_translator=False — no separate model. Main model receives the language
	instruction directly and responds in user's language.

	Args:
	message: Current user input.
	character_id: Which character to respond as.
	client: Initialised HuggingFace InferenceClient.
	persona: Optional user persona dict.
	model: HuggingFace model string.
	max_tokens: Maximum tokens to generate.
	temperature: Sampling temperature.
	top_p: Nucleus sampling probability.
	language: UI language selected by the user.
	translation_model: Model used when use_translator=True.
	use_translator: True = separate translator, False = main model handles language.
	"""


	logger.info(f"[Chat] ▶ character='{character_id}' \| language='{language}' \| model='{model}'")

	# ── Load character ─────────────────────────────────────────────────────────

	try:
	character = load_character(character_id)
	logger.debug(f"[Chat] Character loaded: '{character.get('name', character_id)}'")
	except FileNotFoundError:
	logger.error(f"[Chat] Character not found: {character_id}")
	yield f"(Character '{character_id}' could not be loaded.)"
	return
	except Exception as e:
	logger.error(f"[Chat] Unexpected error loading character: {e}", exc_info=True)
	yield "(Something went wrong loading the character. Check the logs.)"
	return

	# Synchronize the visible transcript before assembling model context.
	try:
	conversation = sync_history(session_id, character_id, history)
	except Exception:
	logger.exception("Could not synchronize conversation memory.")
	yield "(Conversation history could not be loaded. Please try again.)"
	return
	recent_messages, context_omitted = bounded_recent_context(conversation)
	summary = conversation['summary']
	long_term_memory = conversation['long_term_memory']

	# ── Translation mode? ──────────────────────────────────────────────────────
	is_non_english = language.strip().lower() != "english"
	use_translation = use_translator and is_non_english and bool(translation_model)
	translation_client = client
	if use_translation and isinstance(model, str) and model.endswith(":featherless-ai"):
	from huggingface_hub import InferenceClient
	translation_client = InferenceClient(token=getattr(client, "api_key", None), timeout=60)

	if is_non_english and not use_translator:
	logger.info(f"[Chat] Translation: OFF — main model will respond in '{language}' directly")
	elif use_translation:
	logger.info(f"[Chat] Translation: ON → model='{translation_model}'")
	else:
	logger.info("[Chat] Translation: not needed (language is English)")

	if use_translation:
	# Translate user message → English before sending to main model
	logger.debug(f"[Chat] Translating user message: {message[:80]!r}")
	en_message = translate_to_english(message, language, translation_client, translation_model)
	logger.debug(f"[Chat] User message (EN): {en_message[:80]!r}")
	prompt_language = "English" # main model always responds in English
	else:
	en_message = message
	prompt_language = language

	# ── Build prompt ───────────────────────────────────────────────────────────
	messages = build_messages(
	character=character,
	persona=persona,
	summary=summary,
	long_term_memory=long_term_memory,
	recent_messages=recent_messages,
	current_message=en_message,
	language=prompt_language,
	)
	messages[0]['content'] += '\n\n' + relationship_to_prompt(conversation['relationship'])
	if context_omitted:
	messages[0]['content'] += ('\n\nSome unsummarized history is temporarily omitted to keep context bounded. '
	'Do not invent details from missing turns.')
	if use_translation:
	messages[0]['content'] += "\n\nReply in English, even if earlier messages use another language."
	logger.debug(
	f"[Chat] Prompt built: {len(messages[0]['content'])} chars system \| "
	f"{len(messages)} total messages"
	)

	# ── LLM call ──────────────────────────────────────────────────────────────
	logger.info(f"[Chat] LLM call: model='{model}' max_tokens={max_tokens} temp={temperature:.2f} top_p={top_p:.2f}")

	raw_response = ""

	try:
	if use_translation:
	# Collect silently — yield placeholder while model generates
	yield "(…)"
	for chunk in _chat_completion(
	client,
	messages,
	model=model,
	max_tokens=max_tokens,
	stream=True,
	temperature=temperature,
	top_p=top_p,
	frequency_penalty=0.3,
	):
	if chunk.choices and chunk.choices[0].delta.content:
	raw_response += chunk.choices[0].delta.content
	else:
	# Stream tokens in real time
	for chunk in _chat_completion(
	client,
	messages,
	model=model,
	max_tokens=max_tokens,
	stream=True,
	temperature=temperature,
	top_p=top_p,
	frequency_penalty=0.3,
	):
	if chunk.choices and chunk.choices[0].delta.content:
	raw_response += chunk.choices[0].delta.content
	yield raw_response

	except Exception as e:
	logger.error(f"[Chat] LLM error: {e}", exc_info=True)
	error_note = "\n\n(The model did not respond. Please try again.)"
	yield (raw_response + error_note) if raw_response else error_note.strip()
	return

	logger.info(f"[Chat] Stream complete: {len(raw_response)} chars generated")

	if not raw_response.strip():
	yield "(The model returned an empty response. Please try again.)"
	return

	# ── Translate response → user's language, format, yield ───────────────────
	if use_translation:
	logger.debug(f"[Chat] Translating response → {language} ({len(raw_response)} chars)")
	final_text = translate_from_english(raw_response, language, translation_client, translation_model)
	else:
	final_text = raw_response

	segments = parse_response(final_text)
	logger.debug(f"[Chat] Parser: {len(segments)} segments → {[s['type'] for s in segments]}")
	formatted = segments_to_markdown(segments)
	logger.info(f"[Chat] ✓ Response ready: {len(formatted)} chars, {len(segments)} segments")
	# Store exactly what the user sees, including translated text.
	try:
	saved = append_exchange(session_id, character_id, conversation, message, formatted)
	except Exception:
	logger.exception("Could not save conversation memory.")
	yield formatted + "\n\n(This response could not be saved. Please try again.)"
	return
	if not saved:
	yield "(Conversation changed while generating. Please try again.)"
	return
	yield formatted
	schedule_memory_update(session_id, character_id, saved, client, model, persona=persona)