agents/haystack_wrapper.py

"""
Stub Haystack wrapper that routes RAG/document tasks through Brain._call_llm().
"""

from typing import Dict, Any, Tuple, Optional

def _build_messages(brain, query, context, system_prefix=""):
    """Build messages and call LLM directly, bypassing orchestrator to avoid loops."""
    history = (context or {}).get("history", [])
    user_profile = (context or {}).get("profile", {})
    user_model = (context or {}).get("user_model")

    profile_context = brain._format_profile(user_profile) if hasattr(brain, '_format_profile') else ""
    context_snippets, sources, topic = brain._assemble_context(query) if hasattr(brain, '_assemble_context') else ([], [], query)
    system_content = brain._build_system(profile_context, context_snippets, user_model) if hasattr(brain, '_build_system') else ""

    if system_prefix:
        system_content = system_prefix + "\n\n" + system_content

    messages = [{"role": "system", "content": system_content}]
    if history:
        formatted = brain._format_history(history) if hasattr(brain, '_format_history') else []
        messages.extend(formatted)
    messages.append({"role": "user", "content": str(query)})

    return messages


class HaystackWrapper:
    """Stub: routes RAG/document tasks through Brain._call_llm()."""

    def __init__(self, brain=None):
        self.brain = brain

    def run(self, query: str, context: Dict[str, Any] = None) -> Optional[str]:
        if not self.brain:
            return f"[Haystack stub] RAG query: {query[:100]}..."
        prefix = (
            "You are a document analysis expert. Answer based on the provided context. "
            "If the answer isn't in the context, say so clearly.\n\n"
        )
        messages = _build_messages(self.brain, query, context, prefix)
        try:
            response = self.brain._call_llm(messages, stream=False)
            return response.choices[0].message.content
        except Exception as e:
            return f"[Haystack error: {e}]"