diff --git a/bot.py b/bot.py index 5863f30..66db8f7 100644 --- a/bot.py +++ b/bot.py @@ -440,19 +440,8 @@ class Bot: async def _respond_with_ai(self, room, user_message: str): model = self.room_models.get(room.room_id, DEFAULT_MODEL) - # Build conversation context from room timeline - messages = [{"role": "system", "content": SYSTEM_PROMPT}] - - # WildFiles document context - doc_results = await self.rag.search(user_message) - doc_context = self.rag.format_context(doc_results) - if doc_context: - logger.info("RAG found %d docs for: %s", len(doc_results), user_message[:50]) - messages.append({"role": "system", "content": doc_context}) - else: - logger.info("RAG found 0 docs for: %s", user_message[:50]) - - # Fetch last N messages from room via API + # Fetch conversation history FIRST (needed for query rewriting) + history = [] try: resp = await self.client.room_messages( room.room_id, start=self.client.next_batch or "", limit=10 @@ -462,10 +451,27 @@ class Bot: if not hasattr(evt, "body"): continue role = "assistant" if evt.sender == BOT_USER else "user" - messages.append({"role": role, "content": evt.body}) + history.append({"role": role, "content": evt.body}) except Exception: logger.debug("Could not fetch room history, proceeding without context") + # Rewrite query using conversation context for better RAG search + search_query = await self._rewrite_query(user_message, history, model) + + # WildFiles document context + doc_results = await self.rag.search(search_query) + doc_context = self.rag.format_context(doc_results) + if doc_context: + logger.info("RAG found %d docs for: %s (original: %s)", len(doc_results), search_query[:50], user_message[:50]) + else: + logger.info("RAG found 0 docs for: %s (original: %s)", search_query[:50], user_message[:50]) + + # Build conversation context + messages = [{"role": "system", "content": SYSTEM_PROMPT}] + if doc_context: + messages.append({"role": "system", "content": doc_context}) + messages.extend(history) + # Add current user message messages.append({"role": "user", "content": user_message}) @@ -485,6 +491,43 @@ class Bot: logger.exception("LLM call failed") await self._send_text(room.room_id, "Sorry, I couldn't generate a response.") + async def _rewrite_query(self, user_message: str, history: list[dict], model: str) -> str: + """Rewrite user message into a standalone search query using conversation context.""" + if not history or not self.llm: + return user_message + + # Build a compact history summary (last 4 messages max) + recent = history[-4:] + context_lines = [] + for msg in recent: + prefix = "User" if msg["role"] == "user" else "Assistant" + context_lines.append(f"{prefix}: {msg['content'][:200]}") + context_text = "\n".join(context_lines) + + try: + resp = await self.llm.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": ( + "You are a search query rewriter. Given conversation history and a new user message, " + "produce a single standalone search query that resolves all pronouns and references " + "(like 'this house', 'that document', 'it') using context from the conversation. " + "Reply with ONLY the rewritten search query in the same language as the user message. " + "No explanation, no quotes. If the message is already self-contained, return it as-is." + )}, + {"role": "user", "content": f"Conversation:\n{context_text}\n\nNew message: {user_message}"}, + ], + max_tokens=100, + ) + rewritten = resp.choices[0].message.content.strip().strip('"\'') + if rewritten and len(rewritten) < 500: + logger.info("Query rewritten: '%s' -> '%s'", user_message[:50], rewritten[:50]) + return rewritten + except Exception: + logger.debug("Query rewrite failed, using original", exc_info=True) + + return user_message + async def _auto_rename_room(self, room, user_message: str, ai_reply: str): """Generate a short topic title and set it as the room name.""" try: