kairo-code 0.1.0__py3-none-any.whl

kairo/backend/services/chat_service.py

@@ -0,0 +1,501 @@
+import json
+import logging
+from collections.abc import AsyncGenerator
+from datetime import date, datetime, timezone
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
+
+from backend.config import settings
+from backend.services.conversation_service import ConversationService
+from backend.services.llm_service import LLMService
+from backend.services.web_search_service import (
+    web_search,
+    format_search_results,
+)
+
+logger = logging.getLogger(__name__)
+
+WEB_SEARCH_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "web_search",
+        "description": "Search the web for current information, recent events, or to verify facts you are unsure about.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query",
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}
+TOOLS = [WEB_SEARCH_TOOL]
+
+
+def _estimate_tokens(text: str) -> int:
+    """Rough token estimate: ~4 chars per token for English text."""
+    return len(text) // 4
+
+
+def _strip_duplicate_trailing_code(text: str) -> str:
+    """Remove raw code duplicated at the end of a response.
+
+    Quantized models sometimes output a fenced code block and then repeat the
+    same code as unformatted text at the bottom. This detects the pattern and
+    strips the trailing duplicate.
+    """
+    import re as _re
+
+    # Find all fenced code blocks
+    fence_pattern = _re.compile(r"```[\w]*\n(.*?)```", _re.DOTALL)
+    blocks = fence_pattern.findall(text)
+    if not blocks:
+        return text
+
+    # Check if the text after the last fence contains a duplicate of any block
+    last_fence_end = text.rfind("```")
+    if last_fence_end < 0:
+        return text
+    # Move past the closing ```
+    tail_start = last_fence_end + 3
+    tail = text[tail_start:].strip()
+    if not tail or len(tail) < 20:
+        return text
+
+    for block in blocks:
+        block_stripped = block.strip()
+        if not block_stripped:
+            continue
+        # Check if the tail is substantially the same as a code block
+        # Use normalized comparison (collapse whitespace)
+        norm_block = " ".join(block_stripped.split())
+        norm_tail = " ".join(tail.split())
+        # Tail matches or is a prefix/suffix of the block
+        if norm_tail in norm_block or norm_block in norm_tail:
+            logger.info("Stripped duplicate trailing code (%d chars)", len(tail))
+            return text[:tail_start].rstrip()
+        # Also check with high overlap ratio for near-duplicates
+        if len(norm_tail) > 50 and len(norm_block) > 50:
+            shorter = min(len(norm_tail), len(norm_block))
+            # Compare the first N characters
+            if norm_tail[:shorter] == norm_block[:shorter]:
+                logger.info("Stripped near-duplicate trailing code (%d chars)", len(tail))
+                return text[:tail_start].rstrip()
+
+    return text
+
+
+class ChatService:
+    def __init__(self, db: AsyncSession, llm_service: LLMService, user_id: str | None = None):
+        self.db = db
+        self.conversation_service = ConversationService(db, user_id=user_id)
+        self.llm_service = llm_service
+        self.user_id = user_id
+
+    async def stream_response(
+        self,
+        message: str,
+        model: str,
+        conversation_id: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int = 2048,
+        project_id: str | None = None,
+    ) -> AsyncGenerator[str, None]:
+        # Check usage limits before generating
+        if self.user_id:
+            from backend.services.usage_service import UsageService
+            usage_service = UsageService(self.db)
+            allowed, reason = await usage_service.check_limits(self.user_id)
+            if not allowed:
+                yield _sse({"type": "error", "content": reason})
+                return
+
+        # Get or create conversation
+        if conversation_id:
+            conv = await self.conversation_service.get(conversation_id)
+            if not conv:
+                logger.warning("Conversation %s not found", conversation_id)
+                yield _sse({"type": "error", "content": "Conversation not found"})
+                return
+        else:
+            conv = await self.conversation_service.create(model=model, project_id=project_id)
+            logger.info("Created conversation %s with model %s", conv.id, model)
+
+        # Emit meta event with conversation ID
+        yield _sse({"type": "meta", "conversation_id": conv.id})
+
+        # Save user message
+        await self.conversation_service.add_message(conv.id, "user", message)
+
+        # Reload with messages (force fresh from DB)
+        conv = await self.conversation_service.get_fresh(conv.id)
+
+        # Check if we need to summarize older messages before building history
+        await self._maybe_summarize(conv, model)
+        # Reload after potential summary
+        conv = await self.conversation_service.get(conv.id)
+
+        # Build context-aware history
+        history = self._build_history(conv, model)
+
+        logger.info(
+            "Streaming conv=%s model=%s history_msgs=%d",
+            conv.id, model, len(history),
+        )
+
+        # Stream from LLM with tool calling support
+        full_response = ""
+        usage_data = None
+        try:
+            tool_calls_result = None
+            async for chunk in self.llm_service.stream_chat(
+                messages=history,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                tools=TOOLS,
+            ):
+                if isinstance(chunk, dict):
+                    if chunk.get("type") == "fallback":
+                        yield _sse({"type": "status", "content": "Using Nyx Lite — full model resuming shortly"})
+                    elif chunk.get("type") == "tool_calls":
+                        tool_calls_result = chunk["calls"]
+                    else:
+                        usage_data = chunk
+                else:
+                    full_response += chunk
+                    yield _sse({"type": "content", "content": chunk})
+
+            # If the model requested tool calls, execute them
+            if tool_calls_result:
+                for call in tool_calls_result:
+                    if call["name"] == "web_search":
+                        yield _sse({"type": "status", "content": "Searching the web..."})
+                        try:
+                            args = json.loads(call["arguments"])
+                        except json.JSONDecodeError:
+                            logger.warning("Invalid tool call arguments: %s", call["arguments"])
+                            continue
+                        results = await web_search(args.get("query", message))
+                        result_text = format_search_results(results)
+                        logger.info("Tool call web_search(%s) returned %d results", args.get("query"), len(results))
+
+                        # Append tool call + result to history for second LLM call
+                        history.append({
+                            "role": "assistant",
+                            "content": None,
+                            "tool_calls": [{
+                                "id": call["id"],
+                                "type": "function",
+                                "function": {
+                                    "name": call["name"],
+                                    "arguments": call["arguments"],
+                                },
+                            }],
+                        })
+                        history.append({
+                            "role": "tool",
+                            "tool_call_id": call["id"],
+                            "content": result_text or "No results found.",
+                        })
+
+                # Second LLM call — no tools, generate final response using search results
+                async for chunk in self.llm_service.stream_chat(
+                    messages=history,
+                    model=model,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                ):
+                    if isinstance(chunk, dict):
+                        usage_data = chunk
+                    else:
+                        full_response += chunk
+                        yield _sse({"type": "content", "content": chunk})
+
+        except Exception as e:
+            logger.error("LLM streaming error for conv %s: %s", conv.id, e, exc_info=True)
+            yield _sse({"type": "error", "content": "Failed to generate response. Please try again."})
+            return
+
+        # Post-process: strip duplicate trailing code blocks
+        full_response = _strip_duplicate_trailing_code(full_response)
+
+        # Save assistant response
+        if full_response:
+            await self.conversation_service.add_message(conv.id, "assistant", full_response)
+
+        # Record usage
+        if self.user_id and usage_data:
+            try:
+                from backend.services.usage_service import UsageService
+                usage_service = UsageService(self.db)
+                await usage_service.record_usage(
+                    user_id=self.user_id,
+                    conversation_id=conv.id,
+                    model=model,
+                    prompt_tokens=usage_data.get("prompt_tokens", 0),
+                    completion_tokens=usage_data.get("completion_tokens", 0),
+                )
+            except Exception as e:
+                logger.warning("Failed to record usage: %s", e)
+
+        # Auto-title from first exchange
+        if conv.title == "New Conversation" and full_response:
+            title = self._generate_title(message)
+            await self.conversation_service.rename(conv.id, title)
+
+        # Update timestamp
+        conv.updated_at = datetime.now(timezone.utc)
+        await self.db.commit()
+
+        yield "data: [DONE]\n\n"
+
+    def _build_history(self, conv, model: str) -> list[dict[str, str]]:
+        context_limit = settings.CONTEXT_LIMITS.get(model, 6000)
+        history: list[dict[str, str]] = []
+        token_count = 0
+
+        # System prompt with current date
+        system_prompt = settings.SYSTEM_PROMPTS.get(model, "")
+        if system_prompt:
+            today = date.today().strftime("%A, %B %d, %Y")
+            system_prompt += f"\n\nToday's date is {today}."
+            history.append({"role": "system", "content": system_prompt})
+            token_count += _estimate_tokens(system_prompt)
+
+        # If conversation belongs to a project with instructions, inject them
+        if conv.project and conv.project.instructions:
+            project_msg = f"[Project Instructions]\n{conv.project.instructions}"
+            history.append({"role": "system", "content": project_msg})
+            token_count += _estimate_tokens(project_msg)
+
+        # If we have a summary of older messages, inject it
+        if conv.summary:
+            summary_msg = (
+                f"[Context from earlier in this conversation]\n{conv.summary}"
+            )
+            history.append({"role": "system", "content": summary_msg})
+            token_count += _estimate_tokens(summary_msg)
+
+        # Add recent messages, working backwards to prioritize the latest
+        messages = list(conv.messages)
+        to_include = []
+        for msg in reversed(messages):
+            msg_tokens = _estimate_tokens(msg.content)
+            if token_count + msg_tokens > context_limit:
+                break
+            to_include.append({"role": msg.role, "content": msg.content})
+            token_count += msg_tokens
+
+        # Reverse back to chronological order
+        to_include.reverse()
+
+        # Ensure we at least include the very last message (the new user msg)
+        if not to_include and messages:
+            last = messages[-1]
+            to_include = [{"role": last.role, "content": last.content}]
+
+        history.extend(to_include)
+
+        logger.debug(
+            "Built history: %d msgs, ~%d tokens (limit %d)",
+            len(history), token_count, context_limit,
+        )
+        return history
+
+    async def _maybe_summarize(self, conv, model: str) -> None:
+        messages = list(conv.messages)
+        total_tokens = sum(_estimate_tokens(m.content) for m in messages)
+
+        if total_tokens < settings.SUMMARY_TRIGGER_TOKENS:
+            return
+
+        midpoint = len(messages) // 2
+        if midpoint < 2:
+            return
+
+        old_messages = messages[:midpoint]
+        old_text = "\n".join(
+            f"{m.role}: {m.content[:500]}" for m in old_messages
+        )
+
+        summary_prompt = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a summarization assistant. Condense the following "
+                    "conversation into a brief summary (2-4 sentences) that "
+                    "captures the key topics, decisions, and context. "
+                    "Focus on information the AI would need to continue "
+                    "the conversation coherently."
+                ),
+            },
+            {"role": "user", "content": f"Summarize this conversation:\n\n{old_text}"},
+        ]
+
+        logger.info(
+            "Summarizing %d old messages for conv %s (~%d tokens)",
+            len(old_messages), conv.id, sum(_estimate_tokens(m.content) for m in old_messages),
+        )
+
+        try:
+            summary = ""
+            async for token in self.llm_service.stream_chat(
+                messages=summary_prompt,
+                model=model,
+                temperature=0.3,
+                max_tokens=300,
+            ):
+                if isinstance(token, dict):
+                    continue
+                summary += token
+
+            if summary:
+                conv.summary = summary
+                await self.db.commit()
+                logger.info("Saved summary for conv %s: %s", conv.id, summary[:100])
+        except Exception as e:
+            logger.warning("Failed to summarize conv %s: %s", conv.id, e)
+
+    async def regenerate_response(
+        self,
+        conversation_id: str,
+        temperature: float = 0.7,
+        max_tokens: int = 2048,
+    ) -> AsyncGenerator[str, None]:
+        """Delete last assistant message and re-stream a new response."""
+        conv = await self.conversation_service.get(conversation_id)
+        if not conv:
+            yield _sse({"type": "error", "content": "Conversation not found"})
+            return
+
+        # Delete last assistant message
+        await self.conversation_service.delete_last_assistant_message(conversation_id)
+
+        # Reload conversation
+        conv = await self.conversation_service.get(conversation_id)
+        if not conv or not conv.messages:
+            yield _sse({"type": "error", "content": "No messages to regenerate from"})
+            return
+
+        model = conv.model
+        yield _sse({"type": "meta", "conversation_id": conv.id})
+
+        # Build history and stream
+        await self._maybe_summarize(conv, model)
+        conv = await self.conversation_service.get(conv.id)
+        history = self._build_history(conv, model)
+
+        # Find the last user message (for fallback query if tool args parse fails)
+        messages_sorted = sorted(conv.messages, key=lambda m: m.created_at)
+        last_user_msg = ""
+        for m in reversed(messages_sorted):
+            if m.role == "user":
+                last_user_msg = m.content
+                break
+
+        full_response = ""
+        usage_data = None
+        try:
+            tool_calls_result = None
+            async for chunk in self.llm_service.stream_chat(
+                messages=history,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                tools=TOOLS,
+            ):
+                if isinstance(chunk, dict):
+                    if chunk.get("type") == "fallback":
+                        yield _sse({"type": "status", "content": "Using Nyx Lite — full model resuming shortly"})
+                    elif chunk.get("type") == "tool_calls":
+                        tool_calls_result = chunk["calls"]
+                    else:
+                        usage_data = chunk
+                else:
+                    full_response += chunk
+                    yield _sse({"type": "content", "content": chunk})
+
+            if tool_calls_result:
+                for call in tool_calls_result:
+                    if call["name"] == "web_search":
+                        yield _sse({"type": "status", "content": "Searching the web..."})
+                        try:
+                            args = json.loads(call["arguments"])
+                        except json.JSONDecodeError:
+                            logger.warning("Invalid tool call arguments: %s", call["arguments"])
+                            continue
+                        results = await web_search(args.get("query", last_user_msg))
+                        result_text = format_search_results(results)
+
+                        history.append({
+                            "role": "assistant",
+                            "content": None,
+                            "tool_calls": [{
+                                "id": call["id"],
+                                "type": "function",
+                                "function": {
+                                    "name": call["name"],
+                                    "arguments": call["arguments"],
+                                },
+                            }],
+                        })
+                        history.append({
+                            "role": "tool",
+                            "tool_call_id": call["id"],
+                            "content": result_text or "No results found.",
+                        })
+
+                async for chunk in self.llm_service.stream_chat(
+                    messages=history,
+                    model=model,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                ):
+                    if isinstance(chunk, dict):
+                        usage_data = chunk
+                    else:
+                        full_response += chunk
+                        yield _sse({"type": "content", "content": chunk})
+
+        except Exception as e:
+            logger.error("LLM streaming error for conv %s: %s", conv.id, e, exc_info=True)
+            yield _sse({"type": "error", "content": "Failed to generate response. Please try again."})
+            return
+
+        full_response = _strip_duplicate_trailing_code(full_response)
+
+        if full_response:
+            await self.conversation_service.add_message(conv.id, "assistant", full_response)
+
+        if self.user_id and usage_data:
+            try:
+                from backend.services.usage_service import UsageService
+                usage_service = UsageService(self.db)
+                await usage_service.record_usage(
+                    user_id=self.user_id,
+                    conversation_id=conv.id,
+                    model=model,
+                    prompt_tokens=usage_data.get("prompt_tokens", 0),
+                    completion_tokens=usage_data.get("completion_tokens", 0),
+                )
+            except Exception as e:
+                logger.warning("Failed to record usage: %s", e)
+
+        conv.updated_at = datetime.now(timezone.utc)
+        await self.db.commit()
+        yield "data: [DONE]\n\n"
+
+    def _generate_title(self, first_message: str) -> str:
+        title = first_message.strip()
+        if len(title) > 50:
+            title = title[:47] + "..."
+        return title
+
+
+def _sse(data: dict) -> str:
+    return f"data: {json.dumps(data)}\n\n"