retriever/generator.py

#changed the prompt to output as markdown, plus some formating details
#also added get answer stream for incremental token rendering on the frontend
# --@Qamar


# called in get_answer_stream, to truncate to last full stop
class RAGGenerator:
    def truncate_incomplete_tail(self, answer: str) -> str:
        """Trim incomplete trailing text so responses end on a full stop."""
        if not answer:
            return answer

        trimmed = answer.rstrip()
        if not trimmed:
            return trimmed

        if trimmed.endswith("."):
            return trimmed

        last_full_stop = trimmed.rfind(".")
        if last_full_stop == -1:
            # If no sentence boundary exists, keep original text.
            return trimmed

        return trimmed[: last_full_stop + 1].rstrip()

    def generate_prompt(self, query, retrieved_contexts, context_urls=None):
        if context_urls:
            context_text = "\n\n".join([f"[Source {i+1}] {url}: {c}" for i, (c, url) in enumerate(zip(retrieved_contexts, context_urls))])
        else:
            context_text = "\n\n".join([f"[Source {i+1}]: {c}" for i, c in enumerate(retrieved_contexts)])
        
        return f"""You are an empathetic Cognitive Behavioral Therapy (CBT) therapist speaking directly to a client. **Your task is to provide a therapeutic, helpful response based ONLY on the provided clinical documents and excerpts**.

INSTRUCTIONS:
1. THERAPEUTIC DIALOGUE: Respond directly to the user as your client. Start by briefly validating their feelings, then gently apply CBT concepts, psychoeducation, or interventions found STRICTLY in the provided documents.
2. PATIENT EXAMPLES & NAMES (CRITICAL): The provided documents contain transcripts and examples of other patients and therapists (e.g., Abe, Judith, Joseph). These are illustrative case studies ONLY. DO NOT assume the user is "Abe" or any other person mentioned in the text. NEVER address or refer to the user by these names. Extract the CBT concepts/techniques demonstrated in these transcripts and apply them to the current user's unique situation.
3. GROUNDING (NO OPINIONS): Do not give your own opinions, general life advice, or use outside knowledge. Every therapeutic concept, identified cognitive distortion, or suggested exercise must come directly from the provided text.
4. FORMAT: Use clear Markdown formatting. Use paragraphs for conversational tone, and bullet points if you are breaking down specific steps, questions, or exercises found in the text.
5. MISSING INFO: If the provided excerpts do not contain relevant CBT concepts to address the client's specific statement, explicitly state: "While I hear how difficult this is for you, the clinical materials I have right now do not contain specific steps to address this." Do not invent therapeutic advice.

RETRIEVED CLINICAL CONTEXT:
{context_text}

CLIENT STATEMENT: {query}

THERAPEUTIC RESPONSE (GROUNDED IN SOURCES):"""

    def get_answer(self, model_instance, query, retrieved_contexts, context_urls=None, **kwargs):
        """Uses a specific model instance to generate the final answer."""
        prompt = self.generate_prompt(query, retrieved_contexts, context_urls)    
        answer = model_instance.generate(prompt, **kwargs)
        return self.truncate_incomplete_tail(answer)

    def get_answer_stream(self, model_instance, query, retrieved_contexts, context_urls=None, **kwargs):
        """Streams model output token-by-token for incremental UI updates."""
        prompt = self.generate_prompt(query, retrieved_contexts, context_urls)

        if hasattr(model_instance, "generate_stream"):
            for token in model_instance.generate_stream(prompt, **kwargs):
                if token:
                    yield token
            return

        # Fallback for model wrappers that only expose sync generation.
        answer = model_instance.generate(prompt, **kwargs)
        if answer:
            yield self.truncate_incomplete_tail(answer)